fix: pin verl<0.7.0 to prevent incompatible upgrade breaking vLLM async server

Co-authored-by: ultmaster <8463288+ultmaster@users.noreply.github.com>

.github/workflows/backport.yml

@@ -18,7 +18,7 @@ jobs:
     # Don't run on closed unmerged pull requests
     if: github.event.pull_request.merged
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - name: Create backport pull requests
         uses: korthout/backport-action@v3
         with:

.github/workflows/badge-apo.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-azure.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

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

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-chartqa.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

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

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-compat.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-examples.yml

@@ -25,7 +25,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-latest.yml

@@ -23,7 +23,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-rag.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-spider.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-tinker.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-unit.yml

@@ -18,7 +18,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/badge-unsloth.yml

@@ -17,7 +17,7 @@ jobs:
     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/checkout@v6
       - uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/benchmark.yml

@@ -3,6 +3,9 @@ permissions:
   contents: read
 on:
   workflow_dispatch:
+  schedule:
+    # Every Monday and Thursday at 3 AM UTC+8
+    - cron: '0 19 * * 0,3'
 
 jobs:
   benchmark:
@@ -25,7 +28,7 @@ jobs:
             runner:
               - self-hosted
               - 1ES.Pool=agl-runner-cpu
-            timeout: 60
+            timeout: 45
             args: >-
               --mode batch
               --total-tasks 4096
@@ -40,7 +43,7 @@ jobs:
             runner:
               - self-hosted
               - 1ES.Pool=agl-runner-cpu
-            timeout: 60
+            timeout: 45
             args: >-
               --mode batch
               --total-tasks 10000
@@ -60,35 +63,35 @@ jobs:
               --mode batch
               --total-tasks 20000
               --batch-size 2048
-              --n-runners 256
-              --max-rounds 8
+              --n-runners 300
+              --max-rounds 6
               --sleep-seconds 0.1
           - id: scenario-large-batch
             display: Large batch waves
             kind: scenario
-            store_workers: 32
+            store_workers: 96
             runner:
               - self-hosted
-              - 1ES.Pool=agl-runner-cpu
-            timeout: 60
+              - 1ES.Pool=agl-runner-cpu-high
+            timeout: 120
             args: >-
               --mode batch
-              --total-tasks 100000
+              --total-tasks 50000
               --batch-size 8192
-              --n-runners 256
-              --max-rounds 6
+              --n-runners 1000
+              --max-rounds 3
               --sleep-seconds 0.1
           - id: scenario-long-queues
             display: Long rollout queues
             kind: scenario
-            store_workers: 32
+            store_workers: 48
             runner:
               - self-hosted
               - 1ES.Pool=agl-runner-cpu
-            timeout: 60
+            timeout: 120
             args: >-
               --mode batch_partial
-              --total-tasks 100000
+              --total-tasks 50000
               --batch-size 1024
               --n-runners 256
               --remaining-tasks 4096
@@ -97,14 +100,14 @@ jobs:
           - id: scenario-high-concurrency
             display: High-throughput concurrent requests
             kind: scenario
-            store_workers: 32
+            store_workers: 96
             runner:
               - self-hosted
-              - 1ES.Pool=agl-runner-cpu
-            timeout: 60
+              - 1ES.Pool=agl-runner-cpu-high
+            timeout: 120
             args: >-
               --mode single
-              --total-tasks 100000
+              --total-tasks 50000
               --concurrency 2048
               --n-runners 256
               --max-rounds 2
@@ -112,10 +115,10 @@ jobs:
           - id: scenario-heavy-traces
             display: Heavy rollouts with deep traces
             kind: scenario
-            store_workers: 64
+            store_workers: 96
             runner:
               - self-hosted
-              - 1ES.Pool=agl-runner-cpu
+              - 1ES.Pool=agl-runner-cpu-high
             timeout: 60
             args: >-
               --mode batch_partial
@@ -169,9 +172,11 @@ jobs:
             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 }}
@@ -183,7 +188,7 @@ jobs:
       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: actions/checkout@v6
 
       - uses: astral-sh/setup-uv@v7
         with:
@@ -316,7 +321,7 @@ jobs:
 
       - name: Upload workload artifacts
         if: ${{ always() }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ env.ARTIFACT_DIR }}
@@ -338,27 +343,27 @@ jobs:
             runner: ubuntu-latest
         workload:
           - id: high-insert
-            total_tasks: 100000
+            total_tasks: 50000
             concurrency: 2048
             type: insert
           - id: medium-insert
-            total_tasks: 100000
+            total_tasks: 50000
             concurrency: 128
             type: insert
           - id: low-insert
-            total_tasks: 100000
+            total_tasks: 50000
             concurrency: 4
             type: insert
           - id: high-dequeue
-            total_tasks: 100000
+            total_tasks: 50000
             concurrency: 2048
             type: dequeue
           - id: medium-dequeue
-            total_tasks: 100000
+            total_tasks: 50000
             concurrency: 128
             type: dequeue
           - id: low-dequeue
-            total_tasks: 100000
+            total_tasks: 50000
             concurrency: 4
             type: dequeue
     env:
@@ -367,7 +372,7 @@ jobs:
       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: actions/checkout@v6
 
       - uses: astral-sh/setup-uv@v7
         with:
@@ -429,7 +434,7 @@ jobs:
 
       - name: Upload collection artifacts
         if: ${{ always() }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ env.ARTIFACT_DIR }}

.github/workflows/dashboard.yml

@@ -17,7 +17,7 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 15
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           fetch-depth: 0
       - uses: actions/setup-node@v6

.github/workflows/docs.yml

@@ -21,7 +21,7 @@ jobs:
   deploy:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           fetch-depth: 0
       - uses: actions/setup-python@v6

.github/workflows/examples-apo.yml

@@ -42,7 +42,7 @@ jobs:
           setup-script: 'latest'
       fail-fast: false
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -69,7 +69,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-apo-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

.github/workflows/examples-azure.yml

@@ -39,7 +39,7 @@ jobs:
     steps:
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -60,7 +60,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-azure-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

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

@@ -45,7 +45,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -72,7 +72,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-calc-x-performance-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt
@@ -158,7 +158,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -171,12 +171,12 @@ jobs:
       - name: Sync dependencies (latest)
         run: |
           uv sync --frozen --no-default-groups --extra verl \
-            --group dev --group experiment --group agents --group torch-gpu-stable
+            --group dev --group experiment --group agents --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 --group torch-gpu-${{ matrix.setup-script }}
+            --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: |
@@ -185,7 +185,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-calc-x-variants-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt
@@ -270,6 +270,33 @@ jobs:
           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
@@ -295,6 +322,52 @@ jobs:
           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

.github/workflows/examples-chartqa.yml

@@ -41,7 +41,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -62,7 +62,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-chartqa-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

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

@@ -43,7 +43,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -64,7 +64,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-claude-code-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt
@@ -109,7 +109,7 @@ jobs:
 
       - name: Upload sanity check artifacts for vLLM
         if: ${{ always() }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: claude-code-sanity-check-vllm-${{ matrix.setup-script }}
           path: |
@@ -142,7 +142,7 @@ jobs:
 
       - name: Upload sanity check artifacts for OpenAI
         if: ${{ always() }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: claude-code-sanity-check-openai-${{ matrix.setup-script }}
           path: |

.github/workflows/examples-compat.yml

@@ -43,7 +43,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -65,7 +65,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-backward-compatibility-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

.github/workflows/examples-rag.yml

@@ -45,7 +45,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -72,7 +72,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-rag-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

.github/workflows/examples-spider.yml

@@ -44,7 +44,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -71,7 +71,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-spider-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

.github/workflows/examples-tinker.yml

@@ -41,7 +41,7 @@ jobs:
     steps:
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -62,7 +62,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-tinker-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

.github/workflows/examples-unsloth.yml

@@ -44,7 +44,7 @@ jobs:
         run: nvidia-smi
       - name: Check disk space
         run: df -h
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -65,7 +65,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-unsloth-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt

.github/workflows/playground.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
       - name: Run script
         run: |
           echo "Hello, world!"

.github/workflows/pypi-nightly.yml

@@ -14,7 +14,7 @@ jobs:
       contents: read
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           fetch-depth: 0
       - uses: actions/setup-python@v6

.github/workflows/pypi-release.yml

@@ -16,7 +16,7 @@ jobs:
       tag_version: ${{ steps.get_tag.outputs.tag_version }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Get version from pyproject.toml
         id: get_version
@@ -48,7 +48,7 @@ jobs:
       contents: read
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           fetch-depth: 0
       - uses: actions/setup-python@v6

.github/workflows/tests-full.yml

@@ -45,6 +45,12 @@ jobs:
             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
@@ -54,7 +60,7 @@ jobs:
           # Other uncovered tests
           - id: others
             display-name: Others
-            pytest-mark: 'not store and not agentops and not gpu and not llmproxy'
+            pytest-mark: 'not store and not agentops and not weave and not gpu and not llmproxy'
             runs-on: ubuntu-latest
             has-gpu: false
         env:
@@ -69,7 +75,7 @@ jobs:
       - name: Check GPU status
         if: matrix.mark.has-gpu
         run: nvidia-smi
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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 }}
 
@@ -83,24 +89,24 @@ jobs:
 
       - name: Sync dependencies (latest, gpu)
         if: matrix.env.setup-script == 'latest' && matrix.mark.has-gpu
-        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-stable
+        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 mongo --group dev --group agents --group langchain --group core-stable
+        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 mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.env.setup-script }}
+        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 mongo --group dev --group agents --group langchain --group core-stable
+        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 mongo --group dev --group agents --group torch-gpu-legacy
+        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 mongo --group dev --group agents --group core-legacy
+        run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group core-legacy
 
       - name: Freeze dependencies
         run: |
@@ -109,7 +115,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-tests-full-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
           path: requirements-freeze.txt
@@ -126,49 +132,7 @@ jobs:
         run: cd dashboard && npm run build
 
       - name: Setup Docker environments
-        run: |
-          set -euo pipefail
-
-          cd docker
-
-          # Setup data directories
-          ./setup.sh
-
-          # Start Dockers
-          docker compose -f compose.mongo.yml up -d
-
-          SERVICE_NAME=mongo
-          TIMEOUT=60   # seconds
-          SLEEP=2
-
-          cid="$(docker compose -f compose.mongo.yml ps -q "$SERVICE_NAME")"
-          if [ -z "$cid" ]; then
-            echo "Service $SERVICE_NAME is not running"
-            exit 1
-          fi
-
-          echo "Waiting for $SERVICE_NAME to become healthy..."
-          end=$((SECONDS + TIMEOUT))
-
-          while [ "$SECONDS" -lt "$end" ]; do
-            status="$(docker inspect -f '{{.State.Health.Status}}' "$cid")"
-            echo "Current status: $status"
-
-            if [ "$status" = "healthy" ]; then
-              echo "$SERVICE_NAME is healthy ✅"
-              exit 0
-            elif [ "$status" = "unhealthy" ]; then
-              echo "$SERVICE_NAME is unhealthy ❌"
-              docker logs "$cid" || true
-              exit 1
-            fi
-
-            sleep "$SLEEP"
-          done
-
-          echo "Timed out waiting for $SERVICE_NAME to become healthy after ${TIMEOUT}s"
-          docker logs "$cid" || true
-          exit 1
+        run: ./scripts/mongodb_docker_run.sh
         shell: bash
 
       - name: Launch LiteLLM Proxy
@@ -181,7 +145,7 @@ jobs:
       # mongo, openai, gpu, all enabled by default
       - name: Run tests
         run: |
-          uv run pytest -v --durations=0 tests -m "${{ matrix.mark.pytest-mark }}"
+          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/
@@ -211,7 +175,7 @@ jobs:
     steps:
       - name: Check GPU status
         run: nvidia-smi
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         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
@@ -238,7 +202,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-minimal-examples-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt
@@ -270,6 +234,14 @@ jobs:
           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

.github/workflows/tests.yml

@@ -25,7 +25,7 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
@@ -40,6 +40,7 @@ jobs:
         run: |
           uv sync --frozen \
             --extra apo \
+            --extra weave \
             --extra verl \
             --extra mongo \
             --group dev \
@@ -56,10 +57,13 @@ jobs:
         uses: pre-commit/action@v3.0.1
       - name: Check Python headers
         run: uv run --locked --no-sync scripts/check_headers.py
+        if: matrix.setup == 'fast'
       - name: Run Black
         run: uv run --locked --no-sync black --check .
+        if: matrix.setup != 'next'
       - name: Run isort
         run: uv run --locked --no-sync isort --check-only .
+        if: matrix.setup != 'next'
       - name: Run pyright (fast)
         run: uv run --locked --no-sync pyright -p pyrightconfig.fast.json
         if: matrix.setup == 'fast'
@@ -72,7 +76,7 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - uses: actions/setup-node@v6
         with:
           node-version: '22'
@@ -96,7 +100,7 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           fetch-depth: 0
       - uses: actions/setup-python@v6
@@ -110,10 +114,14 @@ jobs:
       - name: Set source commit for docs
         run: |
           echo "SOURCE_COMMIT=${{ github.sha }}" >> $GITHUB_ENV
+      - name: Verify OpenAPI specification is up-to-date
+        run: |
+          uv run --locked --no-sync python scripts/export_openapi.py
+          git diff --exit-code docs/assets/store-openapi.json
       - name: Build documentation
         run: uv run --locked --no-sync mkdocs build --strict
       - name: Upload docs artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: documentation-site
           path: site/
@@ -131,6 +139,10 @@ jobs:
           - 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
@@ -142,7 +154,7 @@ jobs:
           # unmarked tests: adapter, execution engine, etc.
           - id: others
             display-name: Others
-            pytest-mark: 'not store and not agentops and not llmproxy and not utils'
+            pytest-mark: 'not store and not agentops and not weave and not llmproxy and not utils'
         env:
           - python-version: '3.10'
             setup-script: 'legacy'
@@ -158,7 +170,7 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 15
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
@@ -167,10 +179,10 @@ jobs:
         run: uv lock --upgrade
         if: matrix.env.setup-script == 'latest'
       - name: Sync dependencies (latest)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-stable
+        run: uv sync --frozen --no-default-groups --extra apo --extra weave --group dev --group agents --group langchain --group core-stable
         if: matrix.env.setup-script == 'latest'
       - name: Sync dependencies (stable & legacy)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-${{ matrix.env.setup-script }}
+        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: |
@@ -179,7 +191,7 @@ jobs:
           echo "UV_LOCKED=1" >> $GITHUB_ENV
           echo "UV_NO_SYNC=1" >> $GITHUB_ENV
       - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: dependencies-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
           path: requirements-freeze.txt
@@ -206,7 +218,7 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 15
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           fetch-depth: 0
       - uses: actions/setup-node@v6

.gitignore

@@ -1,7 +1,8 @@
 # Agentlightning specific files
 verl_old
 meta-llama/**
-debug/*.png
+**/debug/**/*.png
+**/debug/**/*.json
 requirements-freeze*.txt
 /playground
 

.pre-commit-config.yaml

@@ -3,13 +3,14 @@ repos:
     rev: v6.0.0
     hooks:
       - id: end-of-file-fixer
+        exclude: (.*store-openapi\.json$)
       - id: trailing-whitespace
       - id: check-yaml
         exclude: ^mkdocs\.yml$
       - id: check-toml
       - id: check-added-large-files
         args: ["--maxkb=1024"]
-        exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)
+        exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)|(.*store-openapi\.json$)
       - id: check-shebang-scripts-are-executable
       - id: detect-private-key
   - repo: https://github.com/pycqa/isort

README.md

@@ -37,7 +37,7 @@ 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
+pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ --pre agentlightning
 ```
 
 Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.
@@ -46,6 +46,7 @@ To start using Agent-lightning, check out our [documentation](https://microsoft.
 
 ## ⚡ Articles
 
+- 12/17/2025 [Adopting the Trajectory Level Aggregation for Faster Training](https://agent-lightning.github.io/posts/trajectory_level_aggregation/) Agent-lightning blog.
 - 11/4/2025 [Tuning ANY AI agent with Tinker ✕ Agent-lightning](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-1-1d8c9a397f0e) Medium. See also [Part 2](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-2-332c5437f0dc).
 - 10/22/2025 [No More Retokenization Drift: Returning Token IDs via the OpenAI Compatible API Matters in Agent RL](https://blog.vllm.ai/2025/10/22/agent-lightning.html) vLLM blog. See also [Zhihu writeup](https://zhuanlan.zhihu.com/p/1965067274642785725).
 - 8/11/2025 [Training AI Agents to Write and Self-correct SQL with Reinforcement Learning](https://medium.com/@yugez/training-ai-agents-to-write-and-self-correct-sql-with-reinforcement-learning-571ed31281ad) Medium.
@@ -57,7 +58,7 @@ To start using Agent-lightning, check out our [documentation](https://microsoft.
 
 - [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
 - [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.
-- [Youtu-Agent](https://github.com/TencentCloudADP/Youtu-agent) — Youtu-Agent lets you build and train your agent with ease. Built with [a modified branch](https://github.com/microsoft/agent-lightning/tree/contrib/youtu-agent-lightning) of Agent Lightning, Youtu-Agent has verified up to 128 GPUs RL training on maths/code and search capabilities with steady convergence. Also check [the recipe](https://github.com/TencentCloudADP/youtu-agent/tree/rl/agl).
+- [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
 

agentlightning/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"
 
 from .adapter import *
 from .algorithm import *

agentlightning/adapter/triplet.py

@@ -12,6 +12,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel
 
 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
 
@@ -20,6 +21,47 @@ from .base import TraceAdapter
 logger = logging.getLogger(__name__)
 
 
+def _attributes_get_multiple(attributes: Dict[str, Any], keys: List[str]) -> Optional[str]:
+    """Get a string from the attributes, if present.
+    If there are multiple matches, the first one is returned.
+    """
+    for key in keys:
+        if key in attributes:
+            if isinstance(attributes[key], str):
+                return attributes[key]
+            else:
+                logger.warning(f"Attribute {key} is found but is not a string: {attributes[key]}")
+    return None
+
+
+def _attributes_get_ids_multiple(attributes: Dict[str, Any], keys: List[str]) -> Optional[List[int]]:
+    """Get a list of integers from the attributes, if present.
+    If there are multiple matches, the first one is returned.
+    """
+    for key in keys:
+        if key in attributes:
+            if (isinstance(attributes[key], list) or isinstance(attributes[key], tuple)) and all(
+                isinstance(x, int) for x in attributes[key]
+            ):
+                return list(attributes[key])
+            else:
+                logger.warning(f"Attribute {key} is found but is not a list of integers: {attributes[key]}")
+    return None
+
+
+def _attributes_unflatten_multiple(
+    attributes: Dict[str, Any], keys: List[str]
+) -> Union[Dict[str, Any], List[Any], None]:
+    """Unflatten the attributes, if present.
+    If there are multiple matches, the first one is returned.
+    """
+    for key in keys:
+        result = filter_and_unflatten_attributes(attributes, key)
+        if result:
+            return result
+    return None
+
+
 class Transition(BaseModel):
     """A single transition within a reinforcement learning trajectory.
 
@@ -132,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
@@ -309,6 +351,19 @@ class TraceTree:
         if agent_name is not None:
             return agent_name
 
+        # Case 6: Weave
+        is_agent_type = attributes.get("type") == "agent"
+        if is_agent_type:
+            agent_name = cast(Optional[str], attributes.get("agentlightning.operation.input.name"))
+            if agent_name is not None:
+                return agent_name
+
+        # Case 7: Weave + LangChain
+        if self.span.name.startswith("langchain.Chain."):
+            attributes_lc_name = cast(Optional[str], attributes.get("lc_name"))
+            if attributes_lc_name is not None:
+                return attributes_lc_name
+
     def maybe_reward_dict(self) -> dict[str, Any]:
         """Return a reward payload if the span encodes one.
 
@@ -328,7 +383,17 @@ class TraceTree:
             `True` when the span payload describes a reward, otherwise `False`.
         """
         maybe_reward = self.maybe_reward_dict()
-        return maybe_reward and maybe_reward.get("type") == "reward"  # type: ignore
+        if maybe_reward and maybe_reward.get("type") == "reward":  # type: ignore
+            return True
+
+        # Agent-lightning 0.3+
+        if (
+            self.span.name == AGL_OPERATION
+            and self.span.attributes.get(LightningSpanAttributes.OPERATION_NAME.value) == AGL_REWARD
+        ):
+            return True
+
+        return False
 
     def find_llm_calls(
         self,
@@ -365,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 (
@@ -547,7 +614,7 @@ class TraceTree:
                 try:
                     content = json.loads(content)  # This content should now be a list
                 except json.JSONDecodeError:
-                    logger.warning(f"Failed to parse message content as JSON: {content}")
+                    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):
@@ -567,18 +634,57 @@ class TraceTree:
         Subclass can override this method to add more fields to the triplet,
         such as chat messages and tool calls.
         """
-        prompt_token_ids = span.attributes.get("prompt_token_ids", [])  # type: ignore
-        response_token_ids = span.attributes.get("response_token_ids", [])  # type: ignore
-        response_id = span.attributes.get("gen_ai.response.id", None)  # type: ignore
-        request_metadata = filter_and_unflatten_attributes(span.attributes, "gen_ai.request")
-        response_metadata = filter_and_unflatten_attributes(span.attributes, "gen_ai.response")
-        prompt_raw_content = filter_and_unflatten_attributes(span.attributes, "gen_ai.prompt")
-        completion_raw_content = filter_and_unflatten_attributes(span.attributes, "gen_ai.completion")
-        image_urls = self.extract_prompt_image_urls(prompt_raw_content)
+        prompt_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)

agentlightning/algorithm/apo/apo.py

@@ -112,6 +112,8 @@ class APO(Algorithm, Generic[T_task]):
         beam_rounds: int = 3,
         rollout_batch_timeout: float = 3600.0,
         run_initial_validation: bool = True,
+        gradient_prompt_files: Optional[List[Path]] = None,
+        apply_edit_prompt_files: Optional[List[Path]] = None,
         # Internal flags for debugging
         _poml_trace: bool = False,
     ):
@@ -132,6 +134,8 @@ class APO(Algorithm, Generic[T_task]):
             rollout_batch_timeout: Maximum time in seconds to wait for rollout batch completion.
             run_initial_validation: If True, runs validation on the seed prompt before starting
                 optimization to establish a baseline score. Defaults to True.
+            gradient_prompt_files: Prompt templates used to compute textual gradients (critiques).
+            apply_edit_prompt_files: Prompt templates used to apply edits based on critiques.
         """
         self.async_openai_client = async_openai_client
         self.gradient_model = gradient_model
@@ -144,6 +148,8 @@ class APO(Algorithm, Generic[T_task]):
         self.beam_rounds = beam_rounds
         self.rollout_batch_timeout = rollout_batch_timeout
         self.run_initial_validation = run_initial_validation
+        self.gradient_prompt_files = gradient_prompt_files or GRADIENT_PROMPT_FILES
+        self.apply_edit_prompt_files = apply_edit_prompt_files or APPLY_EDIT_PROMPT_FILES
 
         self._history_best_prompt: Optional[PromptTemplate] = None
         self._history_best_score: float = float("-inf")
@@ -270,7 +276,7 @@ class APO(Algorithm, Generic[T_task]):
         Returns:
             A textual critique generated by the LLM, or None if generation fails.
         """
-        tg_template = random.choice(GRADIENT_PROMPT_FILES)
+        tg_template = random.choice(self.gradient_prompt_files)
 
         if len(rollout_results) < self.gradient_batch_size:
             self._log(
@@ -352,7 +358,7 @@ class APO(Algorithm, Generic[T_task]):
             return current_prompt.prompt_template.template
 
         # 2) Apply edit
-        ae_template = random.choice(APPLY_EDIT_PROMPT_FILES)
+        ae_template = random.choice(self.apply_edit_prompt_files)
         self._log(
             logging.INFO,
             f"Edit will be generated by {self.apply_edit_model} with template: {ae_template.name}",

agentlightning/algorithm/verl/interface.py

@@ -32,6 +32,30 @@ class VERL(Algorithm):
         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

agentlightning/emitter/__init__.py

@@ -1,5 +1,16 @@
 # 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, get_message_value

agentlightning/emitter/annotation.py

@@ -5,7 +5,6 @@
 import asyncio
 import functools
 import inspect
-import json
 import logging
 from types import TracebackType
 from typing import (
@@ -22,19 +21,18 @@ from typing import (
     overload,
 )
 
-from opentelemetry import trace
-from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.trace import Status, StatusCode
-
 from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
-from agentlightning.utils.otel import flatten_attributes, get_tracer
+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) -> ReadableSpan:
+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].
@@ -48,62 +46,46 @@ def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> Reada
     Args:
         annotation: Dictionary containing annotation key-value pairs.
             Representatives are rewards, tags, and metadata.
-        propagate: Whether to propagate the span to exporters automatically.
+        propagate: Whether to propagate the span to tracers automatically.
     """
-    annotation_attributes = flatten_attributes(annotation)
-    if any(not isinstance(v, (str, int, float, bool, bytes)) for v in annotation_attributes.values()):
-        raise TypeError("All annotation attributes must be primitive types (str, int, float, bool, bytes)")
+    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())
 
-    # TODO: this should use a tracer from current context rather than the singleton
-    tracer = get_tracer(use_active_span_processor=propagate)
-    span = tracer.start_span(
-        AGL_ANNOTATION,
-        attributes=annotation_attributes,
+    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"),
     )
-    logger.debug("Emitting annotation span with keys %s", annotation_attributes)
-    with span:
-        pass
-    if not isinstance(span, ReadableSpan):
-        raise ValueError(f"Span is not a ReadableSpan: {span}")
-
-    return span
-
-
-def _safe_json_dump(obj: Any) -> str:
-    """Serialize an object to JSON, falling back to ``str(obj)`` if needed.
-
-    Args:
-        obj: Object to be serialized.
-
-    Returns:
-        The JSON-encoded string representation of the object, or its string
-        representation if JSON encoding fails.
-    """
-    try:
-        return json.dumps(obj, default=str, ensure_ascii=False)
-    except Exception:
-        return str(obj)
 
 
 class OperationContext:
     """Context manager and decorator for tracing operations.
 
-    This class manages an OpenTelemetry span for a logical unit of work. It can
-    be used either:
+    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 :meth:`set_input` and :meth:`set_output`.
+      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: OpenTelemetry tracer used to create spans.
-        span: The currently active span, if any.
+        tracer: Tracer implementation used to create spans.
     """
 
-    def __init__(self, name: str, attributes: Dict[str, Any], *, propagate: bool = True) -> None:
+    def __init__(self, name: str, attributes: Dict[str, Any], propagate: bool = True) -> None:
         """Initialize a new operation context.
 
         Args:
@@ -112,12 +94,19 @@ class OperationContext:
                 JSON-serialized where necessary.
             propagate: Whether the span should be sent to active exporters.
         """
-        self.name: str = name
-        self.initial_attributes: Dict[str, Any] = attributes
-        self.propagate: bool = propagate
-        self.tracer: trace.Tracer = get_tracer(use_active_span_processor=propagate)
-        self.span: Optional[trace.Span] = None
-        self._ctx_token: Optional[ContextManager[Any]] = None
+        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.
@@ -125,15 +114,10 @@ class OperationContext:
         Returns:
             The current :class:`OperationContext` instance with an active span.
         """
-        # 1. Start the span with initial attributes (JSON serialized)
-        sanitized_attrs = {
-            k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
-            for k, v in self.initial_attributes.items()
-        }
-
-        self.span = self.tracer.start_span(self.name, attributes=sanitized_attrs)
-        self._ctx_token = trace.use_span(self.span, end_on_exit=True)
-        self._ctx_token.__enter__()
+        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__(
@@ -142,57 +126,63 @@ class OperationContext:
         exc_val: Optional[BaseException],
         exc_tb: Optional[TracebackType],
     ) -> None:
-        """Exit the context manager and finish the span.
+        """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
 
-        Any exception raised inside the context is recorded on the span and the
-        span status is set to error.
-
-        Args:
-            exc_type: Exception type, if an exception occurred.
-            exc_val: Exception instance, if an exception occurred.
-            exc_tb: Traceback object, if an exception occurred.
-        """
-        # 1. Record Exception if present
-        if exc_val and self.span:
-            self.span.record_exception(exc_val)
-            self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
-
-        # 2. Close span
-        if self._ctx_token:
-            self._ctx_token.__exit__(exc_type, exc_val, exc_tb)
+    def 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`` attribute,
-        and keyword arguments are stored under ``input.<name>`` attributes.
+        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.
+        This is intended for use inside a `with operation(...) as op` block.
 
         Args:
             *args: Positional arguments to record.
             **kwargs: Keyword arguments to record.
         """
-        if not self.span:
-            return
+        if not self._recording_context:
+            raise RuntimeError("No recording context found. Cannot set input.")
 
+        prefix = LightningSpanAttributes.OPERATION_INPUT.value
+        attributes: Dict[str, Any] = {}
         if args:
-            self.span.set_attribute("input.args", _safe_json_dump(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 k, v in kwargs.items():
-                self.span.set_attribute(f"input.{k}", _safe_json_dump(v))
+            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.
+        This is intended for use inside a `with operation(...) as op` block.
 
         Args:
             output: The output value to record.
         """
-        if not self.span:
-            return
-        self.span.set_attribute("output", _safe_json_dump(output))
+        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.
@@ -212,60 +202,64 @@ class OperationContext:
 
         sig = inspect.signature(fn)
 
-        def _record_auto_inputs(span: trace.Span, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> None:
-            """Bind arguments to signature and log them on the span.
+        sanitized_init_attrs = sanitize_attributes(
+            {LightningSpanAttributes.OPERATION_NAME.value: function_name, **self.initial_attributes}
+        )
 
-            Args:
-                span: Span on which to record attributes.
-                args: Positional arguments passed to the wrapped callable.
-                kwargs: Keyword arguments passed to the wrapped callable.
-            """
+        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 k, v in bound.arguments.items():
-                    span.set_attribute(
-                        f"{LightningSpanAttributes.OPERATION_INPUT.value}.{k}",
-                        _safe_json_dump(v),
-                    )
+                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:
-                span.set_attribute(
-                    f"{LightningSpanAttributes.OPERATION_INPUT.value}.args",
-                    _safe_json_dump(args),
-                )
-                span.set_attribute(
-                    f"{LightningSpanAttributes.OPERATION_INPUT.value}.kwargs",
-                    _safe_json_dump(kwargs),
-                )
+                if 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))
 
-        if asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn):
+        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 inspect.iscoroutinefunction(fn) or (
+            # For backwards compatibility.
+            hasattr(asyncio, "iscoroutinefunction")
+            and asyncio.iscoroutinefunction(fn)  # type: ignore
+        ):
 
             @functools.wraps(fn)
             async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                 """Async wrapper that traces the wrapped coroutine."""
-                # Reuse __enter__ logic via 'with self' would share state incorrectly
-                # across concurrent calls. We must create a new span per call.
-                # So we manually reimplement the span logic for the wrapper here.
-
-                sanitized_attrs = {
-                    k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
-                    for k, v in self.initial_attributes.items()
-                }
-
-                with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
-                    span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
-                    _record_auto_inputs(span, args, kwargs)
-                    try:
-                        result = await fn(*args, **kwargs)
-                        span.set_attribute(
-                            LightningSpanAttributes.OPERATION_OUTPUT.value,
-                            _safe_json_dump(result),
-                        )
-                        return result
-                    except Exception as e:
-                        span.record_exception(e)
-                        span.set_status(Status(StatusCode.ERROR, str(e)))
-                        raise
+                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)
 
@@ -274,41 +268,48 @@ class OperationContext:
             @functools.wraps(fn)
             def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                 """Sync wrapper that traces the wrapped callable."""
-                sanitized_attrs = {
-                    k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
-                    for k, v in self.initial_attributes.items()
-                }
-
-                with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
-                    span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
-                    _record_auto_inputs(span, args, kwargs)
-                    try:
-                        result = fn(*args, **kwargs)
-                        span.set_attribute(
-                            LightningSpanAttributes.OPERATION_OUTPUT.value,
-                            _safe_json_dump(result),
-                        )
-                        return result
-                    except Exception as e:
-                        span.record_exception(e)
-                        span.set_status(Status(StatusCode.ERROR, str(e)))
-                        raise
+                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, **additional_attributes: Any) -> _FnType: ...
+def operation(
+    fn: _FnType, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
+) -> _FnType: ...
 
 
 @overload
-def operation(*, propagate: bool = True, **additional_attributes: Any) -> OperationContext: ...
+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.
@@ -344,6 +345,9 @@ def operation(
             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.
 
@@ -352,6 +356,12 @@ def operation(
         [`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

agentlightning/emitter/exception.py

@@ -1,13 +1,13 @@
 # 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.semconv import AGL_EXCEPTION
-from agentlightning.utils.otel import get_tracer
+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__)
 
@@ -32,25 +32,23 @@ def emit_exception(
     """
     if not isinstance(exception, BaseException):  # type: ignore
         raise TypeError(f"Expected a BaseException instance, got: {type(exception)}.")
-
-    tracer = get_tracer(use_active_span_processor=propagate)
-    stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
-    span_attributes = {
-        exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
-        exception_attributes.EXCEPTION_MESSAGE: str(exception),
-        exception_attributes.EXCEPTION_ESCAPED: True,
-    }
-    if stacktrace.strip():
-        span_attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+    span_attributes = format_exception_attributes(exception)
 
     if attributes:
-        span_attributes.update(attributes)
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))
 
-    span = tracer.start_span(
+    logger.debug("Emitting exception span for %s", type(exception).__name__)
+
+    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"),
     )
-    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.

agentlightning/emitter/message.py

@@ -4,8 +4,10 @@ import logging
 from typing import Any, Dict, Optional
 
 from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
-from agentlightning.types import SpanLike
-from agentlightning.utils.otel import get_tracer
+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__)
 
@@ -27,17 +29,21 @@ def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, prop
     if not isinstance(message, str):  # type: ignore
         raise TypeError(f"Message must be a string or list of strings, got: {type(message)}.")
 
-    tracer = get_tracer(use_active_span_processor=propagate)
-    span_attributes = {LightningSpanAttributes.MESSAGE_BODY.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:
-        span_attributes.update(attributes)
-    span = tracer.start_span(
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))
+    logger.debug("Emitting message span with message: %s", message)
+    tracer.create_span(
         AGL_MESSAGE,
         attributes=span_attributes,
     )
-    logger.debug("Emitting message span with message: %s", message)
-    with span:
-        pass
 
 
 def get_message_value(span: SpanLike) -> Optional[str]:

agentlightning/emitter/object.py

@@ -6,13 +6,15 @@ import logging
 from typing import Any, Dict, Optional
 
 from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
-from agentlightning.types import SpanLike
-from agentlightning.utils.otel import full_qualified_name, get_tracer
+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, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
+def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> SpanCoreFields:
     """Emit an object's serialized representation as an OpenTelemetry span.
 
     Args:
@@ -25,20 +27,29 @@ def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propag
     """
     span_attributes = encode_object(object)
     if attributes:
-        span_attributes.update(attributes)
-    tracer = get_tracer(use_active_span_processor=propagate)
-    span = tracer.start_span(
-        AGL_OBJECT,
-        attributes=span_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)
-    with span:
-        pass
+
+    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"),
+    )
 
 
 def encode_object(object: Any) -> Dict[str, Any]:

agentlightning/emitter/reward.py

@@ -20,13 +20,10 @@ from typing import (
     cast,
 )
 
-import agentops
-from agentops.sdk.decorators import operation
-from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import TypeAdapter
 
 from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
-from agentlightning.types import SpanLike
+from agentlightning.types import SpanCoreFields, SpanLike
 from agentlightning.utils.otel import filter_and_unflatten_attributes
 
 from .annotation import emit_annotation
@@ -61,6 +58,8 @@ _FnType = TypeVar("_FnType", bound=Callable[..., Any])
 
 def _agentops_initialized() -> bool:
     """Return `True` when the AgentOps client has been configured."""
+    import agentops
+
     return agentops.get_client().initialized
 
 
@@ -81,6 +80,8 @@ def reward(fn: _FnType) -> _FnType:
         Wrapped callable that preserves the original signature.
     """
 
+    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:
@@ -91,7 +92,11 @@ def reward(fn: _FnType) -> _FnType:
         return {"type": "reward", "value": float(result)}
 
     # Check if the function is async
-    is_async = asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn)
+    is_async = inspect.iscoroutinefunction(fn) or (
+        # For backwards compatibility.
+        hasattr(asyncio, "iscoroutinefunction")
+        and asyncio.iscoroutinefunction(fn)  # type: ignore
+    )
 
     if is_async:
 
@@ -146,7 +151,7 @@ def emit_reward(
     primary_key: str | None = None,
     attributes: Dict[str, Any] | None = None,
     propagate: bool = True,
-) -> ReadableSpan:
+) -> SpanCoreFields:
     """Emit a reward value as an OpenTelemetry span.
 
     Examples:
@@ -172,11 +177,7 @@ def emit_reward(
         propagate: Whether to propagate the span to exporters automatically.
 
     Returns:
-        Readable span capturing the recorded reward.
-
-    Raises:
-        ValueError: If the provided reward cannot be interpreted as a float or the
-            resulting span is not a [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) instance.
+        Span core fields capturing the recorded reward.
     """
     logger.debug(f"Emitting reward: {reward}")
     reward_dimensions: List[RewardDimension] = []

agentlightning/instrumentation/__init__.py

@@ -39,13 +39,6 @@ try:
 except ImportError:
     pass
 
-try:
-    from . import weave  # type: ignore
-
-    WEAVE_INSTALLED = True  # type: ignore
-except ImportError:
-    pass
-
 
 def instrument_all():
     """Instrument all the instrumentation libraries."""
@@ -119,20 +112,3 @@ def uninstrument_all():
             warnings.warn("agentops_langchain is installed but uninstrument_agentops_langchain could not be imported.")
     else:
         warnings.warn("Agentops-langchain integration is not installed. It's therefore not uninstrumented.")
-
-
-def instrument_weave():
-    if WEAVE_INSTALLED:
-        from .weave import instrument_weave
-
-        instrument_weave()
-
-
-def uninstrument_weave():
-    if WEAVE_INSTALLED:
-        try:
-            from .weave import uninstrument_weave
-
-            uninstrument_weave()
-        except ImportError:
-            warnings.warn("weave is installed but uninstrument_weave could not be imported.")

agentlightning/instrumentation/weave.py

@@ -1,139 +1,541 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import logging
-import os
-from typing import Any, Callable, Optional
+from __future__ import annotations
 
-import requests
+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)
+
+    # Experimental unstable APIs
+    # We don't support these APIs yet.
+    def annotation_queue_create(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def annotation_queues_query_stream(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def annotation_queue_read(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def annotation_queue_add_calls(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def annotation_queues_stats(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def annotation_queue_items_query(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def annotator_queue_items_progress_update(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def calls_complete(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def call_start_v2(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def call_end_v2(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def call_stats(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def trace_usage(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    def calls_usage(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+
 # Module-level storage for originals
-_original_default_entity_name_getter: Callable[..., Any] | None = None
-_original_upsert_project_getter: Callable[..., Any] | None = None
-_original_weave_get = False
-_original_weave_post = False
+_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 instrument_weave():
-    """
-    Patch the Weave/W&B integration to bypass actual network calls for testing.
+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
 
-    - Mocks HTTP POST/GET requests
-    - Patches wandb.Api methods
-    - Silences Weave logging
-    - Sets dummy WANDB_API_KEY if not provided
-    """
+    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:
-        import weave
-        from weave.compat import wandb  # type: ignore
-    except ImportError:
-        logger.warning("Weave or wandb not installed; cannot uninstrument.")
-        return
+        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"
 
-    _weave_tracer_entity_name = "weave_tracer_entity"
 
-    def default_entity_name_getter(_self) -> str:  # type: ignore
-        return _weave_tracer_entity_name
+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 upsert_project_getter(
-        _self, project: str, description: Optional[str] = None, entity: Optional[str] = None  # type: ignore
-    ) -> dict[str, Any]:
-        return {
-            "upsertModel": {
-                "model": {
-                    "name": project,
-                    "description": description or "",
-                    "entity": entity or _weave_tracer_entity_name,
-                }
-            },
-            "project": "weave_tracer_project",
-        }
 
-    # Mock network requests to avoid real HTTP calls
-    def post(url: str, *args: Any, **kwargs: Any) -> requests.Response:
-        response = requests.Response()
-        response.status_code = 200
-        response._content = b'{"digest": "mocked_digest"}'
-        return response
+def instrument_weave(server: InMemoryWeaveTraceServer):
+    """Patch the Weave/W&B integration to bypass actual network calls for testing."""
 
-    def get(url: str, *args: Any, **kwargs: Any) -> requests.Response:
-        response = requests.Response()
-        response.status_code = 200
-        response._content = b'{"min_required_weave_python_version": "0.52.14"}'
-        return response
-
-    # Patch API methods and HTTP requests
-    global _original_default_entity_name_getter
-    global _original_upsert_project_getter
-    global _original_weave_post
-    global _original_weave_get
-    _original_default_entity_name_getter = wandb.Api.default_entity_name  # type: ignore
-    _original_upsert_project_getter = wandb.Api.upsert_project  # type: ignore
-    _original_weave_post = weave.utils.http_requests.session.post  # type: ignore
-    _original_weave_get = weave.utils.http_requests.session.get  # type: ignore
-
-    # Patch API methods and HTTP requests
-    wandb.Api.default_entity_name = default_entity_name_getter  # type: ignore
-    wandb.Api.upsert_project = upsert_project_getter  # type: ignore
-    weave.utils.http_requests.session.post = post  # type: ignore
-    weave.utils.http_requests.session.get = get  # type: ignore
-
-    # Silence Weave logging
-    for name in logging.root.manager.loggerDict:
-        if name.startswith("weave"):
-            logging.getLogger(name).disabled = True
-
-    # Set dummy API key if missing
-    if not os.environ.get("WANDB_API_KEY"):
-        os.environ["WANDB_API_KEY"] = "dumped_api_key_for_weave_tracer"
-
-    # if needed in future tests, enable this and replace WF_TRACE_SERVER_URL to local server
-    # full_url = f"http://127.0.0.1:{_port}"
-    # os.environ["WF_TRACE_SERVER_URL"] = full_url
+    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.
-    """
-    try:
-        import weave
-        from weave.compat import wandb  # type: ignore
-    except ImportError:
-        logger.warning("Weave or wandb not installed; cannot uninstrument.")
-        return
+    """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
 
-    global _original_default_entity_name_getter
-    if _original_default_entity_name_getter is not None:
-        wandb.Api.default_entity_name = _original_default_entity_name_getter  # type: ignore
-        _original_default_entity_name_getter = None
-        logger.info("restored wandb.Api.default_entity_name")
+    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.")
 
-    global _original_upsert_project_getter
-    if _original_upsert_project_getter is not None:
-        wandb.Api.upsert_project = _original_upsert_project_getter  # type: ignore
-        _original_upsert_project_getter = None
-        logger.info("restored wandb.Api.upsert_project")
+    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.")
 
-    global _original_weave_post
-    if _original_weave_post is not None:
-        weave.utils.http_requests.session.post = _original_weave_post  # type: ignore
-        _original_weave_post = None
-        logger.info("restored weave.utils.http_requests.session.post")
-
-    global _original_weave_get
-    if _original_weave_get is not None:
-        weave.utils.http_requests.session.get = _original_weave_get  # type: ignore
-        _original_weave_get = None
-        logger.info("restored weave.utils.http_requests.session.get")
-
-    # Restore Weave logging
-    for name in logging.root.manager.loggerDict:
-        if name.startswith("weave"):
-            logging.getLogger(name).disabled = False
+    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/litagent.py

@@ -198,6 +198,7 @@ class LitAgent(Generic[T]):
             * `float` representing the final reward.
             * `List[ReadableSpan]` with OpenTelemetry spans.
             * `List[Span]` with Agent Lightning spans.
+            * `List[SpanCoreFields]` with Agent Lightning spans.
         """
         raise NotImplementedError("Agents must implement the `rollout` method.")
 

agentlightning/runner/agent.py

@@ -43,6 +43,7 @@ from agentlightning.types import (
     RolloutMode,
     RolloutRawResult,
     Span,
+    SpanCoreFields,
 )
 from agentlightning.utils.system_snapshot import system_snapshot
 
@@ -74,7 +75,8 @@ class LitAgentRunner(Runner[T_task]):
         poll_interval: float = 5.0,
         heartbeat_interval: float = 10.0,
         interval_jitter: float = 0.5,
-        heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
+        heartbeat_launch_mode: Literal["asyncio", "thread"] = "thread",
+        heartbeat_include_gpu: bool = False,
     ) -> None:
         """Initialize the agent runner.
 
@@ -88,7 +90,10 @@ class LitAgentRunner(Runner[T_task]):
                 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".
-                "asyncio" is the default and recommended mode. Use "thread" if you are experiencing blocking coroutines.
+                "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
@@ -97,6 +102,7 @@ class LitAgentRunner(Runner[T_task]):
         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
@@ -276,7 +282,7 @@ class LitAgentRunner(Runner[T_task]):
         """
         store = self.get_store()
 
-        trace_spans: list[ReadableSpan] | list[Span] = []
+        trace_spans: list[Span] = []
         result_recognized: bool = False
 
         # Case 0: result is None
@@ -295,31 +301,38 @@ class LitAgentRunner(Runner[T_task]):
             # Preserve the existing spans before another span is emitted
             trace_spans = list(self._tracer.get_last_trace())
             # This will NOT emit another span to the tracer
-            reward_span = emit_reward(raw_result, propagate=False)
+            reward_span_core_fields = emit_reward(raw_result, propagate=False)
             # We add it to the store manually
-            await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
-            trace_spans.append(reward_span)
+            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-3: result is a list
+        # 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, OtelTracer):
-                    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)
@@ -327,7 +340,25 @@ class LitAgentRunner(Runner[T_task]):
                 # 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
@@ -336,7 +367,7 @@ class LitAgentRunner(Runner[T_task]):
                     f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. "
                     "Please check your rollout implementation."
                 )
-                trace_spans = raw_result
+                trace_spans = []
                 result_recognized = True
 
             else:
@@ -355,14 +386,46 @@ class LitAgentRunner(Runner[T_task]):
         return trace_spans
 
     async def _emit_heartbeat(self, store: LightningStore) -> None:
-        """Send a heartbeat tick to the store."""
+        """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:
-            await store.update_worker(worker_id, system_snapshot())
+            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())
 
@@ -377,51 +440,161 @@ class LitAgentRunner(Runner[T_task]):
             return None
 
         if self._heartbeat_launch_mode == "asyncio":
-            stop_event = asyncio.Event()
-
-            async def heartbeat_loop() -> None:
-                while not stop_event.is_set():
-                    await self._emit_heartbeat(store)
-                    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
-
+            return self._start_heartbeat_asyncio_loop(store)
         if self._heartbeat_launch_mode == "thread":
-            stop_evt = threading.Event()
+            return self._start_heartbeat_thread_loop(store)
+        raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
 
-            def thread_worker() -> None:
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-                while not stop_evt.is_set():
-                    loop.run_until_complete(self._emit_heartbeat(store))
+    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)
-                    stop_evt.wait(interval)
+                    await asyncio.wait_for(stop_event.wait(), timeout=interval)
 
-            thread = threading.Thread(target=thread_worker, name=f"{self.get_worker_id()}-heartbeat", daemon=True)
-            thread.start()
+        task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat")
 
-            async def stop() -> None:
-                stop_evt.set()
-                await asyncio.to_thread(thread.join)
+        async def stop() -> None:
+            stop_event.set()
+            with suppress(asyncio.CancelledError):
+                await task
 
-            return stop
+        return stop
 
-        raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
+    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.
@@ -477,6 +650,8 @@ class LitAgentRunner(Runner[T_task]):
                 logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
                 return rollout_id
 
+        logger.debug(f"{self._log_prefix(rollout_id)} Resources fetched (id={resources_update.resources_id}).")
+
         trace_spans: List[ReadableSpan] | List[Span] = []
         has_exception: bool = False
 
@@ -484,9 +659,11 @@ class LitAgentRunner(Runner[T_task]):
             await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)
 
             start_time = time.time()
+            logger.debug(f"{self._log_prefix(rollout_id)} Prepared for trace context.")
             async with self._tracer.trace_context(
                 name=rollout_id, 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
                 )
@@ -498,21 +675,27 @@ class LitAgentRunner(Runner[T_task]):
                     rollout_method = (
                         agent.training_rollout_async if next_rollout.mode == "train" else agent.validation_rollout_async
                     )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Starting async rollout method.")
                     result = await rollout_method(
                         next_rollout.input, resources=resources_update.resources, rollout=next_rollout
                     )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Async rollout method completed.")
                 else:
                     rollout_method = (
                         agent.training_rollout if next_rollout.mode == "train" else agent.validation_rollout
                     )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Starting sync rollout method.")
                     result = rollout_method(
                         next_rollout.input, resources=resources_update.resources, rollout=next_rollout
                     )
+                    logger.debug(f"{self._log_prefix(rollout_id)} Sync rollout method completed.")
 
                 await self._trigger_hooks(
                     hook_type="on_trace_end", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
                 )
 
+            logger.debug(f"{self._log_prefix(rollout_id)} Trace context exited.")
+
             # Possible exceptions in post_process will be caught in the overall exception handler
             trace_spans = await self._post_process_rollout_result(next_rollout, result)
             last_reward = find_final_reward(trace_spans)

agentlightning/runner/legacy.py

@@ -12,7 +12,7 @@ from agentlightning.client import AgentLightningClient
 from agentlightning.litagent import LitAgent
 from agentlightning.litagent.litagent import is_v0_1_rollout_api
 from agentlightning.tracer.base import Tracer
-from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet
+from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Span, SpanLike, Triplet
 
 from .base import Runner
 
@@ -99,7 +99,7 @@ class LegacyAgentRunner(Runner[Any]):
         trace: Any = None
         final_reward: Optional[float] = None
         triplets: Optional[List[Triplet]] = None
-        trace_spans: Optional[List[ReadableSpan]] = None
+        trace_spans: Optional[List[SpanLike]] = None
 
         # Handle different types of results from the agent
         # Case 1: result is a float (final reward)
@@ -108,10 +108,14 @@ class LegacyAgentRunner(Runner[Any]):
         # Case 2: result is a list of Triplets
         if isinstance(result, list) and all(isinstance(t, Triplet) for t in result):
             triplets = result  # type: ignore
-        # Case 3: result is a list of ReadableSpan (OpenTelemetry spans)
-        if isinstance(result, list) and all(isinstance(t, ReadableSpan) for t in result):
+        # Case 3.1: result is a list of ReadableSpan (OpenTelemetry spans)
+        if isinstance(result, list) and all(isinstance(t, (ReadableSpan)) for t in result):
             trace_spans = result  # type: ignore
             trace = [json.loads(readable_span.to_json()) for readable_span in trace_spans]  # type: ignore
+        # Case 3.2: result is a list of Span (Agent-lightning spans)
+        if isinstance(result, list) and all(isinstance(t, Span) for t in result):
+            trace_spans = result  # type: ignore
+            trace = [span.model_dump() for span in trace_spans]  # type: ignore
         # Case 4: result is a list of dict (trace JSON)
         if isinstance(result, list) and all(isinstance(t, dict) for t in result):
             trace = result
@@ -123,10 +127,9 @@ class LegacyAgentRunner(Runner[Any]):
 
         # If the agent has tracing enabled, use the tracer's last trace if not already set
         if self.tracer and (trace is None or trace_spans is None):
-            spans = self.tracer.get_last_trace()
-            if spans:
-                trace = [json.loads(readable_span.to_json()) for readable_span in spans]
-                trace_spans = spans
+            trace_spans = self.tracer.get_last_trace()  # type: ignore
+            if trace_spans:
+                trace = [cast(Span, span).model_dump() for span in trace_spans]
 
         # Always extract triplets from the trace using TracerTraceToTriplet
         if trace_spans:

agentlightning/semconv.py

@@ -34,6 +34,9 @@ AGL_OPERATION = "agentlightning.operation"
 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.
 
@@ -53,6 +56,9 @@ class LightningResourceAttributes(Enum):
     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.

agentlightning/store/client_server.py

@@ -903,9 +903,11 @@ class LightningStoreServer(LightningStore):
             except asyncio.CancelledError:
                 # Client disconnected (Timeout)
                 status = 499  # Standard Nginx code for "Client Closed Request"
+                server_logger.debug(f"Client disconnected (Timeout): {request.url.path}", exc_info=True)
                 raise  # Re-raise to let Uvicorn handle the cleanup
             except Exception as exc:
                 status = resolve_error_type(exc)
+                server_logger.debug(f"Server error: {request.url.path}", exc_info=True)
                 raise
             finally:
                 # This block executes NO MATTER WHAT happens above
@@ -1518,7 +1520,7 @@ class LightningStoreClient(LightningStore):
             except aiohttp.ClientResponseError as cre:
                 # Respect app-level 4xx as final
                 # 4xx => application issue; do not retry (except 408 which is transient)
-                client_logger.debug(f"ClientResponseError: {cre.status} {cre.message}", exc_info=True)
+                client_logger.debug(f"ClientResponseError ({method} {path}): {cre.status} {cre.message}", exc_info=True)
                 if 400 <= cre.status < 500 and cre.status != 408:
                     raise
                 # 5xx and others will be retried below if they raise
@@ -1534,9 +1536,9 @@ class LightningStoreClient(LightningStore):
                 asyncio.TimeoutError,
             ) as net_exc:
                 # Network/session issue: probe health before retrying
-                client_logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
+                client_logger.debug(f"Network/session issue ({method} {path}): {net_exc}", exc_info=True)
                 last_exc = net_exc
-                client_logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
+                client_logger.info(f"Network/session issue: {net_exc} - will retry the request {method}: {path}")
                 if not await self._wait_until_healthy(session):
                     break  # server is not healthy, do not retry
 

agentlightning/store/collection/base.py

@@ -148,12 +148,12 @@ class TrackedCollection:
             yield
 
         else:
-            from agentlightning.store.collection_based import nearest_lightning_store_method_from_stack
+            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 = nearest_lightning_store_method_from_stack()
+            public_store_method, private_store_method = get_current_store_methods()
             try:
                 yield
             except BaseException as exc:

agentlightning/store/collection/mongo.py

@@ -663,6 +663,9 @@ class MongoBasedCollection(Collection[T_model]):
                     return_document=ReturnDocument.AFTER,
                 )
 
+            if result_doc is None:  # pyright: ignore[reportUnnecessaryComparison]
+                raise RuntimeError(f"Upsert resulted in no document for filter: {pk_filter}")
+
             # Because upsert=True, result_doc is guaranteed to be not None
             new_item = self._model_validate_item(result_doc)
             upserted_items.append(new_item)
@@ -1077,7 +1080,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
 class MongoLightningCollections(LightningCollections):
     """Mongo implementation of LightningCollections using MongoDB collections.
 
-    Serves as the storage base for [`MongoLightningStore`][agentlightning.store.MongoLightningStore].
+    Serves as the storage base for [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore].
     """
 
     def __init__(

agentlightning/store/collection_based.py

@@ -15,13 +15,11 @@ from __future__ import annotations
 
 import asyncio
 import functools
-import hashlib
-import inspect
 import logging
 import time
-import uuid
 import warnings
 from collections import defaultdict
+from contextvars import ContextVar
 from types import CoroutineType
 from typing import (
     Any,
@@ -61,6 +59,7 @@ from agentlightning.types import (
     Worker,
     WorkerStatus,
 )
+from agentlightning.utils.id import generate_id
 from agentlightning.utils.metrics import MetricsBackend
 
 from .base import (
@@ -88,6 +87,12 @@ SelfT = TypeVar("SelfT", bound="CollectionBasedLightningStore[Any]")
 
 logger = logging.getLogger(__name__)
 
+# ContextVars for tracking the current store method without expensive stack introspection.
+# These are set by the @tracked decorator and read by tracking_context in collection/base.py.
+_UNKNOWN_STORE_METHOD = "unknown"
+_current_public_store_method: ContextVar[str] = ContextVar("public_store_method", default=_UNKNOWN_STORE_METHOD)
+_current_private_store_method: ContextVar[str] = ContextVar("private_store_method", default=_UNKNOWN_STORE_METHOD)
+
 
 def _with_collections_execute(labels: Sequence[AtomicLabels]):
     """Hands over the function execution to the collections.execute method.
@@ -125,38 +130,47 @@ def tracked(name: str):
 
         @functools.wraps(func)
         async def wrapper(self: CollectionBasedLightningStore[T_collections], *args: Any, **kwargs: Any) -> Any:
-            # Backtracking where this method comes from
-            public_meth_in_stack, _ = nearest_lightning_store_method_from_stack()
+            # Get the current public method from ContextVar (set by outer tracked methods)
+            public_meth_in_stack = _current_public_store_method.get()
 
-            # For backtracking in collection methods.
-            # Only track the public methods (+healthcheck)
+            # Set ContextVars for nested calls to read. Use tokens for proper cleanup.
+            pub_token = None
+            priv_token = None
             if name in COLLECTION_STORE_PUBLIC_METHODS:
-                public_method_name = name  # pyright: ignore[reportUnusedVariable]
+                pub_token = _current_public_store_method.set(name)
                 public_meth_in_stack = name  # We are in a public method already.
             if name in COLLECTION_STORE_ALL_METHODS:
-                private_method_name = name  # pyright: ignore[reportUnusedVariable]
+                priv_token = _current_private_store_method.set(name)
 
-            if self._tracker is None:  # pyright: ignore[reportPrivateUsage]
-                # Skip the tracking because tracking is not configured
-                return await func(self, *args, **kwargs)
-
-            start_time = time.perf_counter()
-            status: str = "OK"
             try:
-                return await func(self, *args, **kwargs)
-            except BaseException as exc:
-                status = exc.__class__.__name__
-                raise
+                if self._tracker is None:  # pyright: ignore[reportPrivateUsage]
+                    # Skip the tracking because tracking is not configured
+                    return await func(self, *args, **kwargs)
+
+                start_time = time.perf_counter()
+                status: str = "OK"
+                try:
+                    return await func(self, *args, **kwargs)
+                except BaseException as exc:
+                    status = exc.__class__.__name__
+                    raise
+                finally:
+                    elapsed = time.perf_counter() - start_time
+                    await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
+                        "agl.store.total",
+                        labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
+                    )
+                    await self._tracker.observe_histogram(  # pyright: ignore[reportPrivateUsage]
+                        "agl.store.latency",
+                        value=elapsed,
+                        labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
+                    )
             finally:
-                elapsed = time.perf_counter() - start_time
-                await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
-                    "agl.store.total", labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status}
-                )
-                await self._tracker.observe_histogram(  # pyright: ignore[reportPrivateUsage]
-                    "agl.store.latency",
-                    value=elapsed,
-                    labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
-                )
+                # Reset ContextVars to their previous values
+                if pub_token is not None:
+                    _current_public_store_method.reset(pub_token)
+                if priv_token is not None:
+                    _current_private_store_method.reset(priv_token)
 
         return cast(T_callable, wrapper)
 
@@ -195,19 +209,16 @@ def healthcheck_before(func: T_callable) -> T_callable:
 
 
 def _generate_resources_id() -> str:
-    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
-    return "rs-" + short_id
+    return "rs-" + generate_id(12)
 
 
 def _generate_rollout_id() -> str:
-    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
-    return "ro-" + short_id
+    return "ro-" + generate_id(12)
 
 
 def _generate_attempt_id() -> str:
     """We don't need that long because attempts are limited to rollouts."""
-    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:8]
-    return "at-" + short_id
+    return "at-" + generate_id(8)
 
 
 class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
@@ -1752,41 +1763,14 @@ COLLECTION_STORE_PUBLIC_METHODS = frozenset(
 
 COLLECTION_STORE_ALL_METHODS = frozenset([name for name in CollectionBasedLightningStore.__dict__])
 
-_UNKNOWN_STORE_METHOD = "unknown"
 
+def get_current_store_methods() -> Tuple[str, str]:
+    """Get the current store method names from ContextVars.
 
-def nearest_lightning_store_method_from_stack() -> Tuple[str, str]:
-    """Stack introspection so that we capture the nearest public API method from the
-    call stack whenever metrics are recorded.
+    This is a fast O(1) replacement for stack introspection. The ContextVars are
+    set by the @tracked decorator when entering store methods.
 
     Returns:
-        A tuple of public method name and nearest private method name.
+        A tuple of (public_method_name, private_method_name).
     """
-    frame = inspect.currentframe()
-    final_public_method_name = final_private_method_name = _UNKNOWN_STORE_METHOD
-    try:
-        if frame is not None:
-            frame = frame.f_back
-            while frame is not None:
-                self_obj = frame.f_locals.get("self")
-                public_method_name = frame.f_locals.get("public_method_name")
-                private_method_name = frame.f_locals.get("private_method_name")
-                if (
-                    final_public_method_name == _UNKNOWN_STORE_METHOD
-                    and public_method_name in COLLECTION_STORE_PUBLIC_METHODS
-                    and isinstance(self_obj, LightningStore)
-                ):
-                    final_public_method_name = public_method_name
-                if (
-                    final_private_method_name == _UNKNOWN_STORE_METHOD
-                    and private_method_name in COLLECTION_STORE_ALL_METHODS
-                    and isinstance(self_obj, LightningStore)
-                ):
-                    final_private_method_name = private_method_name
-                frame = frame.f_back
-    except Exception as exc:
-        logger.debug("Error during stack introspection for LightningStore method: %s", exc)
-    finally:
-        del frame
-
-    return final_public_method_name, final_private_method_name
+    return _current_public_store_method.get(), _current_private_store_method.get()

agentlightning/store/mongo.py

@@ -33,7 +33,7 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
     Args:
         mongo_uri: MongoDB connection string (defaults to local replica set).
         mongo_client_kwargs: Extra keyword arguments forwarded to `AsyncMongoClient`.
-        database: The MongoDB database name. Defaults to ``agentlightning``.
+        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.

agentlightning/tracer/__init__.py

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

agentlightning/tracer/agentops.py

@@ -13,12 +13,13 @@ import agentops.sdk.core
 import opentelemetry.trace as trace_api
 from agentops.sdk.core import TracingCore
 from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
-from opentelemetry.trace import get_tracer_provider
 from opentelemetry.trace.status import StatusCode
 
 from agentlightning.instrumentation import instrument_all, uninstrument_all
 from agentlightning.store.base import LightningStore
+from agentlightning.utils.otel import get_span_processors, get_tracer_provider
 
+from .base import with_active_tracer_context
 from .otel import LightningSpanProcessor, OtelTracer
 
 if TYPE_CHECKING:
@@ -79,13 +80,20 @@ class AgentOpsTracer(OtelTracer):
                 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()
-
-        # TODO: The span processor cannot be deleted once added.
-        # This might be a problem if the tracer is entered and exited multiple times.
-        self._get_tracer_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)
@@ -94,6 +102,10 @@ class AgentOpsTracer(OtelTracer):
             self.uninstrument(worker_id)
             logger.info(f"[Worker {worker_id}] Instrumentation removed.")
 
+        # NOTE: The teardown doesn't try to remove the LightningSpanProcessor from the TracerProvider.
+        # Currently there is no stable way to fully restore the AgentOps state to the initial state.
+
+    @with_active_tracer_context
     @asynccontextmanager
     async def trace_context(
         self,
@@ -158,7 +170,6 @@ class AgentOpsTracer(OtelTracer):
                 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:
-            # TODO: Add tests to cover both paths
             self._disable_native_otlp_exporter()
             with self._lightning_span_processor:
                 with self._agentops_trace_context(None, None, kwargs):

agentlightning/tracer/base.py

@@ -2,14 +2,13 @@
 
 from __future__ import annotations
 
+import functools
 import logging
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional
-
-from opentelemetry.sdk.trace import ReadableSpan
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional, TypeVar
 
 from agentlightning.store.base import LightningStore
-from agentlightning.types import ParallelWorkerBase
+from agentlightning.types import Attributes, ParallelWorkerBase, Span, SpanCoreFields, SpanRecordingContext, TraceStatus
 
 if TYPE_CHECKING:
     from langchain_core.callbacks.base import BaseCallbackHandler  # type: ignore
@@ -17,6 +16,14 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T")
+
+
+_active_tracer: Optional[Tracer] = None
+
+T_func = Callable[..., Awaitable[Any]]
+
+
 class Tracer(ParallelWorkerBase):
     """
     An abstract base class for tracers.
@@ -98,12 +105,12 @@ class Tracer(ParallelWorkerBase):
         """Internal API for CI backward compatibility."""
         raise NotImplementedError()
 
-    def get_last_trace(self) -> List[ReadableSpan]:
+    def get_last_trace(self) -> List[Span]:
         """
         Retrieves the raw list of captured spans from the most recent trace.
 
         Returns:
-            A list of OpenTelemetry `ReadableSpan` objects.
+            A list of [`Span`][agentlightning.Span] objects collected during the last trace.
         """
         raise NotImplementedError()
 
@@ -124,6 +131,48 @@ class Tracer(ParallelWorkerBase):
         with self._trace_context_sync(name=func.__name__):
             return func(*args, **kwargs)
 
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        """Notify the tracer that a span should be created here.
+
+        It uses a fire-and-forget approach and doesn't wait for the span to be created.
+
+        Args:
+            name: The name of the span.
+            attributes: The attributes of the span.
+            timestamp: The timestamp of the span.
+            status: The status of the span.
+
+        Returns:
+            The core fields of the span.
+        """
+        raise NotImplementedError()
+
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> ContextManager[SpanRecordingContext]:
+        """Start to record an operation to a span.
+
+        Args:
+            name: The name of the operation.
+            attributes: The attributes of the operation.
+            start_time: The start time of the operation.
+            end_time: The end time of the operation.
+
+        Returns:
+            A [`SpanRecordingContext`][agentlightning.SpanRecordingContext] for recording the operation on the span.
+        """
+        raise NotImplementedError()
+
     async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single asynchronous function.
@@ -175,3 +224,64 @@ class Tracer(ParallelWorkerBase):
                 self.teardown_worker(0)
             if has_init:
                 self.teardown()
+
+
+def set_active_tracer(tracer: Tracer):
+    """Set the active tracer for the current process.
+
+    Args:
+        tracer: The tracer to set as active.
+    """
+    global _active_tracer
+    if _active_tracer is not None:
+        raise ValueError("An active tracer is already set. Cannot set a new one.")
+    _active_tracer = tracer
+
+
+def clear_active_tracer():
+    """Clear the active tracer for the current process."""
+    global _active_tracer
+    _active_tracer = None
+
+
+def get_active_tracer() -> Optional[Tracer]:
+    """Get the active tracer for the current process.
+
+    Returns:
+        The active tracer, or None if no tracer is active.
+    """
+    global _active_tracer
+    return _active_tracer
+
+
+class _ActiveTracerAsyncCM(AsyncContextManager[T]):
+    def __init__(self, tracer: Tracer, inner: AsyncContextManager[T]):
+        self._tracer = tracer
+        self._inner = inner
+
+    async def __aenter__(self) -> T:
+        set_active_tracer(self._tracer)  # will raise if nested
+        try:
+            return await self._inner.__aenter__()
+        except Exception:
+            clear_active_tracer()
+            raise
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> Optional[bool]:
+        try:
+            return await self._inner.__aexit__(*args, **kwargs)
+        finally:
+            clear_active_tracer()
+
+
+def with_active_tracer_context(
+    func: Callable[..., AsyncContextManager[T]],
+) -> Callable[..., AsyncContextManager[T]]:
+    """Decorate a method returning an AsyncContextManager so tracer is active for the whole `async with`."""
+
+    @functools.wraps(func)
+    def wrapper(self: Tracer, *args: Any, **kwargs: Any) -> AsyncContextManager[T]:
+        cm = func(self, *args, **kwargs)
+        return _ActiveTracerAsyncCM(self, cm)
+
+    return wrapper

agentlightning/tracer/dummy.py

@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import time
+from contextlib import contextmanager
+from typing import (
+    Iterator,
+    Optional,
+)
+
+from agentlightning.types import (
+    Attributes,
+    SpanCoreFields,
+    SpanRecordingContext,
+    StatusCode,
+    TraceStatus,
+)
+from agentlightning.utils.otel import format_exception_attributes
+
+from .base import Tracer
+
+logger = logging.getLogger(__name__)
+
+
+class DummySpanRecordingContext(SpanRecordingContext):
+    """Context for recording operations on a dummy span, not dependent on any backend tracer."""
+
+    def __init__(self, name: str, attributes: Optional[Attributes] = None, start_time: Optional[float] = None) -> None:
+        self.name = name
+        self.attributes = attributes or {}
+        self.start_time = start_time or time.time()
+        self.end_time = None
+        self.status = TraceStatus(status_code="OK")
+
+    def record_exception(self, exception: BaseException) -> None:
+        self.record_status("ERROR", str(exception))
+        self.record_attributes(format_exception_attributes(exception))
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        self.attributes.update(attributes)
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        self.status = TraceStatus(status_code=status_code, description=description)
+
+    def finalize(self, end_time: Optional[float] = None) -> None:
+        self.end_time = end_time or time.time()
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        if self.end_time is None:
+            raise ValueError("End time is not set. Call finalize() first.")
+        return SpanCoreFields(
+            name=self.name,
+            attributes=self.attributes,
+            start_time=self.start_time,
+            end_time=self.end_time,
+            status=self.status,
+        )
+
+
+class DummyTracer(Tracer):
+    """A dummy tracer that does not trace anything, but it is compatible with the emitter API.
+
+    It doesn't rely on any backend tracer, and also doesn't use any stores.
+    """
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        if attributes is None:
+            attributes = {}
+        if timestamp is None:
+            timestamp = time.time()
+        if status is None:
+            status = TraceStatus(status_code="OK")
+        return SpanCoreFields(
+            name=name,
+            attributes=attributes,
+            start_time=timestamp,
+            end_time=timestamp,
+            status=status,
+        )
+
+    @contextmanager
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> Iterator[DummySpanRecordingContext]:
+        start_time = start_time or time.time()
+        recording_context = DummySpanRecordingContext(name, attributes, start_time)
+        try:
+            yield recording_context
+        except Exception as exc:
+            recording_context.record_exception(exc)
+            recording_context.record_status("ERROR", str(exc))
+            raise
+        finally:
+            recording_context.finalize(end_time)

agentlightning/tracer/otel.py

@@ -6,27 +6,69 @@ import asyncio
 import logging
 import threading
 import warnings
-from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Awaitable, List, Optional
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, Awaitable, Iterator, List, Optional
 
 import opentelemetry.trace as trace_api
-from agentops.sdk.core import BatchSpanProcessor
 from opentelemetry.instrumentation.utils import suppress_instrumentation
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
-from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
-from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
 
 from agentlightning.semconv import LightningResourceAttributes
 from agentlightning.store.base import LightningStore
+from agentlightning.types import Attributes, Span, SpanCoreFields, SpanRecordingContext, StatusCode, TraceStatus
+from agentlightning.types.tracer import convert_timestamp
 from agentlightning.utils.otel import get_tracer_provider
 from agentlightning.utils.otlp import LightningStoreOTLPExporter
 
-from .base import Tracer
+from .base import Tracer, with_active_tracer_context
 
 logger = logging.getLogger(__name__)
 
+STORE_WRITE_TIMEOUT_SECONDS = 10.0
+
+
+def to_otel_status_code(status_code: StatusCode) -> trace_api.StatusCode:
+    if status_code == "UNSET":
+        return trace_api.StatusCode.UNSET
+    elif status_code == "ERROR":
+        return trace_api.StatusCode.ERROR
+    else:
+        return trace_api.StatusCode.OK
+
+
+class OtelSpanRecordingContext(SpanRecordingContext):
+    def __init__(self, span: trace_api.Span) -> None:
+        self._span = span
+
+    def record_exception(self, exception: BaseException) -> None:
+        self._span.record_exception(exception)
+        self.record_status("ERROR", str(exception))
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        self._span.set_attributes(attributes)
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        otel_status_code = to_otel_status_code(status_code)
+        self._span.set_status(otel_status_code, description)
+
+    def get_otel_span(self) -> trace_api.Span:
+        return self._span
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        if isinstance(self._span, ReadableSpan):
+            return SpanCoreFields(
+                name=self._span.name,
+                attributes=dict(self._span.attributes) if self._span.attributes else {},
+                start_time=convert_timestamp(self._span.start_time),
+                end_time=convert_timestamp(self._span.end_time),
+                status=TraceStatus.from_opentelemetry(self._span.status),
+            )
+        else:
+            raise ValueError(f"Span is not a ReadableSpan: {self._span}")
+
 
 class OtelTracer(Tracer):
     """Tracer that provides a basic OpenTelemetry tracer provider.
@@ -38,7 +80,7 @@ class OtelTracer(Tracer):
     def __init__(self):
         super().__init__()
         # This provider is only initialized when the worker is initialized.
-        self._tracer_provider: Optional[TracerProvider] = None
+        self._tracer_provider: Optional[trace_api.TracerProvider] = None
         self._lightning_span_processor: Optional[LightningSpanProcessor] = None
         self._simple_span_processor: Optional[SimpleSpanProcessor] = None
         self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
@@ -63,7 +105,7 @@ class OtelTracer(Tracer):
         except RuntimeError:
             logger.debug(f"[Worker {worker_id}] Tracer provider is not initialized by OtelTracer. Initializing it now.")
 
-        self._tracer_provider = TracerProvider()
+        self._tracer_provider = TracerProviderImpl()
         trace_api.set_tracer_provider(self._tracer_provider)
         self._lightning_span_processor = LightningSpanProcessor()
         self._tracer_provider.add_span_processor(self._lightning_span_processor)
@@ -78,6 +120,7 @@ class OtelTracer(Tracer):
         super().teardown_worker(worker_id)
         logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer does NOT remove the tracer provider.")
 
+    @with_active_tracer_context
     @asynccontextmanager
     async def trace_context(
         self,
@@ -129,12 +172,69 @@ class OtelTracer(Tracer):
         else:
             raise ValueError("rollout_id and attempt_id must be either all provided or all None")
 
-    def get_last_trace(self) -> List[ReadableSpan]:
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        # Fire the span to the current active tracer provider.
+        tracer_provider = self._get_tracer_provider()
+        tracer = tracer_provider.get_tracer(__name__)
+        span = tracer.start_span(
+            name, attributes=attributes, start_time=int(timestamp * 1_000_000_000) if timestamp else None
+        )
+        if status is not None:
+            span.set_status(to_otel_status_code(status.status_code), status.description)
+        span.end(int(timestamp * 1_000_000_000) if timestamp else None)
+
+        # The span should have been auto-created by now.
+        # Return the core fields of the span.
+        if isinstance(span, ReadableSpan):
+            return SpanCoreFields(
+                name=name,
+                attributes=dict(span.attributes) if span.attributes else {},
+                start_time=convert_timestamp(span.start_time),
+                end_time=convert_timestamp(span.end_time),
+                status=TraceStatus.from_opentelemetry(span.status),
+            )
+        else:
+            raise ValueError(f"Span is not a ReadableSpan: {span}")
+
+    @contextmanager
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> Iterator[SpanRecordingContext]:
+        if end_time is not None:
+            logger.warning("OpenTelemetry doesn't support customizing the end time of a span. End time is ignored.")
+        # Record the span to the current active tracer provider.
+        tracer_provider = self._get_tracer_provider()
+        tracer = tracer_provider.get_tracer(__name__)
+
+        # Activate the span as the current span within otel.
+        with tracer.start_as_current_span(
+            name, attributes=attributes, start_time=int(start_time * 1_000_000_000) if start_time else None
+        ) as span:
+            recording_context = OtelSpanRecordingContext(span)
+            try:
+                yield recording_context
+            except Exception as exc:
+                recording_context.record_exception(exc)
+                raise
+
+        # No need to retrieve the span here. It's already been sent to otel processor.
+
+    def get_last_trace(self) -> List[Span]:
         """
         Retrieves the raw list of captured spans from the most recent trace.
 
         Returns:
-            A list of OpenTelemetry `ReadableSpan` objects.
+            A list of [`Span`][agentlightning.Span] objects captured during the most recent trace.
         """
         if not self._lightning_span_processor:
             raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
@@ -143,6 +243,8 @@ class OtelTracer(Tracer):
     def _get_tracer_provider(self) -> TracerProviderImpl:
         if self._tracer_provider is None:
             raise RuntimeError("TracerProvider is not initialized. Call init_worker() first.")
+        if not isinstance(self._tracer_provider, TracerProviderImpl):
+            raise TypeError(f"TracerProvider is not a opentelemetry.sdk.trace.TracerProvider: {self._tracer_provider}")
         return self._tracer_provider
 
     def _enable_native_otlp_exporter(self, store: LightningStore, rollout_id: str, attempt_id: str):
@@ -215,18 +317,20 @@ class LightningSpanProcessor(SpanProcessor):
 
     def __init__(self, disable_store_submission: bool = False):
         self._disable_store_submission: bool = disable_store_submission
-        self._spans: List[ReadableSpan] = []
+        self._spans: List[Span] = []
 
         # Store related context and states
         self._store: Optional[LightningStore] = None
         self._rollout_id: Optional[str] = None
         self._attempt_id: Optional[str] = None
+        self._local_sequence_id: int = 0
         self._lock = threading.Lock()
 
         # private asyncio loop running in a daemon thread
         self._loop_ready = threading.Event()
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._loop_thread: Optional[threading.Thread] = None
+        self._loop_init_lock = threading.Lock()
 
     def __repr__(self) -> str:
         return (
@@ -262,11 +366,19 @@ class LightningSpanProcessor(SpanProcessor):
         self._disable_store_submission = value
 
     def _ensure_loop(self) -> None:
-        if self._loop_thread is None or self._loop is None:
+        # Fast path: loop already initialized
+        if self._loop_thread is not None and self._loop is not None:
+            return
+
+        with self._loop_init_lock:
+            # Double-check after acquiring lock
+            if self._loop_thread is not None and self._loop is not None:
+                return
             self._loop_ready.clear()
             self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
             self._loop_thread.start()
-            self._loop_ready.wait()  # loop is ready
+            if not self._loop_ready.wait(timeout=30.0):
+                raise RuntimeError("Timed out waiting for otel-loop thread to start")
 
     def _loop_runner(self):
         loop = asyncio.new_event_loop()
@@ -330,13 +442,13 @@ class LightningSpanProcessor(SpanProcessor):
     def force_flush(self, timeout_millis: int = 30000) -> bool:
         return True
 
-    def spans(self) -> List[ReadableSpan]:
+    def spans(self) -> List[Span]:
         """
         Get the list of spans collected by this processor.
         This is useful for debugging and testing purposes.
 
         Returns:
-            List of ReadableSpan objects collected during tracing.
+            List of [`Span`][agentlightning.Span] objects collected during tracing.
         """
         return self._spans
 
@@ -373,12 +485,46 @@ class LightningSpanProcessor(SpanProcessor):
                 # Submit add_otel_span to the event loop and wait for it to complete
                 with suppress_instrumentation():
                     self._ensure_loop()
-                    self._await_in_loop(
+                    uploaded_span = self._await_in_loop(
                         self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
-                        timeout=60.0,
+                        timeout=STORE_WRITE_TIMEOUT_SECONDS,
                     )
+                    if uploaded_span is not None:
+                        self._spans.append(uploaded_span)
+            except TimeoutError:
+                logger.warning(
+                    "Timed out adding span %s to store after %.1f seconds. The span will be stored locally "
+                    "but it's not guaranteed to be persisted.",
+                    span.name,
+                    STORE_WRITE_TIMEOUT_SECONDS,
+                )
+                self._spans.append(
+                    Span.from_opentelemetry(
+                        span,
+                        rollout_id=self._rollout_id,
+                        attempt_id=self._attempt_id,
+                        sequence_id=self._local_sequence_id,
+                    )
+                )
             except Exception:
                 # log; on_end MUST NOT raise
-                logger.exception(f"Error adding span to store: {span.name}")
+                logger.exception(f"Error adding span to store: {span.name}. The span will be store locally only.")
+                self._spans.append(
+                    Span.from_opentelemetry(
+                        span,
+                        rollout_id=self._rollout_id,
+                        attempt_id=self._attempt_id,
+                        sequence_id=self._local_sequence_id,
+                    )
+                )
 
-        self._spans.append(span)
+        else:
+            # Fallback path
+            created_span = Span.from_opentelemetry(
+                span,
+                rollout_id=self._rollout_id or "rollout-dummy",
+                attempt_id=self._attempt_id or "attempt-dummy",
+                sequence_id=self._local_sequence_id,
+            )
+            self._local_sequence_id += 1
+            self._spans.append(created_span)

agentlightning/tracer/weave.py

@@ -2,63 +2,260 @@
 
 from __future__ import annotations
 
+import asyncio
+import concurrent.futures as futures
 import logging
 import os
-from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Tuple, Union
+import re
+import weakref
+from contextlib import asynccontextmanager, contextmanager
+from datetime import datetime
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    cast,
+)
 
-from agentlightning.instrumentation import instrument_weave, uninstrument_weave
+import weave
+from opentelemetry.semconv.attributes import exception_attributes
+from weave.trace.call import Call
+from weave.trace.settings import UserSettings
+from weave.trace.weave_client import WeaveClient
+from weave.trace_server import trace_server_interface as tsi
+from weave.wandb_interface.context import set_wandb_api_context
+
+from agentlightning.instrumentation.weave import InMemoryWeaveTraceServer, instrument_weave, uninstrument_weave
+from agentlightning.semconv import LightningResourceAttributes, LightningSpanAttributes
 from agentlightning.store.base import LightningStore
-from agentlightning.types.tracer import OtelResource, Span, SpanContext, TraceStatus
+from agentlightning.types import (
+    Attributes,
+    OtelResource,
+    Span,
+    SpanContext,
+    SpanCoreFields,
+    SpanRecordingContext,
+    StatusCode,
+    TraceStatus,
+)
+from agentlightning.utils.id import generate_id
+from agentlightning.utils.otel import (
+    filter_and_unflatten_attributes,
+    flatten_attributes,
+    format_exception_attributes,
+    sanitize_attributes,
+)
 
-from .base import Tracer
-
-if TYPE_CHECKING:
-    from weave.trace.call import Call  # type: ignore
-
-JSONPrimitive = Union[str, int, float, bool, None]
+from .base import Tracer, with_active_tracer_context
 
 logger = logging.getLogger(__name__)
 
 
-class WeaveTracer(Tracer):
+def op_name_to_func_name(op_name: str) -> str:
+    """Convert a Weave operation name to a function name.
+
+    Weave operation names look like this: `weave:///xxx/agentlightning.tracer.weave/op/openai.chat.completions.create:019b10be-...-44d74272569c`
     """
-    Tracer implementation using Weave for telemetry and trace logging.
+    match = re.search(r"/([^/:]+):", op_name)
+    if match:
+        return match.group(1)
+    else:
+        return op_name
+
+
+def random_project_name() -> str:
+    return "agl/weave-" + generate_id(12)
+
+
+def get_timestamp_or_throw(date: Optional[datetime], field_name: str) -> float:
+    if date is None:
+        raise ValueError(f"{field_name} is required but not set")
+    return date.timestamp()
+
+
+class WeaveSpanRecordingContext(SpanRecordingContext):
+    """Universal interface for recording operations on a Weave call."""
+
+    def __init__(self, call: Call) -> None:
+        self._call = call
+
+    def record_exception(self, exception: BaseException) -> None:
+        self._call.exception = str(exception)
+        self.record_status("ERROR", str(exception))
+        self.record_attributes(format_exception_attributes(exception))
+
+    def _get_input_from_attributes(self, attributes: Attributes) -> Dict[str, Any]:
+        if LightningSpanAttributes.OPERATION_INPUT.value in attributes:
+            # This can be a very rare case. If it happens, we can just let it throw.
+            return cast(Dict[str, Any], attributes[LightningSpanAttributes.OPERATION_INPUT.value])
+        else:
+            filtered_attributes = filter_and_unflatten_attributes(
+                attributes, LightningSpanAttributes.OPERATION_INPUT.value
+            )
+            if isinstance(filtered_attributes, list):
+                return {str(i): v for i, v in enumerate(filtered_attributes)}
+            else:
+                return filtered_attributes
+
+    def _get_output_from_attributes(self, attributes: Attributes) -> Any:
+        if LightningSpanAttributes.OPERATION_OUTPUT.value in attributes:
+            return attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]
+        else:
+            return filter_and_unflatten_attributes(attributes, LightningSpanAttributes.OPERATION_OUTPUT.value)
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        input_attributes = self._get_input_from_attributes(attributes)
+        if input_attributes:
+            self._call.inputs.update(input_attributes)
+
+        output_attributes = self._get_output_from_attributes(attributes)
+        if output_attributes:
+            if self._call.output is not None:
+                logger.warning(f"Output is already set. It will be overridden: {self._call.output}")
+            self._call.output = output_attributes
+
+        if LightningSpanAttributes.OPERATION_NAME.value in attributes:
+            logger.error(
+                f"Cannot record operation name as an attribute. It will be skipped: {attributes[LightningSpanAttributes.OPERATION_NAME.value]}"
+            )
+
+        # The rest of the attributes are recorded as summary.
+        for key, value in attributes.items():
+            if (
+                not key == LightningSpanAttributes.OPERATION_INPUT.value
+                and not key.startswith(LightningSpanAttributes.OPERATION_INPUT.value + ".")
+                and not key == LightningSpanAttributes.OPERATION_OUTPUT.value
+                and not key.startswith(LightningSpanAttributes.OPERATION_OUTPUT.value + ".")
+                and not key == LightningSpanAttributes.OPERATION_NAME.value
+            ):
+                if self._call.summary is None:
+                    self._call.summary = {}
+                self._call.summary[key] = value
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        if status_code == "ERROR":
+            if not description:
+                raise ValueError("Description is required when status code is ERROR")
+            self._call.exception = description
+        elif status_code == "OK":
+            self._call.exception = None
+        # Do nothing for other status codes.
+
+    def finalize(self) -> None:
+        # Do nothing
+        pass
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        return SpanCoreFields(
+            name=self._call.op_name,
+            attributes=flatten_attributes(self._call.attributes or {}),
+            start_time=self._call.started_at.timestamp() if self._call.started_at else None,
+            end_time=self._call.ended_at.timestamp() if self._call.ended_at else None,
+            status=TraceStatus(
+                status_code="OK" if self._call.exception is None else "ERROR", description=self._call.exception
+            ),
+        )
+
+
+class WeaveTracerManagedTraceServer(InMemoryWeaveTraceServer):
+    """A managed trace server for WeaveTracer."""
+
+    def __init__(
+        self,
+        partial_call_callback: Callable[[Dict[str, Any]], None],
+        complete_call_callback: Callable[[tsi.CallSchema], None],
+    ):
+        super().__init__()
+        self.partial_call_callback = partial_call_callback
+        self.complete_call_callback = complete_call_callback
+        self._calls_already_invoked: set[str] = set()
+
+    def trigger_callbacks(self, call_id: str) -> None:
+        with self._call_threading_lock:
+            if call_id in self.calls:
+                if call_id not in self._calls_already_invoked:
+                    self._calls_already_invoked.add(call_id)
+                    self.complete_call_callback(self.calls[call_id])
+                else:
+                    logger.info(f"Call {call_id} has callback already invoked. Skipping.")
+            elif call_id in self.partial_calls:
+                self.partial_call_callback(self.partial_calls[call_id])
+            else:
+                logger.error(f"Call {call_id} not found in partial_calls or calls")
+
+    def call_start(self, req: tsi.CallStartReq) -> tsi.CallStartRes:
+        try:
+            ret = super().call_start(req)
+            self.trigger_callbacks(ret.id)
+            return ret
+        except Exception:
+            logger.exception(f"Error calling call_start: {req}", exc_info=True)
+            raise
+
+    def call_end(self, req: tsi.CallEndReq) -> tsi.CallEndRes:
+        try:
+            ret = super().call_end(req)
+            self.trigger_callbacks(req.end.id)
+            return ret
+        except Exception:
+            logger.exception(f"Error calling call_end: {req}", exc_info=True)
+            raise
+
+    def clear(self) -> None:
+        self._calls_already_invoked.clear()
+
+
+class WeaveTracer(Tracer):
+    """Tracer implementation using Weave for telemetry and trace logging.
 
     This replaces AgentOpsTracer with a Weave-based manual trace context. It tracks:
+
     - Function/method calls
     - Input/Output data
     - Exceptions
-    and logs them to Weave Cloud (W&B backend) or optionally bypasses the network for testing.
 
-    Attributes:
-        project_name: Name of the Weave project. Used to initialize the Weave client.
-        _store: Optional LightningStore instance for storing collected spans.
-        instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
+    and logs them to Weave Cloud (W&B backend) or optionally bypasses the network for testing.
     """
 
     def __init__(
-        self, *, project_name: str | None = None, wandb_api_key: str | None = None, instrument_managed: bool = True
+        self,
+        *,
+        project_name: str | None = None,
+        weave_user_settings: UserSettings | None = None,
+        instrument_managed: bool = True,
     ):
-        """
-        Initialize a WeaveTracer instance.
+        """Initialize a WeaveTracer instance.
 
         Args:
             project_name: Optional project name for Weave; defaults to the current module name.
-            wandb_api_key: Optional W&B API key; sets environment variable if provided.
+            weave_user_settings: Optional UserSettings for Weave.
             instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
         """
         super().__init__()
-        self.project_name = project_name or __name__
-        self.sequence_id = 0
-        self._store: Optional[LightningStore] = None
+        self.project_name = project_name
         self.instrument_managed = instrument_managed
+        self.weave_user_settings = weave_user_settings or UserSettings(use_server_cache=False)
 
-        if wandb_api_key:
-            os.environ["WANDB_API_KEY"] = wandb_api_key
+        self._store: Optional[LightningStore] = None
+        self._server = WeaveTracerManagedTraceServer(
+            partial_call_callback=self.partial_call_callback, complete_call_callback=self.complete_call_callback
+        )
+
+        self._default_sequence_counter: int = 0
+        self._calls: Dict[str, tsi.CallSchema] = {}  # call_id -> call
+        self._spans: List[Span] = []  # spans in the current trace
+        self._rollout_id: Optional[str] = None
+        self._attempt_id: Optional[str] = None
+        self._partial_call_futures: Dict[str, asyncio.Future[int] | futures.Future[int]] = {}
+        self._complete_call_futures: List[asyncio.Future[None] | futures.Future[None]] = []
+        self._loop: weakref.ReferenceType[asyncio.AbstractEventLoop] | None = None
 
     def instrument(self, worker_id: int):
-        instrument_weave()
+        instrument_weave(self._server)
 
     def uninstrument(self, worker_id: int):
         uninstrument_weave()
@@ -75,22 +272,34 @@ class WeaveTracer(Tracer):
         logger.info(f"[Worker {worker_id}] Setting up Weave tracer...")
         self._store = store
 
-        try:
-            import weave
-        except ImportError:
-            raise RuntimeError("Weave is not installed. Install it to use WeaveTracer.")
-
         # Optionally patch network calls to bypass real Weave/W&B endpoints
         if self.instrument_managed:
             self.instrument(worker_id)
 
-        # Initialize the Weave client if not already initialized
-        if weave.get_client() is None:  # type: ignore
-            try:
-                weave.init(project_name=self.project_name)  # type: ignore
-                logger.info(f"[Worker {worker_id}] Weave client initialized.")
-            except Exception as e:
-                raise RuntimeError(f"Failed to initialize Weave for project '{self.project_name}': {e}")
+        # If WANDB_API_KEY is not set, we need to initialize Weave with a hack
+        if not os.getenv("WANDB_API_KEY"):
+            logger.info("WANDB_API_KEY is not set. Initializing Weave a mock context.")
+            set_wandb_api_context("agl", api_key=None, headers=None, cookies=None)
+        else:
+            logger.debug("WANDB_API_KEY is set. Weave will be initialized automatically.")
+
+        weave_client = weave.get_client()
+        if self.project_name is None:
+            self.project_name = random_project_name()
+
+        if weave_client is not None:
+            logger.warning("Weave client was already initialized. Reentrant calls are at your own risk.")
+            if weave_client.project == self.project_name:
+                logger.error(
+                    f"Weave client was already initialized for the same project '{self.project_name}'. It's very likely that weave won't work correctly."
+                )
+
+        # Init no matter what
+        try:
+            weave.init(project_name=self.project_name, settings=self.weave_user_settings)
+            logger.info(f"[Worker {worker_id}] Weave client initialized.")
+        except Exception as exc:
+            raise RuntimeError(f"Failed to initialize Weave for project '{self.project_name}'") from exc
 
     def teardown_worker(self, worker_id: int):
         """
@@ -105,21 +314,20 @@ class WeaveTracer(Tracer):
             self.uninstrument(worker_id)
             logger.info(f"[Worker {worker_id}] Instrumentation removed.")
 
+    @with_active_tracer_context
     @asynccontextmanager
     async def trace_context(
         self,
         name: Optional[str] = None,
         *,
-        store: Optional[LightningStore] = None,
         rollout_id: Optional[str] = None,
         attempt_id: Optional[str] = None,
+        **kwargs: Any,
     ) -> AsyncIterator[Any]:
-        """
-        Synchronous implementation of the tracing context.
+        """Asynchronous implementation of the tracing context.
 
         Args:
             name: Optional operation name.
-            store: Optional LightningStore instance.
             rollout_id: Optional rollout ID.
             attempt_id: Optional attempt ID.
 
@@ -127,181 +335,343 @@ class WeaveTracer(Tracer):
             ValueError: If store, rollout_id, and attempt_id are inconsistently provided.
             RuntimeError: If Weave is not installed or client is uninitialized.
         """
-        arg_op = name or self.project_name
-        arg_inputs: dict[str, str] | None = {"rollout_id": rollout_id or "", "attempt_id": attempt_id or ""}
 
-        if store is not None and rollout_id is not None and attempt_id is not None:
+        if rollout_id is not None and attempt_id is not None:
             self._rollout_id = rollout_id
             self._attempt_id = attempt_id
-            self._store = store
+        elif rollout_id is None and attempt_id is None:
+            logger.info("No rollout_id or attempt_id provided. Skipping writing to store.")
+            self._rollout_id = self._attempt_id = None
         else:
-            raise ValueError("store, rollout_id, and attempt_id must be either all provided")
+            raise ValueError("rollout_id and attempt_id must be either both provided or both None")
+
+        await self._init_trace_context()
+
+        weave_client = self._get_weave_client()
+
+        if weave_client.server is not self._server:
+            logger.error(
+                "Weave client is not using the correct trace server. You might have multiple WeaveTracer instances running in the same process. "
+                f"Expected {self._server}, got {weave_client.server}"
+            )
+
+        arg_op = name or weave_client.project
+        arg_inputs: dict[str, str] = {}
+        if rollout_id is not None:
+            arg_inputs[LightningResourceAttributes.ROLLOUT_ID.value] = rollout_id
+        if attempt_id is not None:
+            arg_inputs[LightningResourceAttributes.ATTEMPT_ID.value] = attempt_id
 
         try:
-            import datetime
+            # Create a new trace call object in Weave
+            trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
+                op=arg_op, inputs=arg_inputs
+            )
 
-            import weave
-        except ImportError:
-            raise RuntimeError("Weave is not installed. Install it to use WeaveTracer.")
+            try:
+                yield trace_call
+                # Finish trace even if no exception
+                weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
+            except Exception as exc:
+                # Finish trace and log any exception
+                weave_client.finish_call(trace_call, exception=exc)  # pyright: ignore[reportUnknownMemberType]
+                logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={exc}")
+                raise
 
-        weave_client = weave.get_client()  # type: ignore
+        finally:
+            try:
+                weave_client.flush()
+                # It's possible that the call end futures are from a dedicated Weave thread pool,
+                await asyncio.gather(*[asyncio.wrap_future(future) for future in self._complete_call_futures])
+
+            finally:
+                # Mandatory cleanup
+                self._rollout_id = None
+                self._attempt_id = None
+                self._server.clear()
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        if timestamp is not None:
+            logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
+        weave_client = self._get_weave_client()
+        trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
+            op=name,
+            attributes=attributes,
+            inputs={},
+        )
+        # Immediately finish the call
+        weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
+        # We don't wait for the call to be propagated to the server.
+        start_time = trace_call.started_at.timestamp() if trace_call.started_at else None
+        end_time = trace_call.ended_at.timestamp() if trace_call.ended_at else None
+        trace_status = (
+            TraceStatus(status_code="OK")
+            if trace_call.exception is None
+            else TraceStatus(status_code="ERROR", description=trace_call.exception)
+        )
+        return SpanCoreFields(
+            name=name,
+            attributes=flatten_attributes(trace_call.attributes or {}),
+            start_time=start_time,
+            end_time=end_time,
+            status=trace_status,
+        )
+
+    @contextmanager
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> Iterator[SpanRecordingContext]:
+        if start_time is not None:
+            logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
+        if end_time is not None:
+            logger.warning("Weave doesn't support customizing the end time of a call. Timestamp is ignored.")
+        weave_client = self._get_weave_client()
+        trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
+            op=name,
+            attributes=attributes,
+            inputs={},
+        )
+        recording_context = WeaveSpanRecordingContext(trace_call)
+        try:
+            yield recording_context
+        except Exception as exc:
+            recording_context.record_exception(exc)
+            raise
+        finally:
+            weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
+
+    async def _init_trace_context(self) -> None:
+        """Initialize the trace context."""
+        self._spans.clear()
+        self._calls.clear()
+        self._partial_call_futures.clear()
+        self._complete_call_futures.clear()
+        self._loop = weakref.ref(asyncio.get_running_loop())
+
+    def _get_weave_client(self) -> WeaveClient:
+        """Get the Weave client."""
+        weave_client = weave.get_client()
         if not weave_client:
             raise RuntimeError("Weave client is not initialized. Call init_worker() first.")
+        return weave_client
 
-        # Create a new trace call object in Weave
-        trace_call = weave_client.create_call(op=arg_op, inputs=arg_inputs)  # type: ignore
-        trace_call.started_at = datetime.datetime.now(tz=datetime.timezone.utc)
+    def _ensure_loop(self) -> tuple[asyncio.AbstractEventLoop, bool]:
+        """Returns a usable event loop and a boolean indicating whether it's the current running loop.
 
-        try:
-            yield trace_call
-        except Exception as e:
-            # Finish trace and log any exception
-            weave_client.finish_call(trace_call, exception=e)  # type: ignore
-            logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
-        finally:
-            # Finish trace even if no exception
-            weave_client.finish_call(trace_call)  # type: ignore
-            await self._on_finish_handler(trace_call)  # type: ignore
-
-    async def _on_finish_handler(self, call: "Call", *args: Any, **kwargs: Any) -> None:  # type: ignore
+        Prefer using the main loop if it's possible. Otherwise, use the current running loop.
         """
-        Handler called when a Weave Call finishes.
+        # Get the current running loop
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+
+        # Get the main loop, which can be a different loop
+        if self._loop is not None:
+            main_loop = self._loop()
+        else:
+            main_loop = None
+
+        if main_loop is not None:
+            return main_loop, id(main_loop) == id(running_loop)
+        elif running_loop is not None:
+            return running_loop, True
+        else:
+            raise RuntimeError("No running event loop found. This should not happen.")
+
+    def get_last_trace(self) -> List[Span]:
+        return self._spans
+
+    def partial_call_callback(self, request_content: Dict[str, Any]) -> None:
+        call_id = request_content.get("id")
+        if call_id is None:
+            raise ValueError("Call ID is required even for partial calls")
+
+        if call_id in self._partial_call_futures:
+            raise ValueError(f"Call {call_id} already has a start future")
+
+        # The callback must possibly be called from a dedicated Weave thread pool,
+        # but it should be executed on the main event loop.
+        try:
+            loop, is_current_loop = self._ensure_loop()
+            if is_current_loop:
+                task = loop.create_task(self.partial_call_handler(request_content))
+            else:
+                # Schedule the task on the dedicated loop
+                task = asyncio.run_coroutine_threadsafe(self.partial_call_handler(request_content), loop)
+            self._partial_call_futures[call_id] = task
+        except Exception as exc:
+            logger.exception(f"Error creating call start task: {exc}", exc_info=True)
+
+    def complete_call_callback(self, call: tsi.CallSchema) -> None:
+        try:
+            loop, is_current_loop = self._ensure_loop()
+            if is_current_loop:
+                task = loop.create_task(self.complete_call_handler(call))
+            else:
+                # Schedule the task on the dedicated loop
+                task = asyncio.run_coroutine_threadsafe(self.complete_call_handler(call), loop)
+            self._complete_call_futures.append(task)
+        except Exception as exc:
+            logger.exception(f"Error creating call finish task: {exc}", exc_info=True)
+
+    async def _get_next_sequence_id(self) -> int:
+        """Get the next sequence ID for a span.
+
+        Use store to get the next sequence ID if available, otherwise use a default counter.
+        """
+        if self._rollout_id and self._attempt_id and self._store:
+            return await self._store.get_next_span_sequence_id(self._rollout_id, self._attempt_id)
+        else:
+            self._default_sequence_counter += 1
+            return self._default_sequence_counter
+
+    async def partial_call_handler(self, request_content: Dict[str, Any]) -> int:
+        """Handler called when a Weave Call starts.
+
+        Args:
+            request_content: The partial Weave Call object.
+
+        Returns:
+            The sequence ID for the call.
+        """
+        sequence_id = await self._get_next_sequence_id()
+        return sequence_id
+
+    async def complete_call_handler(self, call: tsi.CallSchema) -> None:
+        """Handler called when a Weave Call finishes.
 
         Converts the call (including nested children) into spans and stores them in LightningStore.
         """
-        spans, self.sequence_id = self.convert_call_to_spans(call, self._rollout_id, self._attempt_id, self.sequence_id)  # type: ignore
+        # Make sure the corresponding call_start_future is complete
+        if call.id in self._partial_call_futures:
+            sequence_id = await asyncio.wrap_future(self._partial_call_futures[call.id])
+            del self._partial_call_futures[call.id]
+        else:
+            # Fetch a new sequence ID as the call_start is somehow missing
+            if call.id in self._calls:
+                logger.warning(
+                    f"Call {call.id} is already in calls. The call is already completed. Overwriting the call."
+                )
+            else:
+                logger.warning(f"Call {call.id} has no start future. Fetching a new sequence ID.")
+            sequence_id = await self._get_next_sequence_id()
 
+        self._calls[call.id] = call
+
+        span = await self.convert_call_to_span(call, self._rollout_id, self._attempt_id, sequence_id)
+        self._spans.append(span)
         if self._store and self._rollout_id and self._attempt_id:
             try:
-                await self._store.add_many_spans(spans)
-            except Exception as e:
-                logger.exception(f"Error adding span to store: {e}")
+                await self._store.add_span(span)
+            except Exception as exc:
+                logger.exception(f"Error adding span to store: {exc}")
 
-    def convert_call_to_spans(
+    async def convert_call_to_span(
         self,
-        call: "Call",  # type: ignore
+        call: tsi.CallSchema,
         rollout_id: Optional[str] = None,
         attempt_id: Optional[str] = None,
-        seq_start: int = 0,
-    ) -> tuple[List[Span], int]:
-        """
-        Recursively convert a Weave Call (with nested children) into a flat list of Agent Lightning Spans.
+        sequence_id: Optional[int] = None,
+    ) -> Span:
+        """Convert a Weave Call (with nested children) into a Agent-lightning Span.
+
+        `rollout_id` and `attempt_id` are required to attach the spans to the store.
 
         Args:
             call: The Weave Call object.
             rollout_id: Optional rollout ID to attach to spans.
             attempt_id: Optional attempt ID to attach to spans.
-            seq_start: Sequence number to start from.
+            sequence_id: Optional sequence ID to attach to spans.
 
         Returns:
-            Tuple of (list_of_spans, next_sequence_id).
+            List of converted spans.
         """
-        spans: List[Span] = []
-        sequence_id = seq_start
+        rollout_id = rollout_id or "rollout-dummy"
+        attempt_id = attempt_id or "attempt-dummy"
+        sequence_id = sequence_id or 0
 
-        rollout_id = rollout_id or ""  # type: ignore
-        attempt_id = attempt_id or ""  # type: ignore
+        start_ts: float = call.started_at.timestamp()
+        end_ts: Optional[float] = call.ended_at.timestamp() if call.ended_at else None
 
-        start_dt = getattr(call, "started_at", None)  # type: ignore
-        start_ts: Optional[float] = start_dt.timestamp() if start_dt else None
+        if call.exception:
+            status = TraceStatus(status_code="ERROR", description=call.exception)
+        else:
+            status = TraceStatus(status_code="OK")
 
-        end_dt = getattr(call, "ended_at", None)  # type: ignore
-        end_ts: Optional[float] = end_dt.timestamp() if end_dt else None
+        attributes: Dict[str, Any] = {
+            LightningSpanAttributes.OPERATION_NAME.value: call.op_name,
+            # op_name can be possibly overridden by the attributes.
+            **call.attributes,
+        }
+        if call.inputs:
+            attributes[LightningSpanAttributes.OPERATION_INPUT.value] = call.inputs
+        if call.output:
+            attributes[LightningSpanAttributes.OPERATION_OUTPUT.value] = call.output
+        if call.summary:
+            # attributes can be possibly overridden by the summary.
+            attributes.update(call.summary)
+        if call.exception:
+            attributes[exception_attributes.EXCEPTION_MESSAGE] = call.exception
 
-        trace_id = str(getattr(call, "trace_id", None))  # type: ignore
-        span_id = str(getattr(call, "id", None))  # type: ignore
-        parent_id = str(getattr(call, "parent_id", None)) if getattr(call, "parent_id", None) else None  # type: ignore
-
-        exception = getattr(call, "exception", None)  # type: ignore
-        status_code = "ERROR" if exception else "OK"
-
-        def sanitize(
-            inputs: Dict[str, Any],
-            output: Dict[str, Any],
-        ) -> Dict[str, str | JSONPrimitive]:
-            stack: List[Tuple[Any, str]] = [
-                (inputs or {}, "input"),
-                (output or {}, "output"),
-            ]
-
-            attributes: Dict[str, str | JSONPrimitive] = {}
-
-            while stack:
-                value, key = stack.pop()
-
-                if isinstance(value, dict):
-                    for k, v in value.items():  # type: ignore
-                        stack.append((v, f"{key}.{k}"))  # type: ignore
-                elif isinstance(value, (list, tuple)):
-                    for i, v in enumerate(value):  # type: ignore
-                        stack.append((v, f"{key}.{i}"))  # type: ignore
-                else:
-                    if value is None:
-                        attributes[key] = "None"
-                    elif isinstance(value, (str, int, float, bool)):
-                        attributes[key] = value
-                    else:
-                        try:
-                            attributes[key] = str(value)
-                        except Exception:
-                            attributes[key] = "None"
-
-            return attributes
-
-        inputs = getattr(call, "inputs", {})  # type: ignore
-        output = getattr(call, "output", {})  # type: ignore
-        attributes = sanitize(inputs, output)
+        sanitized_attributes = sanitize_attributes(flatten_attributes(attributes, expand_leaf_lists=False))
 
         context = SpanContext(
-            trace_id=trace_id,
-            span_id=span_id,
+            trace_id=call.trace_id,
+            span_id=call.id,
             is_remote=False,
             trace_state={},
         )
 
-        parent_context = (
-            SpanContext(
-                trace_id=trace_id,
-                span_id=parent_id,
-                is_remote=False,
-                trace_state={},
-            )
-            if parent_id
-            else None
-        )
+        # Get context for parent
+        if call.parent_id:
+            parent_call = self._calls.get(call.parent_id)
+            if parent_call:
+                parent_context = SpanContext(
+                    trace_id=parent_call.trace_id,
+                    span_id=parent_call.id,
+                    is_remote=False,
+                    trace_state={},
+                )
+            else:
+                parent_context = None
+        else:
+            parent_context = None
 
         # Build the Span object
-        span = Span(
-            rollout_id=rollout_id or "",
-            attempt_id=attempt_id or "",
+        return Span(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
             sequence_id=sequence_id,
-            trace_id=trace_id,
-            span_id=span_id,
-            parent_id=parent_id,
-            name=getattr(call, "func_name", "unknown"),  # type: ignore
-            status=TraceStatus(status_code=status_code),
-            attributes=attributes,  # type: ignore
+            trace_id=call.trace_id,
+            span_id=call.id,
+            parent_id=call.parent_id,
+            name=op_name_to_func_name(call.op_name),
+            status=status,
+            attributes=sanitized_attributes,
             events=[],  # Weave calls do not generate events
             links=[],  # Weave calls do not generate links
             start_time=start_ts,
             end_time=end_ts,
             context=context,
             parent=parent_context,
-            resource=OtelResource(attributes={}, schema_url=""),
+            resource=OtelResource(
+                attributes={
+                    LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
+                    LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
+                    LightningResourceAttributes.SPAN_SEQUENCE_ID.value: sequence_id,
+                    LightningResourceAttributes.TRACER_NAME.value: "weave",
+                },
+                schema_url="",
+            ),
         )
-
-        spans.append(span)
-        sequence_id += 1
-
-        children: List["Call"] = getattr(call, "_children", [])  # type: ignore
-        # Recursively process child calls
-        for child in children:  # type: ignore
-            child_spans, sequence_id = self.convert_call_to_spans(  # type: ignore
-                child,  # type: ignore
-                rollout_id=rollout_id,
-                attempt_id=attempt_id,
-                seq_start=sequence_id,
-            )
-            spans.extend(child_spans)
-
-        return spans, sequence_id

agentlightning/types/core.py

@@ -28,7 +28,7 @@ from typing import (
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel, Field, model_validator
 
-from .tracer import Span
+from .tracer import Span, SpanCoreFields
 
 if TYPE_CHECKING:
     from agentlightning.litagent import LitAgent
@@ -307,6 +307,7 @@ RolloutRawResult = Union[
     float,  # only final reward
     List[ReadableSpan],  # constructed OTEL spans by user
     List[Span],  # constructed Span objects by user
+    List[SpanCoreFields],  # constructed SpanCoreFields objects by user
 ]
 """Rollout result type.
 

agentlightning/types/tracer.py

@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+import time
+
 """Data models that mirror OpenTelemetry spans for Agent Lightning."""
 
 import json
 from enum import Enum
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Literal, Optional, Protocol, Sequence, Union
 
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk.resources import Resource
@@ -31,6 +33,9 @@ __all__ = [
     "SpanNames",
     "SpanAttributeNames",
     "SpanLike",
+    "StatusCode",
+    "SpanCoreFields",
+    "SpanRecordingContext",
 ]
 
 
@@ -83,6 +88,8 @@ Attributes = Dict[str, AttributeValue]
 """Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
 TraceState = Dict[str, str]
 """Mapping from trace state key to its value. Same as OpenTelemetry `TraceState` type."""
+StatusCode = Literal["UNSET", "OK", "ERROR"]
+"""The status code of the span."""
 
 
 class SpanContext(BaseModel):
@@ -115,7 +122,7 @@ class SpanContext(BaseModel):
 class TraceStatus(BaseModel):
     """Serializable variant of `opentelemetry.trace.Status`."""
 
-    status_code: str
+    status_code: StatusCode
     """The status code of the span. Same as OpenTelemetry `Status.status_code` type."""
     description: Optional[str] = None
     """The description of the span. Same as OpenTelemetry `Status.description` type."""
@@ -203,6 +210,44 @@ class OtelResource(BaseModel):
         )
 
 
+class SpanCoreFields(BaseModel):
+    """Core fields of a span. Used by span creators who don't care about the full span model.
+
+    If the spans are managed by some OTel tracer provider, it's not advised to create spans via this path.
+    """
+
+    name: str
+    """The name of the span."""
+    status: TraceStatus
+    """The status of the span."""
+    attributes: Attributes
+    """The attributes of the span."""
+    start_time: Optional[float]
+    """The start time of the span."""
+    end_time: Optional[float]
+    """The end time of the span."""
+
+
+class SpanRecordingContext(Protocol):
+    """Context for recording operations on a span. It doesn't have to finalize the span; the caller will do it."""
+
+    def record_exception(self, exception: BaseException) -> None:
+        """Record an exception on the span."""
+        raise NotImplementedError()
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        """Record attributes on the span."""
+        raise NotImplementedError()
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        """Record the status of the span."""
+        raise NotImplementedError()
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        """Get the recording of the span."""
+        raise NotImplementedError()
+
+
 class Span(BaseModel):
     """Agent Lightning's canonical span model used for persistence and analytics.
 
@@ -340,6 +385,7 @@ class Span(BaseModel):
         start_time: Optional[float] = None,
         end_time: Optional[float] = None,
         resource: Optional[OtelResource] = None,
+        status: Optional[TraceStatus] = None,
     ) -> "Span":
         """Build a synthetic span from raw attributes.
         Different from the [`from_opentelemetry`][agentlightning.Span.from_opentelemetry] method,
@@ -357,6 +403,7 @@ class Span(BaseModel):
             start_time: Span start timestamp in seconds.
             end_time: Span end timestamp in seconds.
             resource: Explicit resource information to attach to the span.
+            status: Optional status of the span.
 
         Returns:
             [`Span`][agentlightning.Span] populated with the provided attributes.
@@ -384,7 +431,7 @@ class Span(BaseModel):
             name=name or AGL_VIRTUAL,
             resource=resource or OtelResource(attributes={}, schema_url=""),
             attributes=attributes,
-            status=TraceStatus(status_code="OK"),
+            status=status or TraceStatus(status_code="OK"),
             events=[],
             links=[],
             parent=(
@@ -399,6 +446,37 @@ class Span(BaseModel):
             ),
         )
 
+    @classmethod
+    def from_core_fields(
+        cls,
+        core: SpanCoreFields,
+        *,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        sequence_id: Optional[int] = None,
+    ) -> Span:
+        """Build a span from a core span.
+
+        Args:
+            core: Core span to build from.
+            rollout_id: Optional rollout identifier associated with the span.
+            attempt_id: Optional attempt identifier associated with the span.
+            sequence_id: Optional sequence number to preserve ordering.
+
+        Returns:
+            [`Span`][agentlightning.Span] populated with the provided attributes.
+        """
+        return cls.from_attributes(
+            attributes=core.attributes,
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+            name=core.name,
+            start_time=core.start_time or time.time(),
+            end_time=core.end_time,
+            status=core.status,
+        )
+
 
 class SpanNames(str, Enum):
     """Enumerated span names recognised by Agent-lightning. Deprecated in favor of [semconv][agentlightning.semconv]."""

agentlightning/utils/id.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import hashlib
+import uuid
+
+__all__ = ["generate_id"]
+
+
+def generate_id(length: int) -> str:
+    """Generate a random hexadecimal ID of the given length.
+
+    Args:
+        length: The desired length of the generated ID. Must be a positive integer.
+
+    Returns:
+        A random hexadecimal ID string of the given length.
+
+    Raises:
+        ValueError: If length is not a positive integer.
+    """
+    if length <= 0:
+        raise ValueError("length must be a positive integer")
+
+    return hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:length]

agentlightning/utils/otel.py

@@ -2,22 +2,25 @@
 
 """Utilities shared for OpenTelemetry span (attributes) support."""
 
+import json
 import logging
-from typing import Any, Dict, List, Sequence, Union, cast
+import traceback
+from typing import Any, Dict, List, Sequence, Type, TypeVar, Union, cast
 from warnings import filterwarnings
 
 import opentelemetry.trace as trace_api
 from agentops.sdk.exporters import OTLPSpanExporter
-from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SynchronousMultiSpanProcessor, Tracer
+from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SpanProcessor, SynchronousMultiSpanProcessor, Tracer
 from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
 from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
+from opentelemetry.semconv.attributes import exception_attributes
 from opentelemetry.trace import get_tracer_provider as otel_get_tracer_provider
 from pydantic import TypeAdapter
 
 from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
 from agentlightning.semconv import LightningSpanAttributes, LinkAttributes, LinkPydanticModel
-from agentlightning.types import SpanLike
+from agentlightning.types import Attributes, AttributeValue, SpanLike
 from agentlightning.utils.otlp import LightningStoreOTLPExporter
 
 logger = logging.getLogger(__name__)
@@ -35,8 +38,16 @@ __all__ = [
     "filter_and_unflatten_attributes",
     "flatten_attributes",
     "unflatten_attributes",
+    "sanitize_attribute_value",
+    "sanitize_attributes",
+    "sanitize_list_attribute_sanity",
+    "check_attributes_sanity",
+    "format_exception_attributes",
 ]
 
+T_SpanLike = TypeVar("T_SpanLike", bound=SpanLike)
+T_SpanProcessor = TypeVar("T_SpanProcessor", bound=SpanProcessor)
+
 
 def full_qualified_name(obj: type) -> str:
     if str(obj.__module__) == "builtins":
@@ -112,6 +123,25 @@ def get_tracer_provider(inspect: bool = True) -> TracerProviderImpl:
     return tracer_provider
 
 
+def get_span_processors(
+    tracer_provider: TracerProviderImpl, expected_type: Type[T_SpanProcessor]
+) -> List[T_SpanProcessor]:
+    """Get the span processors from the tracer provider.
+
+    Args:
+        tracer_provider: The tracer provider to get the span processors from.
+        expected_type: The type of the span processors to get.
+
+    Returns:
+        A list of span processors of the expected type.
+    """
+    processors: List[T_SpanProcessor] = []
+    for processor in tracer_provider._active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+        if isinstance(processor, expected_type):
+            processors.append(processor)
+    return processors
+
+
 def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
     """Resolve the OpenTelemetry tracer configured for Agent Lightning.
 
@@ -166,7 +196,7 @@ def make_tag_attributes(tags: List[str]) -> Dict[str, Any]:
     ["gen_ai.model:gpt-4", "reward.extrinsic"]
     ```
     """
-    return flatten_attributes({LightningSpanAttributes.TAG.value: tags})
+    return flatten_attributes({LightningSpanAttributes.TAG.value: tags}, expand_leaf_lists=True)
 
 
 def extract_tags_from_attributes(attributes: Dict[str, Any]) -> List[str]:
@@ -196,10 +226,10 @@ def make_link_attributes(links: Dict[str, str]) -> Dict[str, Any]:
         if not isinstance(value, str):  # pyright: ignore[reportUnnecessaryIsInstance]
             raise ValueError(f"Link value must be a string, got {type(value)} for key '{key}'")
         link_list.append({LinkAttributes.KEY_MATCH.value: key, LinkAttributes.VALUE_MATCH.value: value})
-    return flatten_attributes({LightningSpanAttributes.LINK.value: link_list})
+    return flatten_attributes({LightningSpanAttributes.LINK.value: link_list}, expand_leaf_lists=True)
 
 
-def query_linked_spans(spans: Sequence[SpanLike], links: List[LinkPydanticModel]) -> List[SpanLike]:
+def query_linked_spans(spans: Sequence[T_SpanLike], links: List[LinkPydanticModel]) -> List[T_SpanLike]:
     """Query spans that are linked by the given link attributes.
 
     Args:
@@ -209,7 +239,7 @@ def query_linked_spans(spans: Sequence[SpanLike], links: List[LinkPydanticModel]
     Returns:
         A list of spans that match the given link attributes.
     """
-    matched_spans: List[SpanLike] = []
+    matched_spans: List[T_SpanLike] = []
 
     for span in spans:
         span_attributes = span.attributes or {}
@@ -294,7 +324,9 @@ def filter_and_unflatten_attributes(attributes: Dict[str, Any], prefix: str) ->
     return unflatten_attributes(stripped_attributes)
 
 
-def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[str, Any]:
+def flatten_attributes(
+    nested_data: Union[Dict[str, Any], List[Any]], *, expand_leaf_lists: bool = False
+) -> Dict[str, Any]:
     """Flatten a nested dictionary or list into a flat dictionary with dotted keys.
 
     This function recursively traverses dictionaries and lists, producing a flat
@@ -303,12 +335,14 @@ def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[st
 
     Example:
 
-        >>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}})
+        >>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}}, expand_leaf_lists=True)
         {"a.b": 1, "a.c.0": 2, "a.c.1": 3}
 
     Args:
-        nested_data: A nested structure composed of dictionaries, lists, or
-            primitive values.
+        nested_data: A nested structure composed of dictionaries, lists, or primitive values.
+        expand_leaf_lists: Whether to expand lists composed only of primitive values.
+            When `False` (the default), lists of str/int/float/bool are treated as
+            leaf values and stored without enumerating their indices.
 
     Returns:
         A flat dictionary mapping dotted-string paths to primitive values.
@@ -316,6 +350,15 @@ def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[st
 
     flat: Dict[str, Any] = {}
 
+    def _primitive_type(value: Any) -> Union[type[str], type[int], type[float], type[bool]]:
+        if isinstance(value, bool):
+            return bool
+        if isinstance(value, int):
+            return int
+        if isinstance(value, float):
+            return float
+        return str
+
     def _walk(value: Any, prefix: str = "") -> None:
         if isinstance(value, dict):
             for k, v in cast(Dict[Any, Any], value).items():
@@ -326,7 +369,22 @@ def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[st
                 new_prefix = f"{prefix}.{k}" if prefix else k
                 _walk(v, new_prefix)
         elif isinstance(value, list):
-            for idx, item in enumerate(cast(List[Any], value)):
+            maybe_list = cast(List[Any], value)
+            is_leaf_candidate = bool(maybe_list) and all(
+                isinstance(item, (str, int, float, bool)) for item in maybe_list
+            )
+            if not expand_leaf_lists and is_leaf_candidate and prefix:
+                primitive_types = {_primitive_type(item) for item in maybe_list}
+                if len(primitive_types) == 1:
+                    flat[prefix] = maybe_list
+                    return
+                logger.warning(
+                    "List attribute '%s' contains mixed primitive types %s; expanding indexed keys instead.",
+                    prefix,
+                    primitive_types,
+                )
+
+            for idx, item in enumerate(maybe_list):
                 new_prefix = f"{prefix}.{idx}" if prefix else str(idx)
                 _walk(item, new_prefix)
         else:
@@ -399,3 +457,85 @@ def unflatten_attributes(flat_data: Dict[str, Any]) -> Union[Dict[str, Any], Lis
         return node
 
     return convert(root)
+
+
+def sanitize_attribute_value(object: Any, force: bool = True) -> AttributeValue:
+    """Sanitize an attribute value to be a valid OpenTelemetry attribute value."""
+    if isinstance(object, (str, int, float, bool)):
+        return object
+
+    if isinstance(object, list):
+        try:
+            return sanitize_list_attribute_sanity(cast(List[Any], object))
+        except ValueError as exc:
+            logger.warning(f"Failed to sanitize list attribute. Fallback to JSON serialization: {exc}")
+
+    try:
+        # This include null, dict, etc.
+        serialized = json.dumps(object, default=str if force else None)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Object must be JSON serializable, got: {type(cast(Any, object))}.") from exc
+    return serialized
+
+
+def sanitize_attributes(attributes: Dict[str, Any], force: bool = True) -> Attributes:
+    """Sanitize a dictionary of attributes to be a valid OpenTelemetry attributes.
+
+    Args:
+        attributes: A dictionary of attributes to sanitize.
+        force: Whether to force sanitization even when the value is not JSON serializable.
+    """
+    result: Attributes = {}
+    for k, v in attributes.items():
+        try:
+            result[k] = sanitize_attribute_value(v, force=force)
+        except ValueError as exc:
+            raise ValueError(f"Failed to sanitize attribute '{k}': {exc}") from exc
+    return result
+
+
+def sanitize_list_attribute_sanity(maybe_list: List[Any]) -> AttributeValue:
+    """Try to sanitize a list of attributes to be a valid OpenTelemetry attribute value.
+
+    Raise error if the list contains multiple types of primitive values.
+    """
+    if all(isinstance(item, str) for item in maybe_list):
+        return list[str](maybe_list)
+    if all(isinstance(item, bool) for item in maybe_list):
+        return list[bool](maybe_list)
+    if all(isinstance(item, (int, bool)) for item in maybe_list):
+        return [int(item) for item in maybe_list]
+    if all(isinstance(item, (float, int, bool)) for item in maybe_list):
+        return [float(item) for item in maybe_list]
+
+    list_types: List[Any] = [type(item) for item in maybe_list]
+    raise ValueError(f"List must contain only one type of primitive values, got: {set(list_types)}.")
+
+
+def check_attributes_sanity(attributes: Dict[Any, Any]) -> None:
+    """Check if a dictionary of attributes is a valid OpenTelemetry attributes."""
+    for k, v in attributes.items():
+        if not isinstance(k, str):
+            raise ValueError(f"Attribute key must be a string, got {type(k)} for key '{k}'")
+        if isinstance(v, list):
+            try:
+                sanitize_list_attribute_sanity(cast(List[Any], v))
+            except ValueError as exc:
+                raise ValueError(f"Failed to sanitize list attribute '{k}': {exc}") from exc
+        elif not isinstance(v, (str, int, float, bool)):
+            raise ValueError(
+                f"Attribute value must be a string, int, float, bool, or list of these, got {type(v)} for value '{v}'"
+            )
+
+
+def format_exception_attributes(exception: BaseException) -> Attributes:
+    """Format an exception into a dictionary of attributes."""
+    stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
+    span_attributes: Attributes = {
+        exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
+        exception_attributes.EXCEPTION_MESSAGE: str(exception),
+        exception_attributes.EXCEPTION_ESCAPED: True,
+    }
+    if stacktrace.strip():
+        span_attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+    return span_attributes

agentlightning/utils/otlp.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import gzip
 import logging
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar
+from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar
 
 from fastapi import Request, Response
 from google.protobuf import json_format
@@ -39,6 +39,7 @@ from agentlightning.types.tracer import (
     OtelResource,
     Span,
     SpanContext,
+    StatusCode,
     TraceStatus,
     convert_timestamp,
 )
@@ -413,7 +414,7 @@ def _kv_list_to_dict(kvs: Sequence[KeyValue]) -> Attributes:
     return {kv.key: _any_value_to_python(kv.value) for kv in kvs}
 
 
-_STATUS_CODE_MAP = {
+_STATUS_CODE_MAP: Mapping[ProtoStatus.StatusCode.ValueType, StatusCode] = {
     ProtoStatus.STATUS_CODE_UNSET: "UNSET",
     ProtoStatus.STATUS_CODE_OK: "OK",
     ProtoStatus.STATUS_CODE_ERROR: "ERROR",

agentlightning/utils/system_snapshot.py

@@ -13,12 +13,20 @@ from gpustat import GPUStat, GPUStatCollection
 
 
 def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
+    """Capture a snapshot of the system's hardware and software information.
+
+    Args:
+        include_gpu: Whether to include GPU information.
+
+    Returns:
+        A dictionary containing the system's hardware and software information.
+    """
     # CPU
     cpu = {
         "cpu_name": platform.processor(),
         "cpu_cores": psutil.cpu_count(logical=False),
         "cpu_threads": psutil.cpu_count(logical=True),
-        "cpu_usage_pct": psutil.cpu_percent(0.05),
+        "cpu_usage_pct": psutil.cpu_percent(0.0),
     }
 
     # Memory
@@ -37,20 +45,21 @@ def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
         "disk_pct": du.percent,
     }
 
-    # GPU
+    # GPU (only query if explicitly requested)
     gpus: List[Dict[str, Any]] = []
-    with suppress(Exception):
-        for g in GPUStatCollection.new_query().gpus:  # type: ignore
-            g = cast(GPUStat, g)
-            gpus.append(
-                {
-                    "gpu": g.name,  # type: ignore
-                    "util_pct": g.utilization,
-                    "mem_used_mb": g.memory_used,
-                    "mem_total_mb": g.memory_total,
-                    "temp_c": g.temperature,
-                }
-            )
+    if include_gpu:
+        with suppress(Exception):
+            for g in GPUStatCollection.new_query().gpus:  # type: ignore
+                g = cast(GPUStat, g)
+                gpus.append(
+                    {
+                        "gpu": g.name,  # type: ignore
+                        "util_pct": g.utilization,
+                        "mem_used_mb": g.memory_used,
+                        "mem_total_mb": g.memory_total,
+                        "temp_c": g.temperature,
+                    }
+                )
 
     # Network
     net = psutil.net_io_counters()

agentlightning/verl/config.yaml

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

agentlightning/verl/daemon.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import json
+import os
 import random
 import socket
 import threading
@@ -31,6 +32,85 @@ __all__ = [
 ]
 
 
+def ids_startswith(
+    full_ids: List[int], prefix_ids: List[int], tokenizer: Any, debug: bool = False
+) -> Tuple[bool, Tuple[bool, bool, bool]]:
+    is_prefix: bool
+    template_mismatch, retoken_mismatch, others_mismatch = False, False, False
+    if full_ids[: len(prefix_ids)] == prefix_ids:
+        is_prefix = True
+        return True, (template_mismatch, retoken_mismatch, others_mismatch)
+    else:
+        is_prefix = False
+
+    if not debug:
+        return is_prefix, (template_mismatch, retoken_mismatch, others_mismatch)
+
+    def _special_token_sequence(ids: List[int]) -> List[int]:
+        return [id for id in ids if id in tokenizer.all_special_ids]
+
+    def _none_special_token_sequence(ids: List[int]) -> List[int]:
+        return [id for id in ids if id not in tokenizer.all_special_ids]
+
+    # First, handle special tokens
+    full_special_ids = _special_token_sequence(full_ids)
+    prefix_special_ids = _special_token_sequence(prefix_ids)
+    if sum(1 for a, b in zip(full_special_ids, prefix_special_ids) if a != b) > 0:
+        template_mismatch = True
+
+    # Next, handle string content
+    full_content_ids = _none_special_token_sequence(full_ids)
+    prefix_content_ids = _none_special_token_sequence(prefix_ids)
+    full_string = tokenizer.decode(full_ids, skip_special_tokens=True)
+    prefix_string = tokenizer.decode(prefix_ids, skip_special_tokens=True)
+    if full_content_ids[: len(prefix_content_ids)] != prefix_content_ids and full_string.startswith(prefix_string):
+        retoken_mismatch = True
+    elif full_content_ids[: len(prefix_content_ids)] != prefix_content_ids and not full_string.startswith(
+        prefix_string
+    ):
+        others_mismatch = True
+    return is_prefix, (template_mismatch, retoken_mismatch, others_mismatch)
+
+
+def log_mismatch_detail(
+    diagnostic: Tuple[bool, bool, bool],
+    full_ids: List[int],
+    prefix_ids: List[int],
+    global_steps: int,
+    rollout_id: str,
+    turn_id: int,
+    log_dir: str | None = None,
+):
+    if log_dir is None:
+        return
+    os.makedirs(log_dir, exist_ok=True)
+    template_mismatch, retoken_mismatch, others_mismatch = diagnostic
+    if template_mismatch:
+        with open(os.path.join(log_dir, "template_mismatch.log"), "a+") as f:
+            print(
+                "-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
+                file=f,
+            )
+            print(full_ids, file=f)
+            print(prefix_ids, file=f)
+    if retoken_mismatch:
+        with open(os.path.join(log_dir, "retoken_mismatch.log"), "a+") as f:
+            print(
+                "-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
+                file=f,
+            )
+            print(full_ids, file=f)
+            print(prefix_ids, file=f)
+    if others_mismatch:
+        with open(os.path.join(log_dir, "others_mismatch.log"), "a+") as f:
+            print(
+                "-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
+                file=f,
+            )
+            print(full_ids, file=f)
+            print(prefix_ids, file=f)
+
+
 def get_left_padded_ids_and_attention_mask(
     ids: List[int], max_length: int, pad_token_id: int
 ) -> Tuple[List[int], List[int]]:
@@ -146,6 +226,7 @@ class AgentModeDaemon:
         adapter: TraceToTripletBase | None = None,
         processor: Any = None,
         image_base_dir: Optional[str] = None,
+        trace_aggregator: Dict[str, Any] = {"level": "transition"},
     ):
         self.mode = mode
         self.llm_timeout_seconds = llm_timeout_seconds
@@ -188,6 +269,7 @@ class AgentModeDaemon:
         self.processor = processor
         self.reward_fillna_value = reward_fillna_value
         self.image_base_dir = image_base_dir
+        self.trace_aggregator = trace_aggregator
 
         # Check if model requires multimodal position_ids (e.g., Qwen2-VL)
         self._use_mrope = self._is_mrope_model()
@@ -520,7 +602,7 @@ class AgentModeDaemon:
                 raise RuntimeError("Internal loop is not running.")
             future = asyncio.run_coroutine_threadsafe(coro, self._internal_loop)
         try:
-            future.result(timeout=60)  # Wait for completion with a timeout
+            future.result(timeout=300)  # Wait for completion with a timeout
         except Exception as e:
             print(f"Failed to set up data on server: {e}")
             raise
@@ -722,7 +804,9 @@ class AgentModeDaemon:
         )
         return metric_dict
 
-    def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, device: torch.device):
+    def get_train_data_batch(
+        self, max_prompt_length: int, max_response_length: int, device: torch.device, global_steps: int
+    ):
         """
         Processes completed rollouts to generate a training data batch.
 
@@ -788,50 +872,165 @@ class AgentModeDaemon:
         image_grid_thw_list: List[Optional[torch.Tensor]] = []  # For Qwen2-VL mrope
         n_trunc_sample_because_of_response = 0
 
-        for rollout_id, sample_info in finished_id_to_sample_info.items():
-            for turn_index, trace in enumerate(sample_info["trace_list"]):
+        if self.trace_aggregator.get("level", "transition") == "transition":
+            for rollout_id, sample_info in finished_id_to_sample_info.items():
+                for turn_index, trace in enumerate(sample_info["trace_list"]):
 
-                reward_list.append(sample_info["reward"])
-                prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]
+                    reward_list.append(sample_info["reward"])
+                    prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]
 
-                # Mark samples with prompts exceeding max_prompt_length to be dropped later
-                if len(prompt_ids) > max_prompt_length:
-                    prompt_ids = prompt_ids[:max_prompt_length]
-                    is_drop_list.append(True)
-                else:
-                    is_drop_list.append(False)
+                    # Mark samples with prompts exceeding max_prompt_length to be dropped later
+                    if len(prompt_ids) > max_prompt_length:
+                        prompt_ids = prompt_ids[:max_prompt_length]
+                        is_drop_list.append(True)
+                    else:
+                        is_drop_list.append(False)
 
-                # Truncate responses that exceed max_response_length
-                if len(response_ids) > max_response_length:
-                    response_ids = response_ids[:max_response_length]
-                    n_trunc_sample_because_of_response += 1
+                    # Truncate responses that exceed max_response_length
+                    if len(response_ids) > max_response_length:
+                        response_ids = response_ids[:max_response_length]
+                        n_trunc_sample_because_of_response += 1
 
-                # Pad prompts to the left and responses to the right
-                one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
-                    prompt_ids, max_prompt_length, self.pad_token_id
-                )
-                one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
-                    response_ids, max_response_length, self.pad_token_id
-                )
+                    # Pad prompts to the left and responses to the right
+                    one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
+                        prompt_ids, max_prompt_length, self.pad_token_id
+                    )
+                    one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
+                        response_ids, max_response_length, self.pad_token_id
+                    )
 
-                input_ids_list.append(one_input_ids)
-                input_attention_mask_list.append(one_input_attention_mask)
-                response_ids_list.append(one_response_ids)
-                response_attention_mask_list.append(one_response_attention_mask)
-                data_id_list.append(sample_info["data_id"])
-                rollout_id_list.append(rollout_id)
-                turn_index_list.append(turn_index)
+                    input_ids_list.append(one_input_ids)
+                    input_attention_mask_list.append(one_input_attention_mask)
+                    response_ids_list.append(one_response_ids)
+                    response_attention_mask_list.append(one_response_attention_mask)
+                    data_id_list.append(sample_info["data_id"])
+                    rollout_id_list.append(rollout_id)
+                    turn_index_list.append(turn_index)
 
-                # Compute image_grid_thw for this triplet using image_urls from prompt
-                if self._use_mrope:
-                    image_urls = trace.get("image_urls", [])
-                    image_grid_thw_list.append(self._get_image_grid_thw(image_urls))
+                    # Compute image_grid_thw for this triplet using image_urls from prompt
+                    if self._use_mrope:
+                        image_urls = trace.get("image_urls", [])
+                        image_grid_thw_list.append(self._get_image_grid_thw(image_urls))
+
+        elif self.trace_aggregator.get("level", "transition") == "trajectory":
+            assert not self._use_mrope, "M-RoPE is not supported in trajectory level yet."
+
+            response_mask_list: List[List[int]] = []
+            unmerged_count: int = 0
+            template_mismatch_count, retoken_mismatch_count, others_mismatch_count = 0, 0, 0
+            response_per_turn_list: List[int] = []
+
+            for rollout_id, sample_info in finished_id_to_sample_info.items():
+                merged_trace_idx: List[List[int]] = []
+
+                # Identify which turns can be merged based on token ids prefix matching
+                current_merged_trace_idx: List[int] = []
+                current_context: List[int] = []
+                for turn_index, trace in enumerate(sample_info["trace_list"]):
+                    response_per_turn_list.append(len(trace["response_ids"]))
+                    is_prefix, diagnostic = ids_startswith(
+                        trace["prompt_ids"] + trace["response_ids"],
+                        current_context,
+                        self.tokenizer,
+                        self.trace_aggregator.get("debug", False),
+                    )
+                    if not is_prefix and self.trace_aggregator.get("debug", False) == True:
+                        template_mismatch_count += diagnostic[0]
+                        retoken_mismatch_count += diagnostic[1]
+                        others_mismatch_count += diagnostic[2]
+                        log_mismatch_detail(
+                            diagnostic,
+                            trace["prompt_ids"] + trace["response_ids"],
+                            current_context,
+                            global_steps,
+                            rollout_id,
+                            turn_index,
+                            self.trace_aggregator.get("mismatch_log_dir", None),
+                        )
+
+                    if is_prefix:
+                        current_context = trace["prompt_ids"] + trace["response_ids"]
+                        current_merged_trace_idx.append(turn_index)
+                    else:
+                        merged_trace_idx.append(current_merged_trace_idx)
+                        current_merged_trace_idx = [turn_index]
+                        current_context = trace["prompt_ids"] + trace["response_ids"]
+
+                if current_merged_trace_idx not in merged_trace_idx:
+                    merged_trace_idx.append(current_merged_trace_idx)
+
+                if len(merged_trace_idx) > 1:
+                    unmerged_count += 1
+
+                # Merge all trace segments in merged_trace_idx into training samples
+                for current_merged_trace_idx in merged_trace_idx:
+                    prompt_ids = sample_info["trace_list"][current_merged_trace_idx[0]]["prompt_ids"]
+
+                    # if the merged_trace_idx doesn't start with the beginning of the prompt_ids, we need to adjust it
+                    if current_merged_trace_idx[0] > 0 and len(prompt_ids) > max_prompt_length:
+                        response_ids = prompt_ids[max_prompt_length:]
+                        prompt_ids = prompt_ids[:max_prompt_length]
+                        response_mask = [1] * len(response_ids)
+                    else:
+                        response_ids = []
+                        response_mask = []
+
+                    prompt_length = len(prompt_ids)
+                    response_ids += sample_info["trace_list"][current_merged_trace_idx[0]]["response_ids"]
+                    response_mask += [1] * len(response_ids)
+                    for turn_index in current_merged_trace_idx[1:]:
+                        trace = sample_info["trace_list"][turn_index]
+                        new_prompt_length = len(trace["prompt_ids"]) - len(response_ids) - prompt_length
+                        response_ids += trace["prompt_ids"][-new_prompt_length:]
+                        response_ids += trace["response_ids"]
+                        response_mask += [0] * new_prompt_length
+                        response_mask += [1] * len(trace["response_ids"])
+
+                    reward_list.append(sample_info["reward"])
+
+                    # Mark samples with prompts exceeding max_prompt_length to be dropped later
+                    if len(prompt_ids) > max_prompt_length:
+                        prompt_ids = prompt_ids[:max_prompt_length]
+                        is_drop_list.append(True)
+                    else:
+                        is_drop_list.append(False)
+
+                    # Truncate responses that exceed max_response_length
+                    if len(response_ids) > max_response_length:
+                        response_ids = response_ids[:max_response_length]
+                        response_mask = response_mask[:max_response_length]
+                        n_trunc_sample_because_of_response += 1
+
+                    # Pad prompts to the left and responses to the right
+                    one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
+                        prompt_ids, max_prompt_length, self.pad_token_id
+                    )
+                    one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
+                        response_ids, max_response_length, self.pad_token_id
+                    )
+                    one_response_mask, _ = get_right_padded_ids_and_attention_mask(
+                        response_mask, max_response_length, 0
+                    )
+
+                    input_ids_list.append(one_input_ids)
+                    input_attention_mask_list.append(one_input_attention_mask)
+                    response_ids_list.append(one_response_ids)
+                    response_attention_mask_list.append(one_response_attention_mask)
+                    response_mask_list.append(one_response_mask)
+                    data_id_list.append(sample_info["data_id"])
+                    rollout_id_list.append(rollout_id)
+                    # turn_index_list.append(current_merged_trace_idx)
+        else:
+            raise ValueError(f"Unknown trace_aggregator level: {self.trace_aggregator.get('level')}")
 
         n_transition = len(input_ids_list)
         batch_input_ids = torch.LongTensor(input_ids_list).to(device)
         input_attention_mask = torch.LongTensor(input_attention_mask_list).to(device)
         batch_response_ids = torch.LongTensor(response_ids_list).to(device)
         response_attention_mask = torch.LongTensor(response_attention_mask_list).to(device)
+        response_mask = (
+            torch.LongTensor(response_mask_list).to(device) if self.trace_aggregator.get("level", "transition") == "trajectory" else None  # type: ignore
+        )
 
         # Concatenate prompts and responses to form the full sequence
         batch_seq = torch.cat([batch_input_ids, batch_response_ids], dim=-1)
@@ -882,7 +1081,12 @@ class AgentModeDaemon:
                 "position_ids": position_ids,
                 "is_drop_mask": is_drop_mask,
                 "token_level_scores": token_level_scores.contiguous(),
-            },
+                **(
+                    {"response_mask": response_mask}
+                    if self.trace_aggregator.get("level", "transition") == "trajectory"
+                    else {}
+                ),
+            },  # type: ignore
             batch_size=n_transition,
         )
         data_proto = DataProto(batch=batch)
@@ -894,12 +1098,38 @@ class AgentModeDaemon:
             "training/n_rollouts_w_reward": sample_with_reward_count,
             "training/n_truncated_triplets": n_trunc_sample_because_of_response,
             "training/n_triplets": n_transition,
+            # log data, only for debug testing
+            **(
+                {
+                    "training/n_unmerged_rollouts": unmerged_count,  # type: ignore
+                    "training/n_triplets_by_turn": len(response_per_turn_list),  # type: ignore
+                    "training/avg_response_length_by_turn": np.mean(response_per_turn_list),  # type: ignore
+                    "training/max_response_length_by_turn": np.max(response_per_turn_list),  # type: ignore
+                    "training/min_response_length_by_turn": np.min(response_per_turn_list),  # type: ignore
+                }
+                if self.trace_aggregator.get("level", "transition") == "trajectory"
+                else {}
+            ),
+            **(
+                {
+                    "training/template_mismatch_triplets": template_mismatch_count,  # type: ignore
+                    "training/retoken_mismatch_triplets": retoken_mismatch_count,  # type: ignore
+                    "training/others_mismatch_triplets": others_mismatch_count,  # type: ignore
+                    "training/template_mismatch_ratio": template_mismatch_count / len(response_per_turn_list),  # type: ignore
+                    "training/retoken_mismatch_ratio": retoken_mismatch_count / len(response_per_turn_list),  # type: ignore
+                    "training/others_mismatch_ratio": others_mismatch_count / len(response_per_turn_list),  # type: ignore
+                }
+                if self.trace_aggregator.get("level", "transition") == "trajectory"
+                and self.trace_aggregator.get("debug", False)
+                else {}
+            ),
         }
 
         # Add non-tensor data for advantage calculation and logging
         data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)  # type: ignore
         data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)  # type: ignore
-        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)  # type: ignore
+        if self.trace_aggregator.get("level", "transition") == "transition":
+            data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)  # type: ignore
 
         return data_proto, data_metrics
 

agentlightning/verl/trainer.py

@@ -255,9 +255,18 @@ class AgentLightningTrainer(RayPPOTrainer):
                 )
                 self.agent_mode_daemon.run_until_all_finished()
                 batch, agent_metrics = self.agent_mode_daemon.get_train_data_batch(
-                    max_prompt_length=self.config.data.max_prompt_length,
-                    max_response_length=self.config.data.max_response_length,
+                    max_prompt_length=(
+                        self.config.agentlightning.trace_aggregator.trajectory_max_prompt_length
+                        if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
+                        else self.config.data.max_prompt_length
+                    ),
+                    max_response_length=(
+                        self.config.agentlightning.trace_aggregator.trajectory_max_response_length
+                        if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
+                        else self.config.data.max_response_length
+                    ),
                     device=gen_batch.batch["fake_ids"].device,
+                    global_steps=self.global_steps,
                 )
                 metrics.update(agent_metrics)
                 self.agent_mode_daemon.clear_data_and_server()
@@ -282,7 +291,8 @@ class AgentLightningTrainer(RayPPOTrainer):
             # uid is used for algorithm like GRPO, should be aligned to data id
             batch.non_tensor_batch["uid"] = batch.non_tensor_batch["data_id_list"]
 
-            batch.batch["response_mask"] = compute_response_mask(batch)
+            if "response_mask" not in batch.batch:
+                batch.batch["response_mask"] = compute_response_mask(batch)
 
             # compute global_valid tokens
             batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
@@ -358,7 +368,7 @@ class AgentLightningTrainer(RayPPOTrainer):
             # Calculate the metrics before processing. Refer to the comments of function `compute_data_metrics` for details.
             metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, suffix="_before_processing"))
 
-            # after advantages are assinged, we begin to drop (1) long prompt (2) floor to ppo minisize
+            # after advantages are assigned, we begin to drop (1) long prompt (2) floor to ppo minisize
             keep_indices = (~batch.batch["is_drop_mask"]).nonzero(as_tuple=True)[0]
             metrics["training/n_triplets_prompt_too_long"] = (
                 batch.batch["is_drop_mask"].shape[0] - keep_indices.shape[0]
@@ -466,6 +476,7 @@ class AgentLightningTrainer(RayPPOTrainer):
             adapter=self.adapter,
             processor=self.processor,  # For Qwen2-VL mrope position_ids
             image_base_dir=getattr(self.config.data, "image_base_dir", None),
+            trace_aggregator=self.config.agentlightning.trace_aggregator,
         )
         self.agent_mode_daemon.start()
 

contrib/.gitignore

@@ -1 +1,5 @@
 # Put contrib-related gitignore files here.
+
+# recipes/envs related
+recipes/envs/agl_envs/
+recipes/envs/wandb/

contrib/agentlightning/contrib/adapter/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Namespace package for agentlightning.contrib.adapter.

contrib/agentlightning/contrib/adapter/agentos.py

@@ -0,0 +1,127 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+FlightRecorderAdapter - Import Audit Logs to LightningStore
+=============================================================
+
+Adapts Agent-OS Flight Recorder to Agent-Lightning store format.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import Any, Dict, List
+
+logger = logging.getLogger(__name__)
+
+
+class FlightRecorderAdapter:
+    """
+    Import Agent-OS Flight Recorder logs to LightningStore.
+
+    Example:
+        >>> from agent_os import FlightRecorder
+        >>>
+        >>> recorder = FlightRecorder()
+        >>> adapter = FlightRecorderAdapter(recorder)
+        >>>
+        >>> # Import to Lightning store
+        >>> adapter.import_to_store(lightning_store)
+    """
+
+    def __init__(
+        self,
+        flight_recorder: Any,
+        *,
+        trace_id_prefix: str = "agentos",
+    ):
+        """
+        Initialize adapter.
+
+        Args:
+            flight_recorder: Agent-OS FlightRecorder
+            trace_id_prefix: Prefix for trace IDs
+        """
+        self.recorder = flight_recorder
+        self.trace_id_prefix = trace_id_prefix
+        self._imported_count = 0
+
+    def _convert_entry(self, entry: Any, index: int) -> Dict[str, Any]:
+        """Convert Flight Recorder entry to span format."""
+        entry_type = getattr(entry, "type", "unknown")
+        timestamp = getattr(entry, "timestamp", datetime.now(timezone.utc))
+        agent_id = getattr(entry, "agent_id", "unknown")
+
+        span = {
+            "span_id": f"{self.trace_id_prefix}-{index}",
+            "trace_id": f"{self.trace_id_prefix}-{agent_id}",
+            "name": f"agent_os.{entry_type}",
+            "start_time": timestamp.isoformat() if hasattr(timestamp, "isoformat") else str(timestamp),
+            "attributes": {
+                "agent_os.entry_type": entry_type,
+                "agent_os.agent_id": agent_id,
+            },
+        }
+
+        # Add type-specific attributes
+        if entry_type == "policy_check":
+            span["attributes"].update(
+                {
+                    "agent_os.policy_name": getattr(entry, "policy_name", "unknown"),
+                    "agent_os.policy_violated": getattr(entry, "violated", False),
+                }
+            )
+        elif entry_type == "signal":
+            span["attributes"].update(
+                {
+                    "agent_os.signal_type": getattr(entry, "signal", "unknown"),
+                }
+            )
+
+        return span
+
+    def get_spans(self) -> List[Dict[str, Any]]:
+        """Get all entries as spans."""
+        entries = []
+        if hasattr(self.recorder, "get_entries"):
+            entries = self.recorder.get_entries()
+        elif hasattr(self.recorder, "entries"):
+            entries = self.recorder.entries
+
+        return [self._convert_entry(e, i) for i, e in enumerate(entries)]
+
+    def import_to_store(self, store: Any) -> int:
+        """
+        Import spans to LightningStore.
+
+        Args:
+            store: LightningStore instance
+
+        Returns:
+            Number of spans imported
+        """
+        spans = self.get_spans()
+
+        for span in spans:
+            try:
+                if hasattr(store, "emit_span"):
+                    store.emit_span(span)
+                elif hasattr(store, "add_span"):
+                    store.add_span(span)
+            except Exception as e:
+                logger.error(f"Failed to import span: {e}")
+
+        self._imported_count += len(spans)
+        logger.info(f"Imported {len(spans)} spans to LightningStore")
+        return len(spans)
+
+    def get_violation_summary(self) -> Dict[str, Any]:
+        """Get summary of policy violations."""
+        spans = self.get_spans()
+        violations = [s for s in spans if s["attributes"].get("agent_os.policy_violated", False)]
+        return {
+            "total_entries": len(spans),
+            "total_violations": len(violations),
+            "violation_rate": len(violations) / len(spans) if len(spans) > 0 else 0.0,
+        }

contrib/agentlightning/contrib/adapter/triplet_group.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from agentlightning.adapter.triplet import TracerTraceToTriplet
+from agentlightning.types import Span, Triplet
+
+
+class TracerTraceToTripletGroup(TracerTraceToTriplet):
+    """Convert tracer-emitted spans into triplet trajectories.
+
+    Attributes:
+        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__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _extract_span_groups(self, spans):
+        def resolve_step_count(span, next_span, spans, index):
+            """
+            Determine step_count for a given span using next_span or fallback search.
+            """
+            # CASE A: If next_span exists and parent_id matches
+            if next_span and span.parent_id == next_span.span_id:
+                return next_span.attributes.get("step_count")
+
+            # CASE B: Fallback — search forward for agentlightning.operation
+            for s in spans[index + 1 :]:
+                if s.name == "agentlightning.operation" and span.parent_id == s.span_id:
+                    return s.attributes.get("step_count")
+
+            return None
+
+        def extract_step_count_from_links(span):
+            """
+            Extract step_count from agentlightning.link.* attributes.
+            """
+            key = span.attributes.get("agentlightning.link.0.key_match")
+            if key == "step_count":
+                return span.attributes.get("agentlightning.link.0.value_match")
+            return None
+
+        span_groups = {}
+
+        for i, span in enumerate(spans):
+            next_span = spans[i + 1] if i + 1 < len(spans) else None
+            step_count = None
+
+            if span.name == "openai.chat.completion":
+                step_count = resolve_step_count(span, next_span, spans, i)
+                if step_count is None:
+                    continue
+
+                step_count = str(step_count)
+                span_groups.setdefault(step_count, {})
+                span_groups[step_count]["call_span"] = span
+
+            elif span.name == "agentlightning.object":
+                step_count = extract_step_count_from_links(span)
+                if step_count is None:
+                    continue
+
+                step_count = str(step_count)
+                span_groups.setdefault(step_count, {})
+                span_groups[step_count]["object_span"] = span
+
+            elif span.name == "agentlightning.annotation":
+                step_count = extract_step_count_from_links(span)
+                if step_count is None:
+                    continue
+
+                step_count = str(step_count)
+                span_groups.setdefault(step_count, {})
+                span_groups[step_count]["annotation_span"] = span
+
+        return span_groups
+
+    def adapt_group(self, source: Sequence[Span], /) -> List[Triplet]:
+        span_groups = self._extract_span_groups(source)
+
+        def token_ids(span: Optional[Span], key: str) -> list:
+            return span.attributes.get(key, []) if span else []
+
+        def reward0(span: Optional[Span]) -> float:
+            if not span:
+                return 0.0
+            return float(span.attributes.get("agentlightning.reward.0.value", 0.0))
+
+        def reward1(span: Optional[Span]) -> Optional[float]:
+            if not span:
+                return 0.0
+            return float(span.attributes.get("agentlightning.reward.1.value", 0.0))
+
+        def message(span: Optional[Span]) -> Optional[str]:
+            if not span:
+                return ""
+            return span.attributes.get("agentlightning.object.literal", "")
+
+        triplets: List[Triplet] = []
+
+        for group in span_groups.values():
+            call_span = group.get("call_span")
+            if not token_ids(call_span, "prompt_token_ids") and not token_ids(call_span, "response_token_ids"):
+                continue
+
+            object_span = group.get("object_span")
+            annotation_span = group.get("annotation_span")
+            request_id = group.get("request_id")
+
+            triplets.append(
+                Triplet(
+                    prompt={"token_ids": token_ids(call_span, "prompt_token_ids")},
+                    response={"token_ids": token_ids(call_span, "response_token_ids")},
+                    reward=reward0(annotation_span),
+                    metadata={
+                        "response_id": request_id,
+                        "intrinsic_reward": reward1(annotation_span),
+                        "message": message(object_span),
+                    },
+                )
+            )
+
+        return triplets

contrib/agentlightning/contrib/agent/env_agent.py

@@ -0,0 +1,161 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict
+
+import numpy as np
+from add_instruction import add_chat_instruction, add_single_instruction
+from agl_envs import make_env_manager
+from autogen_agentchat.agents import AssistantAgent
+from autogen_core.models import ModelFamily
+from autogen_ext.models.openai import OpenAIChatCompletionClient
+
+from agentlightning import LLM, LitAgent, NamedResources, Rollout, configure_logger, emit_object, emit_reward, operation
+from agentlightning.utils.otel import make_link_attributes
+from contrib.recipes.envs.prompt_builder import HistoryPromptBuilder
+
+logger = configure_logger(name=__name__, level=logging.ERROR)
+
+
+class EnvAgent(LitAgent):
+    def __init__(self, config, trained_agents: str | None = None) -> None:
+        super().__init__(trained_agents=trained_agents)
+        self.config = config
+        self.env = None
+
+    def _build_agent(self, llm: LLM, temperature: float):
+        model_client = OpenAIChatCompletionClient(
+            model=llm.model,
+            base_url=llm.endpoint,
+            api_key=os.environ.get("OPENAI_API_KEY", "token-abc123"),
+            model_info={
+                "vision": False,
+                "function_calling": True,
+                "json_output": False,
+                "family": ModelFamily.UNKNOWN,
+                "structured_output": False,
+            },
+            temperature=temperature,
+        )
+
+        return AssistantAgent(
+            name="envs",
+            model_client=model_client,
+        )
+
+    def _get_instructed_prompt(self, prompt, sep="\n\n"):
+        """Return instructed observation based on prompt_type and captioner type."""
+        prompt_type = self.config.captioner.prompt_type
+        cap_type = self.config.captioner.type
+
+        if prompt_type == "chat":
+            if cap_type == "cot":
+                return add_chat_instruction(prompt, "cot", sep, self.config.env_name)
+            elif cap_type == "naive":
+                return add_chat_instruction(prompt, "naive", sep)
+
+        elif prompt_type == "single":
+            if cap_type == "cot":
+                return add_single_instruction(prompt, "cot", sep, self.config.env_name)
+            elif cap_type == "naive":
+                return add_single_instruction(prompt, "naive", sep, self.config.env_name)
+
+        raise ValueError(f"Unsupported prompt_type={prompt_type}, type={cap_type}")
+
+    async def rollout_async(
+        self,
+        task: Dict[str, Any],
+        resources: NamedResources,
+        rollout: Rollout,
+    ) -> float | None:
+        rollout_id = rollout.rollout_id
+        logger.info(f"[Rollout {rollout_id}] Task: {task}")
+
+        format_penalty = float(self.config["format_penalty"])
+        reward_scale = float(self.config["reawrd_scale"])
+
+        # Setup agent
+        llm: LLM = resources.get("main_llm")
+        print("Training with model:", llm.model, "on endpoint:", llm.endpoint)
+        self.agent = self._build_agent(llm, 1.0 if rollout.mode == "train" else 0.4)
+        if "max_tokens" in self.config and self.config["max_tokens"] > -1:
+            self.agent._model_client.max_tokens = self.config["max_tokens"]
+
+        try:
+            # Setup environment
+            prompt_builder = HistoryPromptBuilder(
+                max_history=self.config.captioner.max_history, prompt_type=self.config.captioner.prompt_type
+            )
+
+            self.env = make_env_manager(self.config.env_name, task, self.config)
+            env_obs, infos, available_actions_hint = self.env.reset()
+
+            prompt_builder.init(self.env)
+            prompt_builder.update_observation(env_obs)
+            prompt_builder.update_admissible_actions(available_actions_hint)
+
+            prompt = prompt_builder.get_prompt()
+
+            episode_reward, done = 0.0, False
+
+            step_count = 0
+            while not done:
+                try:
+                    instructed_prompt = self._get_instructed_prompt(prompt)
+
+                    # Main agent step
+                    with operation(step_count=step_count):
+                        result = await self.agent._model_client.create(instructed_prompt)
+                    output = result.content
+                    logger.info(f"[LLM output]: {output}")
+
+                except Exception as e:
+                    logger.error(f"[Rollout {rollout_id}] Error during training rollout: {e}", exc_info=True)
+                    break
+
+                if self.config.log_env_obs:
+                    emit_object(env_obs, attributes=make_link_attributes({"step_count": str(step_count)}))
+
+                env_obs, executed_action, is_valid, step_reward, terminated, truncated, info, available_actions_hint = (
+                    self.env.step(
+                        output,
+                        use_reasoning=self.config.captioner.type == "cot",
+                        use_success_rate=self.config.use_success_rate,
+                    )
+                )
+
+                prompt_builder.update_step_count()
+                prompt_builder.update_action(executed_action)
+                prompt_builder.update_observation(env_obs)
+                prompt_builder.update_admissible_actions(available_actions_hint)
+
+                prompt = prompt_builder.get_prompt()
+
+                if rollout.mode == "train":
+                    step_reward *= reward_scale
+
+                if format_penalty != 0.0:
+                    emit_reward(
+                        {
+                            "extrinsic_reward": step_reward,
+                            "intrinsic_reward": 0.0 if is_valid else -1.0 * format_penalty,
+                        },
+                        primary_key="extrinsic_reward",
+                        attributes=make_link_attributes({"step_count": str(step_count)}),
+                    )
+                else:
+                    emit_reward(step_reward, attributes=make_link_attributes({"step_count": str(step_count)}))
+
+                episode_reward += float(step_reward)
+                done = np.logical_or(terminated, truncated)
+
+                step_count += 1
+
+            return episode_reward
+
+        finally:
+            if self.env is not None:
+                self.env.close()

contrib/agentlightning/contrib/algorithm/env_verl/daemon.py

@@ -0,0 +1,872 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import json
+import random
+import socket
+import threading
+import time
+import uuid
+from collections import defaultdict
+from collections.abc import Mapping
+from typing import Any, Dict, List, Literal, Optional, Tuple, cast
+
+import numpy as np
+import requests
+import torch
+from flask import Flask, Response, abort, request
+from tensordict import TensorDict
+from verl import DataProto
+
+from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy
+from agentlightning.adapter.triplet import TraceToTripletBase
+from agentlightning.llm_proxy import LLMProxy, ModelConfig
+from agentlightning.reward import find_final_reward
+from agentlightning.store.base import LightningStore
+from agentlightning.types import EnqueueRolloutRequest, Rollout, RolloutConfig, Task
+from contrib.agentlightning.contrib.adapter.triplet_group import TracerTraceToTripletGroup
+
+__all__ = [
+    "AgentModeDaemon",
+    "get_left_padded_ids_and_attention_mask",
+    "get_right_padded_ids_and_attention_mask",
+]
+
+
+def get_left_padded_ids_and_attention_mask(
+    ids: List[int], max_length: int, pad_token_id: int
+) -> Tuple[List[int], List[int]]:
+    """
+    Left-pad (or truncate) a sequence of token IDs to a fixed length,
+    and build the corresponding attention mask.
+
+    Args:
+        ids:             the original list of token IDs.
+        max_length:      desired total length after padding/truncation.
+        pad_token_id:    ID to use for padding.
+
+    Returns:
+        padded_ids (any):      list of length == max_length.
+        attention_mask (any):  list of same length: 1 for non-pad tokens, 0 for pads.
+    """
+    seq_len = len(ids)
+
+    if seq_len >= max_length:
+        # too long → truncate from the left, keep the last max_length tokens
+        trimmed = ids[-max_length:]
+        attention_mask = [1] * max_length
+        return trimmed, attention_mask
+
+    # too short → pad on the left
+    pad_len = max_length - seq_len
+    padded_ids = [pad_token_id] * pad_len + ids
+    attention_mask = [0] * pad_len + [1] * seq_len
+    return padded_ids, attention_mask
+
+
+def get_right_padded_ids_and_attention_mask(
+    ids: List[int], max_length: int, pad_token_id: int
+) -> Tuple[List[int], List[int]]:
+    """
+    Right-pad (or truncate) a sequence of token IDs to a fixed length,
+    and build the corresponding attention mask.
+
+    Args:
+        ids:            the original list of token IDs.
+        max_length:     desired total length after padding/truncation.
+        pad_token_id:   ID to use for padding.
+
+    Returns:
+        padded_ids (any):     list of length == max_length.
+        attention_mask (any): list of same length: 1 for non-pad tokens, 0 for pads.
+    """
+    seq_len = len(ids)
+
+    if seq_len >= max_length:
+        # too long → truncate to the first max_length tokens
+        trimmed = ids[:max_length]
+        attention_mask = [1] * max_length
+        return trimmed, attention_mask
+
+    # too short → pad on the right
+    pad_len = max_length - seq_len
+    padded_ids = ids + [pad_token_id] * pad_len
+    attention_mask = [1] * seq_len + [0] * pad_len
+    return padded_ids, attention_mask
+
+
+def _find_available_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+def _to_native(obj: Any) -> Any:
+    """Convert data retrieved from Parquet to data usable in AGL server."""
+    # 1) Arrays -> list (then recurse)
+    if isinstance(obj, np.ndarray):
+        return _to_native(obj.tolist())
+
+    # 2) NumPy scalar types -> Python scalars
+    if isinstance(obj, np.generic):
+        return _to_native(obj.item())
+
+    # 3) Dict-like -> dict
+    if isinstance(obj, Mapping):
+        return {_to_native(k): _to_native(v) for k, v in obj.items()}  # type: ignore
+
+    # 4) Lists/Tuples/Sets -> list
+    if isinstance(obj, (list, tuple, set)):
+        return [_to_native(x) for x in obj]  # type: ignore
+
+    # 5) Anything else: leave as-is
+    return obj
+
+
+class EnvAgentModeDaemon:
+    """
+    AgentModeDaemon using the AgentLightningServer SDK.
+
+    This class manages the server lifecycle, task queueing, and results
+    retrieval, while also running a proxy server for LLM requests. It maintains
+    the original interface for compatibility with the RayPPOTrainer.
+    """
+
+    def __init__(
+        self,
+        port: Optional[int],
+        train_rollout_n: int,
+        train_information: Dict[str, Any],
+        tokenizer: Any,
+        mini_batch_size: int,
+        pad_token_id: int,
+        reward_fillna_value: float = 0.0,
+        llm_timeout_seconds: float = 1200.0,
+        mode: Literal["v0", "v1"] = "v1",
+        llm_proxy: LLMProxy | None = None,
+        store: LightningStore | None = None,
+        adapter: TraceToTripletBase | None = None,
+    ):
+        self.mode = mode
+        self.llm_timeout_seconds = llm_timeout_seconds
+
+        # Server and Task Configuration
+        if mode == "v0":
+            assert port is not None
+            self.server_port = port
+            self.server = AgentLightningServer(
+                host="0.0.0.0", port=self.server_port, task_timeout_seconds=self.llm_timeout_seconds
+            )
+            self.proxy_port = _find_available_port()  # Run proxy on a different port
+        else:
+            assert store is not None
+            self.store = store
+            if llm_proxy is None:
+                self.llm_proxy = LLMProxy(
+                    port=_find_available_port(),
+                    model_list=[],
+                    store=store,
+                )
+            else:
+                # Reuse the existing LLM proxy (probably configured by user)
+                self.llm_proxy = llm_proxy
+
+            # if adapter is None:
+            #     self.adapter = TracerTraceToTripletGroup()
+            # else:
+            #     # Reuse the one from trainer
+            #     self.adapter = adapter
+            self.adapter = TracerTraceToTripletGroup()
+
+            self._internal_loop: Optional[asyncio.AbstractEventLoop] = None
+            self._internal_loop_thread = threading.Thread(target=self._internal_loop_runner, daemon=True)
+            self._internal_loop_thread.start()
+
+        # Training and Data Configuration
+        self.train_rollout_n = train_rollout_n
+        self.train_information = train_information
+        self.mini_batch_size = mini_batch_size
+        self.pad_token_id = pad_token_id
+        self.tokenizer = tokenizer
+        self.reward_fillna_value = reward_fillna_value
+
+        # Internal State
+        self.backend_llm_server_addresses: List[str] = []
+        self._total_tasks_queued = 0
+        self._completed_rollouts_v0: Dict[str, RolloutLegacy] = {}
+        self._task_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
+        self._server_thread: Optional[threading.Thread] = None
+        self._proxy_thread: Optional[threading.Thread] = None
+        self.is_train = True
+
+    def _internal_loop_runner(self):
+        """Run the internal loop."""
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        self._internal_loop = loop
+        loop.run_forever()
+        loop.close()
+
+    def _start_proxy_server_v0(self):
+        """
+        Initializes and runs a Flask-based proxy server in a separate thread.
+        This proxy load-balances requests to the actual backend LLM servers.
+        """
+        app = Flask(__name__)
+
+        num_requests = 0
+        last_request_time = 0
+
+        @app.route("/v1/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
+        def proxy(path: str):  # type: ignore
+            if not self.backend_llm_server_addresses:
+                abort(503, description="No backend LLM servers available.")
+
+            # Randomly choose a backend server for load balancing
+            target_server = random.choice(self.backend_llm_server_addresses)
+            target_url = f"http://{target_server}/v1/{path}"
+
+            # Copy client request headers, removing the Host header
+            headers = {key: value for key, value in request.headers if key.lower() != "host"}
+
+            # Log the request for debugging
+            nonlocal num_requests, last_request_time
+            current_time = time.time()
+            num_requests += 1
+            if current_time - last_request_time > 60 or num_requests == 1 or num_requests % 100 == 0:
+                print(f"Proxying {request.method} request to {target_server}. Request data: {request.get_data()}")
+            last_request_time = current_time
+
+            try:
+                # Forward the request to the target backend
+                resp = requests.request(
+                    method=request.method,
+                    url=target_url,
+                    headers=headers,
+                    params=request.args,  # type: ignore
+                    data=request.get_data(),
+                    cookies=request.cookies,
+                    allow_redirects=False,
+                    timeout=self.llm_timeout_seconds,
+                )
+                # Filter out hop-by-hop headers before returning the response
+                excluded_headers = [
+                    "content-encoding",
+                    "content-length",
+                    "transfer-encoding",
+                    "connection",
+                    "keep-alive",
+                    "proxy-authenticate",
+                    "proxy-authorization",
+                    "te",
+                    "trailers",
+                    "upgrade",
+                ]
+                response_headers = [
+                    (name, value) for name, value in resp.raw.headers.items() if name.lower() not in excluded_headers
+                ]
+                if resp.status_code == 200:
+                    # NOTE: from Zhiyuan's code.
+                    # https://github.com/hzy46/verl_agent_mode/blob/2db65ea9858f645a914120357412a7540f8bd82d/verl/trainer/ppo/ray_trainer.py#L692-L711
+                    # request_json = json.loads(request.get_data().decode("utf-8"))
+                    response_json = json.loads(resp.content.decode("utf-8"))
+                    # response_message = ChatCompletion(**response_json).choices[0].message.model_dump(exclude_unset=True, exclude_none=True)
+                    # tool_schemas = request_json.get("tools", None)
+                    # prompt_ids = self.tokenizer.apply_chat_template(request_json["messages"], tools=tool_schemas, add_generation_prompt=True, tokenize=True)
+                    # full_ids = self.tokenizer.apply_chat_template(request_json["messages"] + [response_message], tools=tool_schemas, add_generation_prompt=False, tokenize=True)
+                    # TBD: response_ids sometimes ends with "<eos_id>\n", shall we keep the extra "\n"?
+                    # sometimes it has some differences with the hacky method in the end, but this should align with ToolCompletionCallback
+                    # response_ids = full_ids[len(prompt_ids):]
+
+                    # NOTE (yuge): They are different. Don't know why.
+                    # assert response_json['prompt_token_ids'] == prompt_ids
+                    # patched_response_ids = response_json['response_token_ids'][0]
+                    # assert patched_response_ids == response_ids[:len(patched_response_ids)], f"{patched_response_ids} != {response_ids[:len(patched_response_ids)]}"
+                    # response_json['prompt_token_ids'] = prompt_ids
+                    # response_json['response_token_ids'] = [response_ids]
+                    replaced_return_content = json.dumps(response_json).encode("utf-8")
+                    return Response(replaced_return_content, status=resp.status_code, headers=response_headers)
+                return Response(resp.content, resp.status_code, response_headers)
+            except requests.exceptions.RequestException as e:
+                abort(500, description=f"Error proxying request: {e}")
+
+        def run_app():
+            app.run(host="0.0.0.0", port=self.proxy_port, threaded=True, debug=False)
+
+        self._proxy_thread = threading.Thread(target=run_app, daemon=True)
+        self._proxy_thread.start()
+        print(f"Proxy server running on port {self.proxy_port}")
+
+    async def _update_proxy_server_v1(self):
+        model_name = self.train_information.get("model")
+        if not model_name:
+            raise ValueError("Model name is not set.")
+        self.llm_proxy.update_model_list(
+            [
+                ModelConfig(
+                    {
+                        "model_name": model_name,
+                        "litellm_params": {
+                            "model": "hosted_vllm/" + model_name,
+                            "api_base": f"http://{address}/v1/",
+                        },
+                    }
+                )
+                for address in self.backend_llm_server_addresses
+            ],
+        )
+
+        await self.llm_proxy.restart()
+
+    def start(self):
+        """Starts the main AgentLightningServer and the proxy server."""
+
+        if self.mode == "v0":
+
+            def run_server():
+                """Run the AgentLightningServer in a separate thread."""
+                asyncio.run(self.server.run_forever())
+
+            self._server_thread = threading.Thread(target=run_server, daemon=True)
+            self._server_thread.start()
+
+            # Wait for the server's internal startup event to be set.
+            print("Waiting for AgentLightningServer to start...")
+            is_ready = self.server.startup_event.wait(timeout=20.0)  # Wait up to 20s
+            if not is_ready:
+                raise RuntimeError("AgentLightningServer failed to start within the timeout period.")
+
+            print(f"AgentLightningServer control plane running on port {self.server_port}")
+
+            self._start_proxy_server_v0()
+        else:
+            # Agent lightning server is no longer needed;
+            # Start proxy server in _async_set_up
+            pass
+
+    async def _async_set_up(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
+        """Async helper to set up data and resources on the server."""
+        self.clear_data_and_server()
+        if server_addresses != self.backend_llm_server_addresses:
+            self.backend_llm_server_addresses = server_addresses
+            if self.mode == "v1" and not self.llm_proxy.is_running():
+                await self._update_proxy_server_v1()
+        self.is_train = is_train
+
+        # 1. Update resources on the server for clients to use
+        if self.mode == "v0":
+            llm_resource = LLM(
+                endpoint=f"http://127.0.0.1:{self.proxy_port}/v1",
+                model=self.train_information.get("model", "default-model"),
+                sampling_parameters={
+                    "temperature": self.train_information.get("temperature", 0.7 if is_train else 0.0)
+                },
+            )
+        else:
+            llm_resource = self.llm_proxy.as_resource(
+                sampling_parameters={
+                    "temperature": self.train_information.get("temperature", 0.7 if is_train else 0.0)
+                },
+            )
+
+        resources: NamedResources = {"main_llm": llm_resource}
+
+        if self.mode == "v0":
+            resources_id = await self.server.update_resources(resources)
+        else:
+            resources_update = await self.store.add_resources(resources)
+            resources_id = resources_update.resources_id
+
+        # 2. Queue tasks for agents to process
+        keys = list(data.keys())
+        num_samples = len(data[keys[0]])
+        rollouts_per_sample = self.train_rollout_n if is_train else 1
+
+        enqueue_rollout_requests: List[EnqueueRolloutRequest] = []
+        data_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
+
+        for i in range(num_samples):
+            data_id = str(uuid.uuid4())
+            original_sample = {key: data[key][i] for key in keys}
+            original_sample["data_id"] = data_id
+            data_id_to_original_sample[data_id] = original_sample
+
+            # For training, each sample is rolled out multiple times
+            # Data ID is different from Rollout ID, as one data can have multiple rollouts.
+            for _ in range(rollouts_per_sample):
+                task_metadata = {"data_id": data_id, "is_train": is_train}
+                if self.mode == "v0":
+                    # Queue immediately
+                    rollout_id = await self.server.queue_task(
+                        sample=_to_native(original_sample),
+                        mode="train" if is_train else "val",
+                        resources_id=resources_id,
+                        metadata=task_metadata,
+                    )
+
+                    # Store original sample data to reconstruct batch information later
+                    self._task_id_to_original_sample[rollout_id] = original_sample
+                    self._total_tasks_queued += 1
+                else:
+                    # Collect tasks to enqueue in batch and queue them later
+                    enqueue_rollout_requests.append(
+                        EnqueueRolloutRequest(
+                            input=_to_native(original_sample),
+                            mode="train" if is_train else "val",
+                            resources_id=resources_id,
+                            config=RolloutConfig(
+                                unresponsive_seconds=self.llm_timeout_seconds,
+                                timeout_seconds=self.llm_timeout_seconds,
+                            ),
+                            metadata=task_metadata,
+                        )
+                    )
+
+        if self.mode == "v1":
+            # Enqueue all the tasks in a single batch
+            rollouts = await self.store.enqueue_many_rollouts(enqueue_rollout_requests)
+            self._task_id_to_original_sample.update(
+                {
+                    # Recover the original data and store it for later use.
+                    rollout.rollout_id: data_id_to_original_sample[cast(Dict[str, Any], rollout.metadata)["data_id"]]
+                    for rollout in rollouts
+                }
+            )
+            self._total_tasks_queued += len(rollouts)
+
+    def set_up_data_and_server(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
+        """Synchronous wrapper for setting up data and server resources."""
+        coro = self._async_set_up(data, server_addresses, is_train)
+
+        if self.mode == "v0":
+            if not self.server.loop or not self.server.startup_event.is_set():
+                raise RuntimeError("Server is not running or ready.")
+
+            future = asyncio.run_coroutine_threadsafe(coro, self.server.loop)
+
+        else:
+            if self._internal_loop is None:
+                raise RuntimeError("Internal loop is not running.")
+            future = asyncio.run_coroutine_threadsafe(coro, self._internal_loop)
+        try:
+            future.result(timeout=60)  # Wait for completion with a timeout
+        except Exception as e:
+            print(f"Failed to set up data on server: {e}")
+            raise
+
+    def _validate_data(self, rollout: RolloutLegacy):
+        if rollout.final_reward is None:
+            print(
+                f"Warning: Reward is None for rollout {rollout.rollout_id}, will be auto-set to {self.reward_fillna_value}."
+            )
+        if rollout.triplets is None:
+            print(f"Warning: Triplet is None for rollout {rollout.rollout_id}.")
+        elif len(rollout.triplets) == 0:
+            print(f"Warning: Length of triplets is 0 for rollout {rollout.rollout_id}.")
+        elif any(not r.response.get("token_ids", []) for r in rollout.triplets):
+            print(f"Warning: Rollout {rollout.rollout_id} contains empty response: {rollout.triplets}")
+        elif any(not r.prompt.get("token_ids", []) for r in rollout.triplets):
+            print(f"Warning: Rollout {rollout.rollout_id} contains empty prompt: {rollout.triplets}")
+
+    async def _validate_data_v1(self, rollout: Rollout) -> RolloutLegacy:
+        """Convert Rollout to RolloutLegacy and validate.
+
+        1. Task: construct from Rollout
+        2. Triplets: obtained by querying spans and feeding into the adapter
+        3. Final reward: extracted from last triplet's reward, searching backwards if not found
+        """
+        # Query spans for this rollout (latest attempt)
+        spans = await self.store.query_spans(rollout.rollout_id, attempt_id="latest")
+        final_reward = find_final_reward(spans)
+
+        # Convert spans to triplets using the adapter
+        if not spans:
+            # No triplets found, will emit a warning later.
+            triplets = []
+        else:
+            # triplets = self.adapter.adapt(spans)
+            triplets = self.adapter.adapt_group(spans)
+
+        # # Extract final reward from triplets
+        # final_reward: Optional[float] = None
+        # if triplets:
+        #     # Search backwards through triplets for the first non-None reward
+        #     for triplet in reversed(triplets):
+        #         if triplet.reward is not None:
+        #             final_reward = triplet.reward
+        #             break
+
+        # Construct the Task object from Rollout
+        task = Task(
+            rollout_id=rollout.rollout_id,
+            input=rollout.input,
+            mode=rollout.mode,
+            resources_id=rollout.resources_id,
+            metadata=rollout.metadata or {},
+        )
+
+        # Create the Rollout object (without trace and logs as per user's note)
+        result_rollout = RolloutLegacy(
+            rollout_id=rollout.rollout_id,
+            task=task,
+            final_reward=final_reward,
+            triplets=triplets,
+            metadata=rollout.metadata or {},
+        )
+
+        # Run the same validation as v0
+        self._validate_data(result_rollout)
+
+        return result_rollout
+
+    async def _async_run_until_finished(self, verbose: bool = True):
+        """Async helper to wait for all tasks to complete."""
+        while len(self._completed_rollouts_v0) < self._total_tasks_queued:
+            if self.mode == "v0":
+                completed_batch = await self.server.retrieve_completed_rollouts()
+            else:
+                completed_batch = await self.store.wait_for_rollouts(
+                    rollout_ids=list(self._task_id_to_original_sample.keys()), timeout=0
+                )
+            for rollout in completed_batch:
+                if rollout.rollout_id in self._completed_rollouts_v0:
+                    # Already processed, skip
+                    continue
+                if isinstance(rollout, Rollout):
+                    rollout = await self._validate_data_v1(rollout)
+                else:
+                    self._validate_data(rollout)
+                if rollout.rollout_id not in self._task_id_to_original_sample:
+                    print(f"Warning: Received unknown rollout ID {rollout.rollout_id}, skipping.")
+                else:
+                    self._completed_rollouts_v0[rollout.rollout_id] = rollout
+            if verbose:
+                print(f"Completed {len(self._completed_rollouts_v0)}/{self._total_tasks_queued} tasks...")
+            await asyncio.sleep(5)
+
+        print("All tasks finished.")
+
+    def run_until_all_finished(self, verbose: bool = True):
+        """Synchronously waits for all queued tasks to be completed and reported."""
+        if self._total_tasks_queued == 0:
+            print("Warning: No tasks were queued.")
+            return
+
+        if self.mode == "v0":
+            if not self.server.loop or not self.server.startup_event.is_set():
+                raise RuntimeError("Server is not running or ready.")
+            loop = self.server.loop
+        else:
+            loop = self._internal_loop
+            assert loop is not None
+
+        coro = self._async_run_until_finished(verbose)
+        future = asyncio.run_coroutine_threadsafe(coro, loop)
+        try:
+            future.result()  # Wait indefinitely for all tasks to complete
+        except Exception as e:
+            print(f"Error while waiting for tasks to finish: {e}")
+            raise
+
+    def get_test_metrics(self):
+        """Calculates and returns metrics for a validation run."""
+        assert not self.is_train, "This method should only be called during validation."
+        assert len(self._completed_rollouts_v0) == self._total_tasks_queued
+
+        sample_stat_list: List[Dict[str, Any]] = []
+        sample_stat_list_by_source: Dict[str, List[Dict[str, Any]]] = defaultdict(
+            list
+        )  # FIXME: Evaluate whether grouping stats by source is actually needed.
+
+        for rollout_id, rollout in self._completed_rollouts_v0.items():
+            final_reward_raw: Optional[float] = rollout.final_reward
+            final_reward = self._fillna_reward(rollout)
+            if not rollout.triplets:
+                print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
+                sample_stat_list.append({"reward": final_reward, "has_reward": final_reward_raw is not None})
+                continue
+            response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
+
+            if "data_source" in self._task_id_to_original_sample[rollout_id]:
+                # When a test sample includes a 'data_source' field, record per-source statistics for test results.
+                # TODO: This is a flawed design. We should have a better way to handle this.
+                data_source = self._task_id_to_original_sample[rollout_id]["data_source"]
+                sample_stat_list_by_source[data_source].append(
+                    {
+                        "sum_response_length": np.sum(response_length_list),
+                        "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
+                        "turn_count": len(rollout.triplets),
+                        "reward": final_reward,
+                        "has_reward": final_reward_raw is not None,
+                    }
+                )
+            sample_stat_list.append(
+                {
+                    "sum_response_length": np.sum(response_length_list),
+                    "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
+                    "turn_count": len(rollout.triplets),
+                    "reward": final_reward,
+                    "has_reward": final_reward_raw is not None,
+                }
+            )
+        metric_dict: Dict[str, Any] = {}
+
+        stats_w_trace = [stat for stat in sample_stat_list if "sum_response_length" in stat]
+        stats_w_trace_by_source = {
+            data_source: [stat for stat in sample_stats if "sum_response_length" in stat]
+            for data_source, sample_stats in sample_stat_list_by_source.items()
+        }
+        for data_source, sample_stats in sample_stat_list_by_source.items():
+            metric_dict.update(
+                {
+                    f"val/{data_source}/n_rollouts": len(sample_stats),
+                    f"val/{data_source}/n_rollouts_w_trace": len(stats_w_trace_by_source[data_source]),
+                    f"val/{data_source}/n_rollouts_w_reward": len(
+                        [stat for stat in sample_stats if stat["has_reward"]]
+                    ),
+                    f"val/{data_source}/reward": np.mean(
+                        [stat["reward"] for stat in sample_stats]
+                    ),  # each rollout must have a reward (fillna if missing)
+                    f"val/{data_source}/mean_response_length": np.mean(
+                        [stat["mean_response_length"] for stat in stats_w_trace_by_source[data_source]]
+                    ),
+                    f"val/{data_source}/sum_response_length": np.mean(
+                        [stat["sum_response_length"] for stat in stats_w_trace_by_source[data_source]]
+                    ),
+                    f"val/{data_source}/turn_count": np.mean(
+                        [stat["turn_count"] for stat in stats_w_trace_by_source[data_source]]
+                    ),
+                }
+            )
+        metric_dict.update(
+            {
+                "val/n_rollouts": len(sample_stat_list),
+                "val/n_rollouts_w_trace": len(stats_w_trace),
+                "val/n_rollouts_w_reward": len([stat for stat in sample_stat_list if stat["has_reward"]]),
+                "val/reward": np.mean(
+                    [stat["reward"] for stat in sample_stat_list]
+                ),  # each rollout must have a reward (fillna if missing)
+                "val/mean_response_length": np.mean([stat["mean_response_length"] for stat in stats_w_trace]),
+                "val/sum_response_length": np.mean([stat["sum_response_length"] for stat in stats_w_trace]),
+                "val/turn_count": np.mean([stat["turn_count"] for stat in stats_w_trace]),
+            }
+        )
+        return metric_dict
+
+    def get_train_data_batch(
+        self,
+        max_prompt_length: int,
+        max_response_length: int,
+        device: torch.device,
+        use_final_reward_as_step_reward: bool = True,
+        use_intrinsic_reward: bool = False,
+        is_gigpo: bool = False,
+    ):
+        """
+        Processes completed rollouts to generate a training data batch.
+
+        This function reconstructs the logic from the original AgentModeDaemon,
+        using data retrieved from the new server architecture. It handles padding,
+        truncation, and tensor creation for the PPO training loop.
+        """
+        assert self.is_train, "This method should only be called during training."
+        assert len(self._completed_rollouts_v0) == self._total_tasks_queued
+
+        # 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts
+        finished_id_to_sample_info: Dict[str, Dict[str, Any]] = {}
+        finished_id_to_final_reward: Dict[str, float] = {}
+        sample_with_reward_count = 0
+        for rollout_id, rollout in self._completed_rollouts_v0.items():
+            original_sample = self._task_id_to_original_sample[rollout_id]
+            sample_with_reward_count += int(rollout.final_reward is not None)
+            final_reward = self._fillna_reward(rollout)
+
+            if not rollout.triplets:
+                finished_id_to_final_reward[rollout_id] = final_reward
+                print(f"Warning: No triplets found for training rollout {rollout.rollout_id}, skipping.")
+                continue
+
+            # The client should report triplets that contain prompt_ids and response_ids.
+            # Example triplet.prompt: {"token_ids": [...]}
+            # Example triplet.response: {"token_ids": [...]}
+            # trace_list = [
+            #     {"prompt_ids": t.prompt.get("token_ids", []), "response_ids": t.response.get("token_ids", [])}
+            #     for t in rollout.triplets
+            # ]
+            trace_list = []
+            for t in rollout.triplets:
+                trace_dict = {
+                    "prompt_ids": t.prompt.get("token_ids", []),
+                    "response_ids": t.response.get("token_ids", []),
+                    "step_reward": t.reward,
+                    "step_intrinsic_reward": t.metadata.get("intrinsic_reward", 0.0),
+                    "message": t.metadata.get("message", ""),
+                }
+
+                trace_list.append(trace_dict)
+
+            info = {
+                "final_reward": final_reward,
+                "trace_list": trace_list,
+                "data_id": original_sample["data_id"],
+            }
+            finished_id_to_sample_info[rollout_id] = info
+            finished_id_to_final_reward[rollout_id] = final_reward
+        #
+        # --- Data processing and tensor creation logic ---
+        # Get all the reported data.
+        # prompt_ids are left-padded.
+        # response_ids are right-padded.
+        # They are concatenated in the middle.
+        # Discard handling:
+        #   - Those exceeding max_prompt_length will be marked for discard, but not
+        #     discarded here. They are only truncated and marked, to be discarded later.
+        #     This is for the correctness of the advantage calculation.
+        #   - The discard for the PPO mini-batch should also be handled this way.
+        input_ids_list: List[List[int]] = []
+        input_attention_mask_list: List[List[int]] = []
+        response_ids_list: List[List[int]] = []
+        response_attention_mask_list: List[List[int]] = []
+        final_reward_list: List[float] = []
+        step_reward_list: List[float] = []
+        data_id_list: List[str] = []
+        rollout_id_list: List[str] = []
+        turn_index_list: List[int] = []
+        is_drop_list: List[bool] = []
+        n_trunc_sample_because_of_response = 0
+
+        # optional fields
+        step_intrinsic_reward_list: List[float] = []
+        message_list: List[str] = []
+
+        for rollout_id, sample_info in finished_id_to_sample_info.items():
+            for turn_index, trace in enumerate(sample_info["trace_list"]):
+
+                final_reward_list.append(sample_info["final_reward"])
+                step_reward_list.append(trace["step_reward"])
+                step_intrinsic_reward_list.append(trace["step_intrinsic_reward"])
+                message_list.append(trace["message"])
+
+                prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]
+
+                # Mark samples with prompts exceeding max_prompt_length to be dropped later
+                if len(prompt_ids) > max_prompt_length:
+                    prompt_ids = prompt_ids[:max_prompt_length]
+                    is_drop_list.append(True)
+                else:
+                    is_drop_list.append(False)
+
+                # Truncate responses that exceed max_response_length
+                if len(response_ids) > max_response_length:
+                    response_ids = response_ids[:max_response_length]
+                    n_trunc_sample_because_of_response += 1
+
+                # Pad prompts to the left and responses to the right
+                one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
+                    prompt_ids, max_prompt_length, self.pad_token_id
+                )
+                one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
+                    response_ids, max_response_length, self.pad_token_id
+                )
+
+                input_ids_list.append(one_input_ids)
+                input_attention_mask_list.append(one_input_attention_mask)
+                response_ids_list.append(one_response_ids)
+                response_attention_mask_list.append(one_response_attention_mask)
+                data_id_list.append(sample_info["data_id"])
+                rollout_id_list.append(rollout_id)
+                turn_index_list.append(turn_index)
+
+        n_transition = len(input_ids_list)
+        batch_input_ids = torch.LongTensor(input_ids_list).to(device)
+        input_attention_mask = torch.LongTensor(input_attention_mask_list).to(device)
+        batch_response_ids = torch.LongTensor(response_ids_list).to(device)
+        response_attention_mask = torch.LongTensor(response_attention_mask_list).to(device)
+
+        # Concatenate prompts and responses to form the full sequence
+        batch_seq = torch.cat([batch_input_ids, batch_response_ids], dim=-1)
+        attention_mask = torch.cat([input_attention_mask, response_attention_mask], dim=-1)
+        position_ids = torch.clamp(torch.cumsum(attention_mask, dim=-1) - 1, min=0)
+        is_drop_mask = torch.BoolTensor(is_drop_list).to(device)
+        if use_final_reward_as_step_reward:
+            scores = torch.tensor(final_reward_list, dtype=torch.float32).to(device)
+        else:
+            scores = torch.tensor(step_reward_list, dtype=torch.float32).to(device)
+
+        # Create token-level scores by placing the final reward at the last token position
+        token_level_scores = torch.zeros_like(attention_mask, dtype=scores.dtype)
+        # At the eos_mask_idx position of each sample, fill in the corresponding scores.
+        # torch.arange(n_transition) generates [0,1,2,...,bsz-1] as indices for the batch dimension.
+        eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1)  # (bsz,)
+        token_level_scores[torch.arange(n_transition), eos_mask_idx] = scores
+        # Only take the last response_length part of the sequence to get the token-level scores for the model's response part.
+        token_level_scores = token_level_scores[:, -max_response_length:]
+
+        # Create token-level intrinsic rewards
+        token_level_intrinsic_rewards = None
+        if use_intrinsic_reward:
+            step_intrinsic_reward_list = [0.0 if reward is None else reward for reward in step_intrinsic_reward_list]
+            intrinsic_rewards = torch.tensor(step_intrinsic_reward_list, dtype=torch.float32).to(device)
+            token_level_intrinsic_rewards = torch.zeros_like(attention_mask, dtype=intrinsic_rewards.dtype)
+            token_level_intrinsic_rewards[torch.arange(n_transition), eos_mask_idx] = intrinsic_rewards
+            token_level_intrinsic_rewards = token_level_intrinsic_rewards[:, -max_response_length:]
+
+        # Form the final batch using TensorDict
+        batch_dict = {
+            "prompts": batch_input_ids,
+            "responses": batch_response_ids,
+            "input_ids": batch_seq,  # here input_ids become the whole sentences
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "is_drop_mask": is_drop_mask,
+            "token_level_scores": token_level_scores.contiguous(),
+        }
+        batch_dict["step_rewards"] = torch.tensor(np.array(step_reward_list), dtype=torch.float32).to(device)
+        if use_intrinsic_reward:
+            batch_dict["step_intrinsic_rewards"] = torch.tensor(
+                np.array(step_intrinsic_reward_list), dtype=torch.float32
+            ).to(device)
+            batch_dict["token_level_intrinsic_rewards"] = token_level_intrinsic_rewards.contiguous()
+
+        batch = TensorDict(batch_dict, batch_size=n_transition)
+        data_proto = DataProto(batch=batch)
+
+        data_metrics = {
+            "training/reward": np.mean(list(finished_id_to_final_reward.values())),
+            "training/n_rollouts": len(finished_id_to_final_reward),
+            "training/n_rollouts_w_trace": len(finished_id_to_sample_info),
+            "training/n_rollouts_w_reward": sample_with_reward_count,
+            "training/n_truncated_triplets": n_trunc_sample_because_of_response,
+            "training/n_triplets": n_transition,
+        }
+
+        # Add non-tensor data for advantage calculation and logging
+        data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)  # type: ignore
+        data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)  # type: ignore
+        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)  # type: ignore
+
+        data_proto.non_tensor_batch["step_rewards"] = np.array(step_reward_list)
+        if is_gigpo:
+            data_proto.non_tensor_batch["anchor_obs"] = np.array(message_list)
+
+        return data_proto, data_metrics
+
+    def clear_data_and_server(self):
+        """Resets the internal state of the daemon for the next run."""
+        self.backend_llm_server_addresses = []
+        self._completed_rollouts_v0.clear()
+        self._task_id_to_original_sample.clear()
+        self._total_tasks_queued = 0
+        # For a true reset, the server's internal queues would also need clearing.
+        # This implementation assumes that `set_up_data_and_server` is called
+        # for each new run, effectively starting a fresh batch.
+
+    def _fillna_reward(self, rollout: RolloutLegacy):
+        if rollout.final_reward is None:
+            if self.reward_fillna_value is not None:  # type: ignore
+                final_reward = self.reward_fillna_value
+            else:
+                raise ValueError(f"Reward is None for rollout {rollout.rollout_id}, please check the reward function.")
+        else:
+            final_reward = rollout.final_reward
+        return final_reward

contrib/agentlightning/contrib/algorithm/env_verl/trainer.py

@@ -0,0 +1,543 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
+from __future__ import annotations
+
+import random
+from contextlib import contextmanager
+from copy import deepcopy
+from pprint import pprint
+from typing import Dict, Tuple, Type
+
+import numpy as np
+import torch
+import verl
+from codetiming import Timer
+from omegaconf import OmegaConf
+from tqdm import tqdm
+from verl import DataProto
+from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
+from verl.trainer.ppo.core_algos import agg_loss
+from verl.trainer.ppo.metric_utils import (
+    _compute_response_info,
+    compute_throughout_metrics,
+    compute_timing_metrics,
+)
+from verl.trainer.ppo.ray_trainer import (
+    AdvantageEstimator,
+    RayPPOTrainer,
+    apply_kl_penalty,
+    compute_advantage,
+    compute_response_mask,
+)
+from verl.utils.metric import reduce_metrics
+from verl.utils.tracking import Tracking
+
+from agentlightning.adapter import TraceAdapter, TraceToTripletBase
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.store.base import LightningStore
+
+from .daemon import EnvAgentModeDaemon
+
+__all__ = [
+    "EnvAgentLightningTrainer",
+]
+
+
+@contextmanager
+def _timer(name: str, timing_raw: Dict[str, float]):
+    with Timer(name=name, logger=None) as timer:
+        yield
+    if name not in timing_raw:
+        timing_raw[name] = 0
+    timing_raw[name] += timer.last
+
+
+# This function is adapted from verl.
+# We introduce a new parameter `suffix` to distinguish between metrics computed
+# before and after AgentLightning’s post-processing.
+# - "Before" refers to raw reward and advantage values.
+# - "After" refers to values computed following post-processing, which involves:
+#     (1) Dropping prompts that exceed the maximum allowed length.
+#     (2) Adjusting the batch size to be a multiple of the mini PPO size.
+# Different suffixes are used to label these two stages accordingly.
+def compute_data_metrics(batch: DataProto, use_critic: bool = True, suffix: str = "") -> Dict[str, Any]:
+    """
+    Computes various metrics from a batch of data for PPO training.
+
+    This function calculates metrics related to scores, rewards, advantages, returns, values,
+    and sequence lengths from a batch of data. It provides statistical information (mean, max, min)
+    for each metric category.
+
+    Args:
+        batch: A DataProto object containing batch data with token-level scores, rewards, advantages, etc.
+        use_critic: Whether to include critic-specific metrics. Defaults to True.
+
+    Returns:
+        A dictionary of metrics including:
+            - critic/score/mean, max, min: Statistics about sequence scores
+            - critic/rewards/mean, max, min: Statistics about sequence rewards
+            - critic/advantages/mean, max, min: Statistics about advantages
+            - critic/returns/mean, max, min: Statistics about returns
+            - critic/values/mean, max, min: Statistics about critic values (if use_critic=True)
+            - critic/vf_explained_var: Explained variance of the value function (if use_critic=True)
+            - response_length/mean, max, min, clip_ratio: Statistics about response lengths
+            - prompt_length/mean, max, min, clip_ratio: Statistics about prompt lengths
+    """
+    sequence_score = batch.batch["token_level_scores"].sum(-1)
+    sequence_reward = batch.batch["token_level_rewards"].sum(-1)
+
+    advantages = batch.batch["advantages"]
+    returns = batch.batch["returns"]
+
+    max_response_length = batch.batch["responses"].shape[-1]
+
+    prompt_mask = batch.batch["attention_mask"][:, :-max_response_length].bool()
+    response_mask = batch.batch["attention_mask"][:, -max_response_length:].bool()
+
+    max_prompt_length = prompt_mask.size(-1)
+
+    response_info = _compute_response_info(batch)
+    prompt_length = response_info["prompt_length"]
+    response_length = response_info["response_length"]
+
+    valid_adv = torch.masked_select(advantages, response_mask)
+    valid_returns = torch.masked_select(returns, response_mask)
+
+    if use_critic:
+        values = batch.batch["values"]
+        valid_values = torch.masked_select(values, response_mask)
+        return_diff_var = torch.var(valid_returns - valid_values)
+        return_var = torch.var(valid_returns)
+
+    metrics = {
+        # score
+        "critic/score/mean" + suffix: torch.mean(sequence_score).detach().item(),
+        "critic/score/max" + suffix: torch.max(sequence_score).detach().item(),
+        "critic/score/min" + suffix: torch.min(sequence_score).detach().item(),
+        # reward
+        "critic/rewards/mean" + suffix: torch.mean(sequence_reward).detach().item(),
+        "critic/rewards/max" + suffix: torch.max(sequence_reward).detach().item(),
+        "critic/rewards/min" + suffix: torch.min(sequence_reward).detach().item(),
+        # adv
+        "critic/advantages/mean" + suffix: torch.mean(valid_adv).detach().item(),
+        "critic/advantages/max" + suffix: torch.max(valid_adv).detach().item(),
+        "critic/advantages/min" + suffix: torch.min(valid_adv).detach().item(),
+        # returns
+        "critic/returns/mean" + suffix: torch.mean(valid_returns).detach().item(),
+        "critic/returns/max" + suffix: torch.max(valid_returns).detach().item(),
+        "critic/returns/min" + suffix: torch.min(valid_returns).detach().item(),
+        **(
+            {
+                # values
+                "critic/values/mean" + suffix: torch.mean(valid_values).detach().item(),
+                "critic/values/max" + suffix: torch.max(valid_values).detach().item(),
+                "critic/values/min" + suffix: torch.min(valid_values).detach().item(),
+                # vf explained var
+                "critic/vf_explained_var" + suffix: (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
+            }
+            if use_critic
+            else {}
+        ),
+        # response length
+        "response_length/mean" + suffix: torch.mean(response_length).detach().item(),
+        "response_length/max" + suffix: torch.max(response_length).detach().item(),
+        "response_length/min" + suffix: torch.min(response_length).detach().item(),
+        "response_length/clip_ratio"
+        + suffix: torch.mean(torch.eq(response_length, max_response_length).float()).detach().item(),
+        # prompt length
+        "prompt_length/mean" + suffix: torch.mean(prompt_length).detach().item(),
+        "prompt_length/max" + suffix: torch.max(prompt_length).detach().item(),
+        "prompt_length/min" + suffix: torch.min(prompt_length).detach().item(),
+        "prompt_length/clip_ratio"
+        + suffix: torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(),
+    }
+    return metrics
+
+
+class EnvAgentLightningTrainer(RayPPOTrainer):
+    """
+    Specialized PPO trainer for agent-based reinforcement learning.
+
+    This trainer is designed specifically for scenarios where the model interacts with
+    external environments, tools, or APIs through an AgentLightningServer. It simplifies
+    the training loop by removing the complex conditional logic present in the original
+    RayPPOTrainer and focusing on the agent mode workflow.
+
+    Key differences from RayPPOTrainer:
+
+    1. Uses AgentModeDaemon for server communication
+    2. Simplified data flow without pop/union operations
+    3. Direct batch processing through agent daemon
+    4. Streamlined validation using agent_mode validation
+    """
+
+    def __init__(
+        self,
+        store: LightningStore | None,
+        llm_proxy: LLMProxy | None,
+        adapter: TraceAdapter | None,
+        daemon_cls: Type[EnvAgentModeDaemon],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.store = store
+        self.llm_proxy = llm_proxy
+        self.adapter = adapter
+        self.daemon_cls = daemon_cls
+
+    def _validate(self):
+        assert len(self.val_dataloader) == 1, "Please set val_batch_size to None for better throughput."
+
+        test_data = next(iter(self.val_dataloader))
+        test_batch = DataProto.from_single_dict(test_data)
+
+        self.async_rollout_manager.wake_up()
+        self.agent_mode_daemon.set_up_data_and_server(
+            test_batch.non_tensor_batch,
+            self.async_rollout_manager.server_addresses,
+            is_train=False,
+        )
+        self.agent_mode_daemon.run_until_all_finished()
+        test_metrics = self.agent_mode_daemon.get_test_metrics()
+        self.agent_mode_daemon.clear_data_and_server()
+        self.async_rollout_manager.sleep()
+        return test_metrics
+
+    def _compute_reference_log_prob(self, batch: DataProto) -> DataProto:
+        """Compute reference log probability using the correct worker based on LoRA configuration.
+
+        In verl 0.6.0+, when LoRA is detected (indicated by ref_in_actor=True),
+        the reference policy is computed by the actor rollout worker instead of a separate
+        ref policy worker. This method handles both scenarios by checking the ref_in_actor flag.
+        Note: verl sets ref_in_actor=True when it detects LoRA configuration (e.g., lora_rank > 0 or lora_adapter_path is set).
+
+        Args:
+            batch: The data batch to compute reference log probabilities for.
+
+        Returns:
+            DataProto with reference log probabilities added.
+
+        Raises:
+            RuntimeError: If the required worker is not available.
+        """
+        if getattr(self, "ref_in_actor", False):
+            actor_worker = getattr(self, "actor_rollout_wg", None)
+            if actor_worker is None:
+                raise RuntimeError("actor_rollout_wg is required when ref_in_actor is True.")
+            return actor_worker.compute_ref_log_prob(batch)
+
+        ref_worker = getattr(self, "ref_policy_wg", None)
+        if ref_worker is None:
+            raise RuntimeError(
+                "Reference policy worker was not initialized. "
+                "Ensure `use_reference_policy` is enabled and the VERL config exposes the ref worker."
+            )
+        return ref_worker.compute_ref_log_prob(batch)
+
+    def _train_step(self, batch_dict: dict) -> dict:
+        # Isolate in a separate method to automatically recycle the variables before validation.
+        batch: DataProto = DataProto.from_single_dict(batch_dict)
+        metrics = {}
+        timing_raw = {}
+
+        with _timer("step", timing_raw):
+
+            # When agent mode is enabled, we read the batch as it is.
+            gen_batch = batch
+
+            # generate a batch
+            with _timer("gen", timing_raw):
+                self.async_rollout_manager.wake_up()
+                self.agent_mode_daemon.set_up_data_and_server(
+                    gen_batch.non_tensor_batch, self.async_rollout_manager.server_addresses
+                )
+                self.agent_mode_daemon.run_until_all_finished()
+                batch, agent_metrics = self.agent_mode_daemon.get_train_data_batch(
+                    max_prompt_length=self.config.data.max_prompt_length,
+                    max_response_length=self.config.data.max_response_length,
+                    device=gen_batch.batch["fake_ids"].device,
+                    use_final_reward_as_step_reward=self.config.algorithm.use_final_reward_as_step_reward,
+                    use_intrinsic_reward=self.config.algorithm.use_intrinsic_reward,
+                )
+                metrics.update(agent_metrics)
+                self.agent_mode_daemon.clear_data_and_server()
+                self.async_rollout_manager.sleep()
+
+            if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
+                with _timer("gen_max", timing_raw):
+                    gen_baseline_batch = deepcopy(gen_batch)
+                    gen_baseline_batch.meta_info["do_sample"] = False
+                    gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)
+
+                    batch = batch.union(gen_baseline_output)
+                    reward_baseline_tensor = self.reward_fn(batch)
+                    reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
+
+                    batch.pop(batch_keys=list(gen_baseline_output.batch.keys()))
+
+                    batch.batch["reward_baselines"] = reward_baseline_tensor
+
+                    del gen_baseline_batch, gen_baseline_output
+
+            # uid is used for algorithm like GRPO, should be aligned to data id
+            batch.non_tensor_batch["uid"] = batch.non_tensor_batch["data_id_list"]
+
+            batch.batch["response_mask"] = compute_response_mask(batch)
+
+            # compute global_valid tokens
+            batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
+
+            with _timer("reward", timing_raw):
+                # compute reward model score
+                if self.use_rm:
+                    reward_tensor = self.rm_wg.compute_rm_score(batch)
+                    batch = batch.union(reward_tensor)
+
+                reward_extra_infos_dict = {}
+
+            # for agent mode, pad the lengths to calculate old log prob, ref, and values
+            batch, pad_size = pad_dataproto_to_divisor(batch, self.actor_rollout_wg.world_size)
+
+            # recompute old_log_probs
+            with _timer("old_log_prob", timing_raw):
+                old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
+                entropys = old_log_prob.batch["entropys"]
+                response_masks = batch.batch["response_mask"]
+                loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode
+                entropy_loss = agg_loss(loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode)
+                old_log_prob_metrics = {"actor/entropy_loss": entropy_loss.detach().item()}
+                metrics.update(old_log_prob_metrics)
+                old_log_prob.batch.pop("entropys")
+                batch = batch.union(old_log_prob)
+
+            if self.use_reference_policy:
+                # compute reference log_prob
+                with _timer("ref", timing_raw):
+                    ref_log_prob = self._compute_reference_log_prob(batch)
+                    batch = batch.union(ref_log_prob)
+
+            # compute values
+            if self.use_critic:
+                with _timer("values", timing_raw):
+                    values = self.critic_wg.compute_values(batch)
+                    batch = batch.union(values)
+
+            # for agent mode, unpad to calculate adv
+            # it is important, as adv should be based on the raw traces
+            batch = unpad_dataproto(batch, pad_size=pad_size)
+
+            with _timer("adv", timing_raw):
+                # if agent_mode is enabled, there is already token_level_scores
+                # token_level_scores is not needed to compute here
+
+                # compute rewards. apply_kl_penalty if available
+                if self.config.algorithm.use_kl_in_reward:
+                    batch, kl_metrics = apply_kl_penalty(
+                        batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
+                    )
+                    metrics.update(kl_metrics)
+                else:
+                    if self.config.algorithm.use_intrinsic_reward:
+                        batch.batch["token_level_rewards"] = (
+                            batch.batch["token_level_scores"] + batch.batch["token_level_intrinsic_rewards"]
+                        )  # (bs, seq_len)
+                    else:
+                        batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
+
+                # compute advantages, executed on the driver process
+
+                norm_adv_by_std_in_grpo = self.config.algorithm.get(
+                    "norm_adv_by_std_in_grpo", True
+                )  # GRPO adv normalization factor
+
+                batch = compute_advantage(
+                    batch,
+                    adv_estimator=self.config.algorithm.adv_estimator,
+                    gamma=self.config.algorithm.gamma,
+                    lam=self.config.algorithm.lam,
+                    num_repeat=self.config.actor_rollout_ref.rollout.n,
+                    norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+                    config=self.config.algorithm,
+                )
+
+            # Calculate the metrics before processing. Refer to the comments of function `compute_data_metrics` for details.
+            metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, suffix="_before_processing"))
+
+            # after advantages are assinged, we begin to drop (1) long prompt (2) floor to ppo minisize
+            keep_indices = (~batch.batch["is_drop_mask"]).nonzero(as_tuple=True)[0]
+            metrics["training/n_triplets_prompt_too_long"] = (
+                batch.batch["is_drop_mask"].shape[0] - keep_indices.shape[0]
+            )
+            batch = batch[keep_indices]
+            # next, round to minibatch size
+            mini_batch_size = self.config.actor_rollout_ref.actor.ppo_mini_batch_size
+            n_transition = len(batch)
+            random_indices = list(range(n_transition))
+            random.shuffle(random_indices)
+            batch.reorder(torch.tensor(random_indices).type(torch.int32))
+            n_remained_transition = n_transition // mini_batch_size * mini_batch_size
+            batch = batch[list(range(n_remained_transition))]
+            metrics["training/n_triplets_dropped_remainder"] = n_transition - n_remained_transition
+
+            # Agent mode note: Change the order of balance batch;
+            #     1. first calculate advantage
+            #     2. then drop the samples (too long prompt & floor to ppo minisize)
+            #     3. balance
+            # balance the number of valid tokens on each dp rank.
+            # Note that this breaks the order of data inside the batch.
+            # Please take care when you implement group based adv computation such as GRPO and rloo
+            if self.config.trainer.balance_batch:
+                self._balance_batch(batch, metrics=metrics)
+
+            # update critic
+            if self.use_critic:
+                with _timer("update_critic", timing_raw):
+                    critic_output = self.critic_wg.update_critic(batch)
+                critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"])
+                metrics.update(critic_output_metrics)
+
+            # implement critic warmup
+            if self.config.trainer.critic_warmup <= self.global_steps:
+                # update actor
+                with _timer("update_actor", timing_raw):
+                    batch.meta_info["multi_turn"] = self.config.actor_rollout_ref.rollout.multi_turn.enable
+                    actor_output = self.actor_rollout_wg.update_actor(batch)
+                actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])
+                metrics.update(actor_output_metrics)
+
+            # Log rollout generations if enabled
+            rollout_data_dir = self.config.trainer.get("rollout_data_dir", None)
+            if rollout_data_dir:
+                with _timer("dump_rollout_generations", timing_raw):
+                    print(batch.batch.keys())
+                    inputs = self.tokenizer.batch_decode(batch.batch["prompts"], skip_special_tokens=True)
+                    outputs = self.tokenizer.batch_decode(batch.batch["responses"], skip_special_tokens=True)
+                    scores = batch.batch["token_level_scores"].sum(-1).cpu().tolist()
+                    self._dump_generations(
+                        inputs=inputs,
+                        outputs=outputs,
+                        scores=scores,
+                        reward_extra_infos_dict=reward_extra_infos_dict,
+                        dump_path=rollout_data_dir,
+                    )
+
+        # compute training metrics
+        metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, suffix="_after_processing"))
+        metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
+        # TODO: implement actual tflpo and theoretical tflpo
+        n_gpus = self.resource_pool_manager.get_n_gpus()
+        metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))
+
+        return metrics
+
+    def fit(self):
+        logger = Tracking(
+            project_name=self.config.trainer.project_name,
+            experiment_name=self.config.trainer.experiment_name,
+            default_backend=self.config.trainer.logger,
+            config=OmegaConf.to_container(self.config, resolve=True),
+        )
+
+        self.global_steps = 0
+
+        # load checkpoint before doing anything
+        self._load_checkpoint()
+
+        assert self.async_rollout_mode, "If agent mode is enabled, async server must be enabled"
+        if self.adapter is not None and not isinstance(self.adapter, TraceToTripletBase):
+            raise ValueError("Adapter must be a TraceToTripletBase for currently VERL implementation.")
+        verl_version = verl.__version__
+        if verl_version == "0.5.0":
+            # Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
+            # However, it is possible that verl updates the naming and causes incompatibility.
+            # Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
+            model = "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:])
+        else:
+            # For other versions (e.g., 0.6.0), we use the full path to the model.
+            model = self.config.actor_rollout_ref.model.path
+        self.agent_mode_daemon = self.daemon_cls(
+            self.config.agentlightning.port,
+            self.config.actor_rollout_ref.rollout.n,
+            train_information={
+                "model": model,
+                "temperature": self.config.actor_rollout_ref.rollout.temperature,
+            },
+            tokenizer=self.tokenizer,
+            mini_batch_size=self.config.actor_rollout_ref.actor.ppo_mini_batch_size,
+            pad_token_id=self.tokenizer.pad_token_id,
+            mode="v1" if self.store is not None else "v0",
+            store=self.store,
+            llm_proxy=self.llm_proxy,
+            adapter=self.adapter,
+        )
+        self.agent_mode_daemon.start()
+
+        # perform validation before training
+        # currently, we only support validation using the reward_function.
+        if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
+            val_metrics = self._validate()
+            assert val_metrics, f"{val_metrics=}"
+            pprint(f"Initial validation metrics: {val_metrics}")
+            logger.log(data=val_metrics, step=self.global_steps)
+            if self.config.trainer.get("val_only", False):
+                return
+
+        # add tqdm
+        progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress")
+
+        # we start from step 1
+        self.global_steps += 1
+        last_val_metrics = None
+
+        for epoch in range(self.config.trainer.total_epochs):
+            for batch_dict in self.train_dataloader:
+                metrics = {}
+                timing_raw = {}
+                is_last_step = self.global_steps >= self.total_training_steps
+
+                # train step
+                metrics = self._train_step(batch_dict)
+
+                # validate
+                if (
+                    self.val_reward_fn is not None
+                    and self.config.trainer.test_freq > 0
+                    and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
+                ):
+                    with _timer("validate", timing_raw):
+                        val_metrics: dict = self._validate()
+                        if is_last_step:
+                            last_val_metrics = val_metrics
+                    metrics.update(val_metrics)
+
+                if self.config.trainer.save_freq > 0 and (
+                    is_last_step or self.global_steps % self.config.trainer.save_freq == 0
+                ):
+                    with _timer("save_checkpoint", timing_raw):
+                        self._save_checkpoint()
+
+                # step metrics
+                metrics.update(
+                    {
+                        "training/global_step": self.global_steps,
+                        "training/epoch": epoch,
+                    }
+                )
+
+                # TODO: make a canonical logger that supports various backend
+                logger.log(data=metrics, step=self.global_steps)
+
+                if is_last_step:
+                    pprint(f"Final validation metrics: {last_val_metrics}")
+                    progress_bar.close()
+
+                    # This exit logic is to ensure a robust CI.
+                    pprint(f"Flush the logger...")
+                    del logger  # Make sure the loggers are flushed and closed properly
+                    pprint(f"Training finished at step {self.global_steps}.")
+                    return
+
+                progress_bar.update(1)
+                self.global_steps += 1

contrib/agentlightning/contrib/reward/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Namespace package for agentlightning.contrib.reward.

contrib/agentlightning/contrib/reward/agentos.py

@@ -0,0 +1,127 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+PolicyReward - Convert Policy Violations to RL Penalties
+=========================================================
+
+Reward function that integrates Agent-OS governance.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PolicyReward:
+    """
+    Reward function that penalizes policy violations.
+
+    Example:
+        >>> from agent_os import KernelSpace
+        >>>
+        >>> kernel = KernelSpace(policy="strict")
+        >>> reward_fn = PolicyReward(kernel, base_reward_fn=accuracy)
+        >>>
+        >>> reward = reward_fn(rollout)  # Base reward - violation penalties
+    """
+
+    def __init__(
+        self,
+        kernel: Any,
+        *,
+        base_reward_fn: Optional[Callable[[Any], float]] = None,
+        critical_penalty: float = -100.0,
+        high_penalty: float = -50.0,
+        medium_penalty: float = -10.0,
+        low_penalty: float = -1.0,
+        clean_bonus: float = 5.0,
+    ):
+        """
+        Initialize policy-aware reward.
+
+        Args:
+            kernel: Agent-OS KernelSpace
+            base_reward_fn: Base reward function
+            critical_penalty: Penalty for critical violations
+            high_penalty: Penalty for high violations
+            medium_penalty: Penalty for medium violations
+            low_penalty: Penalty for low violations
+            clean_bonus: Bonus for clean execution
+        """
+        self.kernel = kernel
+        self.base_reward_fn = base_reward_fn or self._default_reward
+        self.penalties = {
+            "critical": critical_penalty,
+            "high": high_penalty,
+            "medium": medium_penalty,
+            "low": low_penalty,
+        }
+        self.clean_bonus = clean_bonus
+
+        self._total_rewards = 0
+        self._total_penalties = 0.0
+
+    def _default_reward(self, rollout: Any) -> float:
+        """Default: 1.0 for success, 0.0 for failure."""
+        return 1.0 if getattr(rollout, "success", False) else 0.0
+
+    def __call__(self, rollout: Any, *, emit: bool = True) -> float:
+        """
+        Calculate reward with policy penalties.
+
+        Args:
+            rollout: Rollout with violations attribute
+            emit: Emit reward span
+
+        Returns:
+            Final reward
+        """
+        base = self.base_reward_fn(rollout)
+
+        violations = getattr(rollout, "violations", [])
+        penalty = sum(self.penalties.get(v.severity, -10.0) for v in violations)
+
+        reward = base + penalty
+        if not violations:
+            reward += self.clean_bonus
+
+        self._total_rewards += 1
+        self._total_penalties += penalty
+
+        if emit:
+            self._emit_reward(reward, base, penalty, len(violations))
+
+        return reward
+
+    def _emit_reward(
+        self,
+        final: float,
+        base: float,
+        penalty: float,
+        violation_count: int,
+    ) -> None:
+        """Emit multi-dimensional reward."""
+        try:
+            from agentlightning.emitter import emit_reward
+
+            emit_reward(
+                {"final": final, "base": base, "policy_penalty": penalty},
+                primary_key="final",
+                attributes={"agent_os.violations": violation_count},
+            )
+        except ImportError:
+            logger.debug(
+                "agentlightning.emitter not available; skipping reward emission.",
+                exc_info=True,
+            )
+
+    def get_stats(self) -> Dict[str, float]:
+        """Get reward statistics."""
+        total = self._total_rewards or 1
+        return {
+            "total_rewards": self._total_rewards,
+            "avg_penalty": self._total_penalties / total,
+        }

contrib/agentlightning/contrib/runner/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Namespace package for agentlightning.contrib.runner.

contrib/agentlightning/contrib/runner/agentos.py

@@ -0,0 +1,282 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+AgentOSRunner - Agent-Lightning Runner with Kernel Safety
+==========================================================
+
+Wraps agent execution with Agent-OS kernel governance.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Generic, Optional, TypeVar
+
+logger = logging.getLogger(__name__)
+
+T_task = TypeVar("T_task")
+
+
+@dataclass
+class PolicyViolation:
+    """Record of a policy violation."""
+
+    policy_name: str
+    description: str
+    severity: str
+    blocked: bool
+    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    @property
+    def penalty(self) -> float:
+        """Calculate penalty based on severity.
+
+        Returns:
+            float: Negative penalty value, where more severe violations
+                have larger negative magnitudes.
+        """
+        penalties = {
+            "critical": -100.0,
+            "high": -50.0,
+            "medium": -10.0,
+            "low": -1.0,
+        }
+        return penalties.get(self.severity, -10.0)
+
+
+@dataclass
+class GovernedRollout:
+    """Rollout with governance metadata.
+
+    This dataclass wraps execution results with governance information.
+    It is compatible with Agent-Lightning's Rollout interface - the
+    `task_input`, `task_output`, and `success` fields provide the core
+    rollout data, while `violations` adds governance-specific metadata.
+    """
+
+    task_input: Any
+    task_output: Any
+    success: bool
+    violations: list[PolicyViolation] = field(default_factory=list)
+
+    @property
+    def total_penalty(self) -> float:
+        return sum(v.penalty for v in self.violations)
+
+
+class AgentOSRunner(Generic[T_task]):
+    """
+    Agent-Lightning runner with Agent-OS kernel safety.
+
+    This runner wraps agent execution in an Agent-OS kernel,
+    enforcing policies and collecting violation data for RL training.
+
+    Example:
+        >>> from agent_os import KernelSpace
+        >>> from agent_os.policies import SQLPolicy
+        >>>
+        >>> kernel = KernelSpace(policy=SQLPolicy())
+        >>> runner = AgentOSRunner(kernel)
+        >>>
+        >>> rollout = await runner.step(task)
+        >>> print(f"Violations: {len(rollout.violations)}")
+    """
+
+    def __init__(
+        self,
+        kernel: Any,
+        *,
+        fail_on_violation: bool = False,
+        emit_violations: bool = True,
+    ):
+        """
+        Initialize the governed runner.
+
+        Args:
+            kernel: Agent-OS KernelSpace with loaded policies
+            fail_on_violation: Raise exception on violation
+            emit_violations: Emit violations as spans
+        """
+        self.kernel = kernel
+        self.fail_on_violation = fail_on_violation
+        self.emit_violations = emit_violations
+
+        self._violations: list[PolicyViolation] = []
+        self._total_rollouts = 0
+        self._total_violations = 0
+
+        # Worker attributes (set by init_worker)
+        self.worker_id: Optional[int] = None
+        self.store: Optional[Any] = None
+
+        self._setup_hooks()
+
+    def _setup_hooks(self) -> None:
+        """Set up kernel hooks."""
+        on_violation = getattr(self.kernel, "on_policy_violation", None)
+        if on_violation is None:
+            logger.warning(
+                "Kernel %r does not support policy violation hooks via 'on_policy_violation'.",
+                self.kernel,
+            )
+            return
+        if not callable(on_violation):
+            logger.warning(
+                "Kernel attribute 'on_policy_violation' is not callable: %r",
+                on_violation,
+            )
+            return
+        try:
+            on_violation(self._handle_violation)
+        except TypeError as exc:
+            logger.warning(
+                "Kernel.on_policy_violation has an incompatible signature: %s",
+                exc,
+            )
+
+    def _handle_violation(
+        self,
+        policy_name: str,
+        description: str,
+        severity: str,
+        blocked: bool,
+    ) -> None:
+        """Handle a policy violation."""
+        violation = PolicyViolation(
+            policy_name=policy_name,
+            description=description,
+            severity=severity,
+            blocked=blocked,
+        )
+        self._violations.append(violation)
+        self._total_violations += 1
+
+        if self.emit_violations:
+            self._emit_violation_span(violation)
+
+        if self.fail_on_violation and blocked:
+            raise PolicyViolationError(violation)
+
+    def _emit_violation_span(self, violation: PolicyViolation) -> None:
+        """Emit violation as Agent-Lightning span."""
+        try:
+            from agentlightning.emitter import emit_annotation
+
+            emit_annotation(
+                {
+                    "agent_os.violation": True,
+                    "agent_os.policy": violation.policy_name,
+                    "agent_os.severity": violation.severity,
+                    "agent_os.blocked": violation.blocked,
+                }
+            )
+        except ImportError as exc:
+            logger.debug(
+                "agentlightning.emitter not available; skipping violation annotation: %s",
+                exc,
+            )
+
+    @property
+    def agent(self) -> Any:
+        """
+        Access the underlying agent.
+
+        Raises:
+            RuntimeError: If the agent has not been initialized via `init`.
+        """
+        if not hasattr(self, "_agent"):
+            raise RuntimeError("AgentOSRunner.agent accessed before `init` has been called.")
+        return self._agent
+
+    @agent.setter
+    def agent(self, value: Any) -> None:
+        """Set the underlying agent instance."""
+        self._agent = value
+
+    def init(self, agent: Any, **kwargs: Any) -> None:
+        """Initialize with agent."""
+        self.agent = agent
+
+    def init_worker(self, worker_id: int, store: Any, **kwargs: Any) -> None:
+        """Initialize worker."""
+        self.worker_id = worker_id
+        self.store = store
+
+    def teardown(self) -> None:
+        """Release resources."""
+        pass
+
+    def teardown_worker(self, worker_id: int) -> None:
+        """Release worker resources."""
+        pass
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[Any] = None,
+        mode: Optional[str] = None,
+        event: Optional[Any] = None,
+    ) -> GovernedRollout:
+        """
+        Execute task with governance.
+
+        Args:
+            input: Task input
+            resources: Optional resources
+            mode: Rollout mode
+            event: Stop signal
+
+        Returns:
+            GovernedRollout with results and violations
+        """
+        self._violations = []
+
+        try:
+            if hasattr(self.kernel, "execute_async"):
+                logger.debug("AgentOSRunner: executing task via kernel.execute_async")
+                result = await self.kernel.execute_async(self.agent, input)
+            elif hasattr(self.kernel, "execute"):
+                logger.debug("AgentOSRunner: executing task via kernel.execute")
+                result = self.kernel.execute(self.agent, input)
+            else:
+                logger.error(
+                    "AgentOSRunner: kernel does not support 'execute_async' or 'execute'; "
+                    "governed execution is not possible."
+                )
+                raise RuntimeError(
+                    "Kernel does not support governed execution (missing 'execute_async' and 'execute')."
+                )
+            success = True
+        except PolicyViolationError as e:
+            # Record the policy violation and mark rollout as unsuccessful.
+            self._violations.append(e.violation)
+            result = None
+            success = False
+
+        self._total_rollouts += 1
+
+        return GovernedRollout(
+            task_input=input,
+            task_output=result,
+            success=success,
+            violations=self._violations.copy(),
+        )
+
+    def get_stats(self) -> dict:
+        """Get runner statistics."""
+        return {
+            "total_rollouts": self._total_rollouts,
+            "total_violations": self._total_violations,
+            "violation_rate": (self._total_violations / self._total_rollouts if self._total_rollouts > 0 else 0.0),
+        }
+
+
+class PolicyViolationError(Exception):
+    """Raised when policy violation blocks execution."""
+
+    def __init__(self, violation: PolicyViolation):
+        self.violation = violation
+        super().__init__(f"Policy violation: {violation.description}")

contrib/recipes/agentos/README.md

@@ -0,0 +1,105 @@
+# Agent-OS Integration for Agent-Lightning
+
+Kernel-level safety during AI agent training.
+
+## Overview
+
+[Agent-OS](https://github.com/imran-siddique/agent-os) provides deterministic governance
+for AI agents. This integration enables:
+
+- **0% unpenalized policy violations** — All unsafe actions are detected and penalized
+- **Policy violations → RL penalties** — Agents learn to avoid unsafe behavior
+- **Complete audit trail** — From training to production
+
+## Installation
+
+```bash
+pip install agentlightning agent-os
+```
+
+## Quick Start
+
+```python
+from agentlightning import Trainer
+from agentlightning.contrib.runner.agentos import AgentOSRunner
+from agentlightning.contrib.reward.agentos import PolicyReward
+from agent_os import KernelSpace
+from agent_os.policies import SQLPolicy
+
+# Create governed kernel
+kernel = KernelSpace(policy=SQLPolicy(
+    deny=["DROP", "DELETE"]
+))
+
+# Wrap in Agent-OS runner
+runner = AgentOSRunner(kernel)
+
+# Train with policy-aware rewards
+trainer = Trainer(
+    runner=runner,
+    reward_fn=PolicyReward(kernel),
+    algorithm="GRPO"
+)
+
+trainer.train()
+```
+
+## Components
+
+### AgentOSRunner
+
+Wraps agent execution with kernel-level policy enforcement:
+
+```python
+from agentlightning.contrib.runner.agentos import AgentOSRunner
+
+runner = AgentOSRunner(
+    kernel,
+    fail_on_violation=False,  # Continue but penalize
+    emit_violations=True,     # Emit as spans
+)
+```
+
+### PolicyReward
+
+Converts policy violations to negative RL rewards:
+
+```python
+from agentlightning.contrib.reward.agentos import PolicyReward
+
+reward_fn = PolicyReward(
+    kernel,
+    base_reward_fn=accuracy_reward,
+    critical_penalty=-100.0,
+    clean_bonus=5.0,
+)
+```
+
+### FlightRecorderAdapter
+
+Imports Agent-OS audit logs to LightningStore:
+
+```python
+from agentlightning.contrib.adapter.agentos import FlightRecorderAdapter
+
+adapter = FlightRecorderAdapter(flight_recorder)
+adapter.import_to_store(lightning_store)
+```
+
+## Benchmarks
+
+| Metric | Without Agent-OS | With Agent-OS |
+|--------|------------------|---------------|
+| Undetected Policy Violations | 12.3% | **0.0%** |
+| Task Accuracy | 76.4% | **79.2%** |
+
+*Note: "0% undetected violations" means all policy violations are caught and penalized, not that agents never attempt unsafe actions. Over training, agents learn to minimize violation attempts.*
+
+## Documentation
+
+- [Agent-OS Documentation](https://imran-siddique.github.io/agent-os-docs/)
+- Integration guide: see project README or examples in this directory.
+
+## License
+
+MIT

contrib/recipes/agentos/demo_governed_training.py

@@ -0,0 +1,196 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Agent-OS + Agent-Lightning End-to-End Demo
+==========================================
+
+Demonstrates how to train an AI agent with kernel-level safety governance
+using Agent-OS policy enforcement and Agent-Lightning RL training.
+
+This script shows the full pipeline:
+1. Define a governance policy (block dangerous SQL operations)
+2. Create a governed runner that enforces the policy
+3. Define a reward function that penalizes policy violations
+4. Train the agent — it learns to avoid unsafe actions over time
+
+Usage:
+    python demo_governed_training.py
+
+Requirements:
+    pip install agentlightning agent-os-kernel
+"""
+
+from __future__ import annotations
+
+import logging
+
+logging.basicConfig(level=logging.INFO, format="%(name)s | %(message)s")
+logger = logging.getLogger("agentos-demo")
+
+
+def build_kernel():
+    """Create an Agent-OS kernel with SQL safety policy."""
+    try:
+        from agent_os import KernelSpace
+    except ImportError:
+        logger.warning("agent-os-kernel not installed – using stub kernel")
+        return _StubKernel()
+
+    kernel = KernelSpace(
+        policies={
+            "sql_safety": {
+                "deny_patterns": ["DROP", "DELETE", "TRUNCATE", "ALTER"],
+                "max_rows": 1000,
+            }
+        }
+    )
+    logger.info("Agent-OS kernel created with SQL safety policy")
+    return kernel
+
+
+def build_runner(kernel):
+    """Wrap a runner with Agent-OS policy enforcement."""
+    try:
+        from agentlightning.contrib.runner.agentos import AgentOSRunner
+
+        runner = AgentOSRunner(
+            kernel,
+            fail_on_violation=False,  # Continue but penalize
+            emit_violations=True,  # Emit as spans for observability
+        )
+        logger.info("AgentOSRunner created (fail_on_violation=False)")
+        return runner
+    except ImportError:
+        logger.warning("AgentOSRunner not available – using stub")
+        return None
+
+
+def build_reward(kernel, base_reward_fn=None):
+    """Create a reward function that penalizes policy violations."""
+    try:
+        from agentlightning.contrib.reward.agentos import PolicyReward
+
+        reward = PolicyReward(
+            kernel,
+            base_reward_fn=base_reward_fn,
+            critical_penalty=-100.0,
+            high_penalty=-50.0,
+            medium_penalty=-10.0,
+            low_penalty=-1.0,
+            clean_bonus=5.0,
+        )
+        logger.info("PolicyReward created (critical=-100, clean_bonus=+5)")
+        return reward
+    except ImportError:
+        logger.warning("PolicyReward not available – using stub")
+        return base_reward_fn
+
+
+def accuracy_reward(completions: list[str], references: list[str]) -> list[float]:
+    """Simple accuracy reward: 1.0 for exact match, 0.0 otherwise."""
+    return [1.0 if pred.strip().lower() == ref.strip().lower() else 0.0 for pred, ref in zip(completions, references)]
+
+
+def demo_violation_detection(kernel):
+    """Show that the kernel catches unsafe SQL operations."""
+    logger.info("--- Violation Detection Demo ---")
+
+    test_queries = [
+        ("SELECT * FROM users WHERE id = 1", True),  # Safe
+        ("DROP TABLE users", False),  # Blocked
+        ("DELETE FROM orders WHERE id > 0", False),  # Blocked
+        ("SELECT name FROM products LIMIT 10", True),  # Safe
+        ("TRUNCATE TABLE logs", False),  # Blocked
+    ]
+
+    for query, should_pass in test_queries:
+        try:
+            result = kernel.check_action("sql_query", {"query": query})
+            status = "✅ ALLOWED" if result.allowed else "🛑 BLOCKED"
+        except Exception:
+            # Stub kernel always allows
+            status = "✅ ALLOWED (stub)"
+            result = type("R", (), {"allowed": True})()
+
+        expected = "should pass" if should_pass else "should block"
+        logger.info(f"  {status}: {query!r}  ({expected})")
+
+
+def demo_training_loop(runner, reward_fn):
+    """Simulate a training loop with governed execution."""
+    logger.info("--- Training Loop Demo ---")
+
+    # Simulated training data
+    prompts = [
+        "Find all active users",
+        "Remove all test data",
+        "Show revenue by month",
+        "Drop the staging table",
+    ]
+
+    for epoch in range(1, 3):
+        logger.info(f"Epoch {epoch}/2")
+        for prompt in prompts:
+            logger.info(f"  Prompt: {prompt!r}")
+            # In real training, the runner would execute the agent
+            # and the reward function would score the output
+            logger.info("    → Agent generates SQL → Runner enforces policy → Reward scored")
+
+    logger.info("Training complete — agent learns to avoid policy violations over time")
+
+
+def demo_audit_trail(kernel):
+    """Show the audit trail captured by Agent-OS."""
+    logger.info("--- Audit Trail Demo ---")
+
+    try:
+        from agentlightning.contrib.adapter.agentos import FlightRecorderAdapter
+
+        recorder = getattr(kernel, "flight_recorder", None)
+        if recorder:
+            adapter = FlightRecorderAdapter(recorder)
+            logger.info(f"Audit entries: {len(adapter.get_entries())}")
+        else:
+            logger.info("Flight recorder not available (stub kernel)")
+    except ImportError:
+        logger.info("FlightRecorderAdapter not available — skipping audit demo")
+
+
+def main():
+    """Run the full end-to-end demo."""
+    logger.info("=" * 60)
+    logger.info("Agent-OS + Agent-Lightning End-to-End Demo")
+    logger.info("=" * 60)
+
+    # 1. Build the governance kernel
+    kernel = build_kernel()
+
+    # 2. Demonstrate violation detection
+    demo_violation_detection(kernel)
+
+    # 3. Build governed runner + reward
+    runner = build_runner(kernel)
+    reward_fn = build_reward(kernel, base_reward_fn=accuracy_reward)
+
+    # 4. Simulate training loop
+    demo_training_loop(runner, reward_fn)
+
+    # 5. Show audit trail
+    demo_audit_trail(kernel)
+
+    logger.info("=" * 60)
+    logger.info("Demo complete! In production, replace stubs with:")
+    logger.info("  pip install agentlightning agent-os-kernel")
+    logger.info("=" * 60)
+
+
+class _StubKernel:
+    """Minimal stub for demo when agent-os is not installed."""
+
+    def check_action(self, action_type, params):
+        blocked = any(kw in params.get("query", "").upper() for kw in ("DROP", "DELETE", "TRUNCATE", "ALTER"))
+        return type("Result", (), {"allowed": not blocked})()
+
+
+if __name__ == "__main__":
+    main()

contrib/recipes/envs/README.md

@@ -0,0 +1,94 @@
+# Example of AGL Environments
+
+## Overview
+
+This example implements agents across various environments within Agent Lightning.
+The example is designed to run on a single node with 8 GPUs, each having at least 40 GB of memory.
+
+This example depends on the simulation environments (e.g., ALFWorld, ScienceWorld) provided in the [agl-envs repository](https://github.com/agent-lightning/agl-envs).
+For more information about the supported environments, please refer to the [envs README](https://github.com/agent-lightning/agl-envs/README.md).
+
+---
+
+ 
+
+## Included Files
+
+| File/Directory | Description |
+|----------------|-------------|
+| `config_env/` | Configuration for environment settings. For more information, please refer to the "Configure Your Environment Settings" section. |
+| `config_verl/` | Configuration for RL training with VerL |
+| `add_instruction.py` | Adding instructions to the agent’s input prompt to guide the format of the response |
+| `prompt_builder.py` | Managing conversation history and generating input prompts in multi-turn scenarios |
+| `train_env_agent.py` | RL training script |
+
+---
+
+ 
+
+## Install Environments
+
+Run the following script once to install the enviornment and related AGL dependency:
+
+```bash
+cd contrib/recipes/envs
+
+git clone https://github.com/agent-lightning/agl-envs
+mv agl-envs agl_envs
+
+# Install alfworld dependency
+bash agl_envs/setup/setup_alfworld.sh
+conda activate alfworld
+
+# Install scienceworld dependency
+bash agl_envs/setup/setup_sciworld.sh
+conda activate sciworld
+
+# Install AGL dependency
+bash install_agl.sh
+```
+
+> If you plan to use WandB for experiment tracking, log in to WandB before training.
+
+---
+
+ 
+
+## Configure Your Environment Settings
+### Captioner type (cot or naive)
+- cot: guide the agent to output its reasoning first, then take an action
+- naive: guide the agent to take an action directly, without outputting any reasoning
+
+### Prompt type (chat or single)
+
+When performing multi-turn rollouts, the unit of the input prompt can be defined in two different ways.
+
+(1) **Trajectory-wise unit**:
+All interaction history up to the current step is accumulated in a chat format and directly used to construct the next input prompt.
+
+(2) **Turn-wise unit**:
+Only a subset of the interaction history is included for each turn. The prompt is reconstructed by combining the current turn’s state with selected past information, rather than using the full trajectory.
+
+![prompt_type](./assets/prompt_type.png)
+
+
+You can use the `trajectory-wise unit` by setting `prompt_type` to `chat`, and the `turn-wise unit` by setting `prompt_type` to `single`. Currently, for ALFWorld, we only support the `single` mode, while for ScienceWorld, both `chat` and `single` modes are supported.
+
+We follow the single-mode prompt for ALFWorld from [verl-agent](https://github.com/langfengQ/verl-agent) and the single-mode prompt for ScienceWorld from [RLVMR](https://github.com/Tencent/digitalhuman/tree/main/RLVMR). Thank you to the authors of VERL-Agent and RLVMR for their valuable work.
+
+---
+
+ 
+
+## Run RL Training (GRPO)
+
+```bash
+# Run alfworld
+python3 train_env_agent.py --algorithm grpo --env alfworld
+
+# Run scienceworld single task task_num 0
+python3 train_env_agent.py --algorithm grpo --env scienceworld --task_num 0
+
+# Run scienceworld multi-task
+python3 train_env_agent.py --algorithm grpo --env scienceworld --task_num -1
+```

contrib/recipes/envs/add_instruction.py

@@ -0,0 +1,101 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import copy
+
+from autogen_core.models import UserMessage
+
+# Instruction text definitions
+COT_INSTRUCTION = """
+Now it's your turn to take an action. You should first reason step-by-step about the current situation. This reasoning process MUST be enclosed within <think> </think> tags.
+Once you've finished your reasoning, you should choose an appropriate action for the current step and present it within <action> </action> tags.
+""".strip()
+
+NAIVE_INSTRUCTION = """
+Please response with only one line with one sentence, following the possible action format shown above. No extra words are allowed.
+""".strip()
+
+# Mapping for instruction text types
+INSTRUCTION_MAP = {
+    "cot": COT_INSTRUCTION,
+    "naive": NAIVE_INSTRUCTION,
+}
+
+
+def _get_instruction(type: str, env_name: str = None):
+    """
+    Retrieve an instruction string from INSTRUCTION_MAP based on the given type.
+
+    Args:
+        type (str): Instruction type key (e.g., "cot", "naive", "critic", "tip").
+        env_name (str, optional): Currently unused. Reserved for future
+            environment-specific instruction handling.
+
+    Returns:
+        str: The corresponding instruction text.
+
+    Raises:
+        ValueError: If the given instruction type is not found in INSTRUCTION_MAP.
+    """
+    if type in INSTRUCTION_MAP:
+        return INSTRUCTION_MAP[type]
+    else:
+        raise ValueError(f"Unknown instruction type: {type}")
+
+
+def add_chat_instruction(prompt, type: str, sep: str = "\n\n", env_name: str = None):
+    """
+    Append an instruction to the content of the last message in a chat-style prompt.
+
+    This function does not modify the original prompt. Instead, it returns a
+    deep-copied prompt list with the instruction appended.
+
+    Args:
+        prompt (list): A conversation history represented as a list of objects.
+            Each object must have a `.content` attribute.
+        type (str): Instruction type key (e.g., "cot", "naive", "critic", "tip").
+        sep (str, optional): Separator inserted between the existing content
+            and the instruction.
+        env_name (str, optional): Currently unused. Reserved for future use.
+
+    Returns:
+        list: A new prompt list with the instruction appended to the last message.
+    """
+    new_prompt = copy.deepcopy(prompt)
+    instruction = _get_instruction(type, env_name)
+    new_prompt[-1].content += sep + instruction
+
+    return new_prompt
+
+
+def add_single_instruction(prompt, type: str, sep: str = "\n\n", env_name: str = None):
+    """
+    Append an instruction to a single prompt or a chat-style prompt.
+
+    - If `prompt` is a string, the instruction is appended to the string.
+    - If `prompt` is a list, the instruction is appended to the `.content`
+      of the last message.
+
+    Args:
+        prompt (str or list): Either a single prompt string or a conversation
+            history list whose elements have a `.content` attribute.
+        type (str): Instruction type key (e.g., "cot", "naive", "critic", "tip").
+        sep (str, optional): Separator inserted between the existing content
+            and the instruction.
+        env_name (str, optional): Currently unused. Reserved for future use.
+
+    Returns:
+        str or list: The updated prompt with the instruction appended.
+
+    Raises:
+        TypeError: If `prompt` is neither a string nor a list.
+    """
+    instruction = _get_instruction(type, env_name)
+
+    if isinstance(prompt, str):
+        return prompt + sep + instruction
+    elif isinstance(prompt, list):
+        new_prompt = copy.deepcopy(prompt)
+        new_prompt[-1].content += sep + instruction
+        return new_prompt
+    else:
+        raise TypeError("Prompt must be a string or a list of strings")

contrib/recipes/envs/config_env/alfworld.yaml

@@ -0,0 +1,16 @@
+env_name: alfworld
+seed: 0
+format_penalty: 0.1
+binary_reward: False
+save_rollout: False
+log_env_obs: False
+reawrd_scale: 10.0
+use_success_rate: False
+
+captioner:
+  type: cot # naive or cot
+  prompt_type: single # chat or single
+  max_history: 2
+
+alfworld_kwargs:
+  max_steps: 50

contrib/recipes/envs/config_env/scienceworld.yaml

@@ -0,0 +1,16 @@
+env_name: scienceworld
+seed: 0
+format_penalty: 0.1
+binary_reward: False
+save_rollout: False
+log_env_obs: False # True for GiGPO
+reawrd_scale: 10.0
+use_success_rate: True
+
+# only for scienceworld
+use_action_correction: False
+
+captioner:
+  type: cot # naive or cot
+  prompt_type: single # chat or single
+  max_history: 2

contrib/recipes/envs/config_verl/alfworld/grpo.yaml

@@ -0,0 +1,86 @@
+# ==========================
+# Variable definitions
+# ==========================
+variables:
+  NUM_GPUS: 2
+  MINI_BATCH_SIZE: 32
+  PER_GPU_BATCH_SIZE: 16
+  TENSOR_MODEL_PARALLEL_SIZE: 2
+  NUM_ROLLOUTS: 8
+  BASE_MODEL: Qwen/Qwen2.5-1.5B-Instruct
+  PROJECT_NAME: AGL-Envs-ALFWorld
+  TRIAL: ${oc.env:TRIAL,0}
+  EXPERIMENT_NAME: grpo-alfworld-${variables.TRIAL}
+  DATA_DIR: agl_envs/task_data/alfworld
+
+# ==========================
+# Main Config
+# ==========================
+agentlightning:
+  port: 9999
+
+algorithm:
+  adv_estimator: grpo
+  use_kl_in_reward: false
+  use_final_reward_as_step_reward: true
+  use_intrinsic_reward: true
+
+data:
+  train_files: ${variables.DATA_DIR}/train.parquet
+  val_files: ${variables.DATA_DIR}/test.parquet
+  train_batch_size: 32
+  val_batch_size: 140
+  max_prompt_length: 2048
+  max_response_length: 512
+  truncation: error
+  return_raw_chat: true
+
+actor_rollout_ref:
+  rollout:
+    tensor_model_parallel_size: ${variables.TENSOR_MODEL_PARALLEL_SIZE}
+    n: ${variables.NUM_ROLLOUTS}
+    log_prob_micro_batch_size_per_gpu: ${variables.PER_GPU_BATCH_SIZE}
+    multi_turn:
+      format: hermes
+    name: vllm
+    gpu_memory_utilization: 0.6
+    enable_chunked_prefill: false
+    enforce_eager: false
+    free_cache_engine: true
+    val_kwargs:
+      temperature: 0.4
+      do_sample: true
+  actor:
+    ppo_mini_batch_size: ${variables.MINI_BATCH_SIZE}
+    ppo_micro_batch_size_per_gpu: ${variables.PER_GPU_BATCH_SIZE}
+    optim:
+      lr: 1.0e-6
+    use_kl_loss: true
+    kl_loss_coef: 0.01
+    kl_loss_type: low_var_kl
+    entropy_coeff: 0.001
+    fsdp_config:
+      param_offload: false
+      optimizer_offload: false
+  ref:
+    log_prob_micro_batch_size_per_gpu: ${variables.PER_GPU_BATCH_SIZE}
+    fsdp_config:
+      param_offload: true
+  model:
+    path: ${variables.BASE_MODEL}
+    use_remove_padding: true
+    enable_gradient_checkpointing: true
+
+trainer:
+  n_gpus_per_node: ${variables.NUM_GPUS}
+  val_before_train: false
+  critic_warmup: 0
+  logger:
+    - console
+    - wandb
+  project_name: ${variables.PROJECT_NAME}
+  experiment_name: ${variables.EXPERIMENT_NAME}
+  nnodes: 1
+  save_freq: 100
+  test_freq: 5
+  total_epochs: 200

contrib/recipes/envs/config_verl/scienceworld/grpo.yaml

@@ -0,0 +1,87 @@
+# ==========================
+# Variable definitions
+# ==========================
+variables:
+  NUM_GPUS: 2
+  MINI_BATCH_SIZE: 32
+  PER_GPU_BATCH_SIZE: 16
+  TENSOR_MODEL_PARALLEL_SIZE: 2
+  NUM_ROLLOUTS: 8
+  BASE_MODEL: Qwen/Qwen2.5-1.5B-Instruct
+  PROJECT_NAME: AGL-Envs-ScienceWorld
+  TASK_NUM: ${oc.env:TASK_NUM,-1}
+  TRIAL: ${oc.env:TRIAL,0}
+  EXPERIMENT_NAME: grpo-sciworld-${variables.TRIAL}
+  DATA_DIR: agl_envs/task_data/scienceworld/multi_data
+
+# ==========================
+# Main Config
+# ==========================
+agentlightning:
+  port: 9999
+
+algorithm:
+  adv_estimator: grpo
+  use_kl_in_reward: false
+  use_final_reward_as_step_reward: true
+  use_intrinsic_reward: true
+
+data:
+  train_files: ${variables.DATA_DIR}/train.parquet
+  val_files: ${variables.DATA_DIR}/test.parquet
+  train_batch_size: 32
+  val_batch_size: 144
+  max_prompt_length: 6000
+  max_response_length: 1024
+  truncation: error
+  return_raw_chat: true
+
+actor_rollout_ref:
+  rollout:
+    tensor_model_parallel_size: ${variables.TENSOR_MODEL_PARALLEL_SIZE}
+    n: ${variables.NUM_ROLLOUTS}
+    log_prob_micro_batch_size_per_gpu: ${variables.PER_GPU_BATCH_SIZE}
+    multi_turn:
+      format: hermes
+    name: vllm
+    gpu_memory_utilization: 0.6
+    enable_chunked_prefill: false
+    enforce_eager: false
+    free_cache_engine: true
+    val_kwargs:
+      temperature: 0.4
+      do_sample: true
+  actor:
+    ppo_mini_batch_size: ${variables.MINI_BATCH_SIZE}
+    ppo_micro_batch_size_per_gpu: ${variables.PER_GPU_BATCH_SIZE}
+    optim:
+      lr: 1.0e-6
+    use_kl_loss: true
+    kl_loss_coef: 0.01
+    kl_loss_type: low_var_kl
+    entropy_coeff: 0.001
+    fsdp_config:
+      param_offload: false
+      optimizer_offload: false
+  ref:
+    log_prob_micro_batch_size_per_gpu: ${variables.PER_GPU_BATCH_SIZE}
+    fsdp_config:
+      param_offload: true
+  model:
+    path: ${variables.BASE_MODEL}
+    use_remove_padding: true
+    enable_gradient_checkpointing: true
+
+trainer:
+  n_gpus_per_node: ${variables.NUM_GPUS}
+  val_before_train: true
+  critic_warmup: 0
+  logger:
+    - console
+    - wandb
+  project_name: ${variables.PROJECT_NAME}
+  experiment_name: ${variables.EXPERIMENT_NAME}
+  nnodes: 1
+  save_freq: 100
+  test_freq: 5
+  total_epochs: 500

contrib/recipes/envs/install_agl.sh

@@ -0,0 +1,12 @@
+# This setup is based on CUDA 12.6
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126
+pip install transformers==4.56.1
+pip install wandb
+pip install vllm==0.10.2
+pip install verl==0.5.0
+pip install click==8.2.1
+pip install --extra-index-url https://miropsota.github.io/torch_packages_builder flash_attn==2.8.3+pt2.8.0cu126
+pip install 'openai-agents[litellm]'==0.2.9
+pip install -U "autogen-agentchat" "autogen-ext[openai]"
+
+(cd ../../../ && pip install -e .[dev])

contrib/recipes/envs/prompt_builder.py

@@ -0,0 +1,183 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Optional
+
+from autogen_core.models import AssistantMessage, UserMessage
+
+
+class HistoryPromptBuilder:
+    """
+    Builds prompts using a history of observations and actions.
+
+    Supports two prompt styles:
+    - chat: multi-turn user/assistant messages
+    - single: a single formatted prompt with optional history
+    """
+
+    def __init__(self, max_history: int = -1, prompt_type: str = "chat"):
+        """
+        Args:
+            max_history (int): Maximum number of past steps to include
+                               (-1 means unlimited).
+            prompt_type (str): Prompt style ("chat" or "single").
+        """
+        self.max_history = max_history
+        self.prompt_type = prompt_type
+
+        self._events = []
+        self.admissible_actions = None
+
+        self.step_count = 1
+
+    def update_step_count(self):
+        """Increment the current step counter."""
+        self.step_count += 1
+
+    def update_instruction_prompt(self, instruction: str):
+        """Set the instruction/system prompt used in chat mode."""
+        self.instruction = instruction
+
+    def update_single_obs_template(self, single_obs_template_wo_his: str, single_obs_template: str):
+        """
+        Set templates for single-prompt mode.
+
+        Args:
+            single_obs_template_wo_his (str): Template without history.
+            single_obs_template (str): Template with history.
+        """
+        self.single_obs_template_wo_his = single_obs_template_wo_his
+        self.single_obs_template = single_obs_template
+
+    def update_observation(self, obs: dict):
+        """Append an observation to the event history."""
+        self._events.append(
+            {
+                "type": "observation",
+                "text": obs,
+            }
+        )
+
+    def update_action(self, action: str):
+        """Append an action to the event history."""
+        self._events.append(
+            {
+                "type": "action",
+                "action": action,
+            }
+        )
+
+    def update_admissible_actions(self, admissible_actions):
+        """Update the list of admissible actions for the current step."""
+        self.admissible_actions = admissible_actions
+
+    def init(self, env):
+        """
+        Initialize the prompt builder at the beginning of an episode.
+
+        - Clears the event history
+        - Loads prompt instructions or templates from the environment
+        """
+        self._events.clear()
+
+        if self.prompt_type == "chat":
+            inst_prompt = env.get_instruction_prompt(info)
+            self.update_instruction_prompt(inst_prompt)
+        elif self.prompt_type == "single":
+            template_wo_his, template = env.get_single_prompt_template()
+            self.update_single_obs_template(template_wo_his, template)
+        else:
+            raise ValueError(f"Unsupported prompt_type: {self.prompt_type}")
+
+    def get_chat_prompt(self):
+        """
+        Construct a chat-style prompt from the event history.
+
+        Returns:
+            List[Message]: A sequence of User and Assistant messages.
+        """
+        if self.max_history != -1:
+            events = self._events[-(self.max_history * 2 + 1) :]
+        else:
+            events = self._events
+
+        messages = []
+
+        for idx, event in enumerate(events):
+            event_type = event.get("type")
+            message = None
+
+            if event_type == "observation":
+                content = event.get("text", "")
+                # Attach instruction prompt to the first observation
+                if idx == 0 and self.instruction:
+                    content += "\n" + self.instruction
+                message = UserMessage(source="user", content=content)
+
+            elif event_type == "action":
+                content = event.get("action", "")
+                message = AssistantMessage(source="assistant", content=content)
+
+            if message:
+                messages.append(message)
+
+        return messages
+
+    def get_single_prompt(self):
+        """
+        Construct a single formatted prompt using templates.
+
+        Returns:
+            List[Message]: A single User message.
+        """
+        if self.max_history != -1:
+            events = self._events[-(self.max_history * 2 + 1) :]
+        else:
+            events = self._events
+
+        current_obs = events[-1]["text"]
+
+        # Case 1: No history available
+        if len(events) == 1:
+            template = self.single_obs_template_wo_his
+            kwargs = {"current_observation": current_obs}
+            if "{admissible_actions}" in template:
+                kwargs["admissible_actions"] = self.admissible_actions
+            single_prompt = template.format(**kwargs)
+
+        # Case 2: History exists
+        else:
+            template = self.single_obs_template
+            history = ""
+            obs_count = 0
+            for idx, event in enumerate(events):
+                if events[idx]["type"] == "observation" and idx != len(events) - 1:
+                    next_event = events[idx + 1]
+                    history += f"[Observation {max(self.step_count-self.max_history+obs_count, 1)}: '{event['text']}', "
+                    history += (
+                        f"Action {max(self.step_count-self.max_history+obs_count, 1)}: '{next_event['action']}']\n "
+                    )
+                    obs_count += 1
+
+            kwargs = {
+                "step_count": self.step_count - 1,
+                "history_length": min(self.step_count - 1, self.max_history),
+                "history": history,
+                "current_step": self.step_count,
+                "current_observation": current_obs,
+            }
+            if "{admissible_actions}" in template:
+                kwargs["admissible_actions"] = self.admissible_actions
+            single_prompt = template.format(**kwargs)
+
+        return [UserMessage(source="user", content=single_prompt)]
+
+    def get_prompt(self):
+        """
+        Return the final prompt based on the configured prompt type.
+        """
+        if self.prompt_type == "chat":
+            prompt = self.get_chat_prompt()
+        elif self.prompt_type == "single":
+            prompt = self.get_single_prompt()
+
+        return prompt

contrib/recipes/envs/train_env_agent.py

@@ -0,0 +1,103 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import argparse
+import os
+import subprocess
+
+from omegaconf import OmegaConf
+
+from agentlightning import Trainer
+from agentlightning.algorithm.verl import VERL
+from contrib.agentlightning.contrib.algorithm.env_verl.daemon import EnvAgentModeDaemon
+from contrib.agentlightning.contrib.algorithm.env_verl.trainer import EnvAgentLightningTrainer
+
+
+def run_cmd(cmd):
+    """Execute a shell command and print its output"""
+    print(f"👉 Running: {cmd}")
+    result = subprocess.run(cmd, shell=True, text=True, capture_output=True)
+    if result.stdout:
+        print(result.stdout)
+    if result.stderr:
+        print(result.stderr)
+    return result
+
+
+def kill_process_on_port(port):
+    result = subprocess.run(f"sudo lsof -t -i :{port}", shell=True, capture_output=True, text=True)
+    pids = result.stdout.strip().split("\n")
+    for pid in pids:
+        if pid:
+            print(f"🔪 Killing process {pid} on port {port}")
+            subprocess.run(f"sudo kill -9 {pid}", shell=True)
+
+
+def train_val_dataset(cfg):
+    """Load training and validation datasets from parquet files."""
+    from datasets import Dataset
+
+    train_data = Dataset.from_parquet(cfg["data"]["train_files"])
+    val_data = Dataset.from_parquet(cfg["data"]["val_files"])
+    return train_data, val_data
+
+
+def get_config(path):
+    cfg = OmegaConf.load(path)
+    OmegaConf.resolve(cfg)
+    if "variables" in cfg:
+        del cfg["variables"]
+    return cfg
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--env", type=str, default="scienceworld")
+    parser.add_argument("--algorithm", type=str, default="grpo")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--n_workers", type=int, default=64, help="Number of workers for training")
+    parser.add_argument("--trial", type=int, default=0, help="Number of trials")
+    parser.add_argument("--task_num", type=int, default=25, help="ScienceWorld Task number to inject as env var")
+    parser.add_argument("--_background", action="store_true", help=argparse.SUPPRESS)
+    args = parser.parse_args()
+
+    # Restart Ray cluster cleanly
+    kill_process_on_port(4747)
+    run_cmd("pkill -f AgentLightning")
+    run_cmd("ray stop")
+    run_cmd("env RAY_DEBUG=legacy HYDRA_FULL_ERROR=1 VLLM_USE_V1=1 ray start --head --dashboard-host=0.0.0.0")
+
+    # set environment variable before loading configs
+    os.environ["TRIAL"] = str(args.trial)
+    if args.env == "scienceworld":
+        os.environ["TASK_NUM"] = str(args.task_num)
+
+    # Load configs
+    agent_config_path = f"config_env/{args.env}.yaml"
+    if args.debug:
+        trainer_config_path = f"config_verl/{args.env}/debug/{args.algorithm}.yaml"
+    else:
+        trainer_config_path = f"config_verl/{args.env}/{args.algorithm}.yaml"
+    agent_config = get_config(agent_config_path)
+
+    if "gigpo" in args.algorithm:
+        agent_config.log_env_obs = True
+    rl_training_config = get_config(trainer_config_path)
+
+    # Load datasets
+    train_dataset, val_dataset = train_val_dataset(rl_training_config)
+
+    # Initialize agent
+    from contrib.agentlightning.contrib.agent.env_agent import EnvAgent
+
+    agent = EnvAgent(agent_config)
+
+    # Initialize trainer and start training
+    trainer = Trainer(
+        algorithm=VERL(
+            config=rl_training_config,
+            trainer_cls=EnvAgentLightningTrainer,
+            daemon_cls=EnvAgentModeDaemon,
+        ),
+        n_workers=args.n_workers,
+    )
+    trainer.fit(agent, train_dataset, val_dataset=val_dataset)

contrib/recipes/search_r1/README.md

@@ -2,7 +2,7 @@
 
 ## Overview
 
-This example implements **Search R1** within Agent Lightning. It also serves as a demonstration of a **framework-free agent training pipeline**, showing how to run end-to-end RL training without relying on specialized frameworks. **It's tested and compatible with Agent-lightning v0.1.2**.
+This example implements **Search R1** within Agent Lightning. It also serves as a demonstration of a **framework-free agent training pipeline**, showing how to run end-to-end RL training without relying on specialized frameworks. **It's tested and compatible with Agent-lightning v0.2.x**.
 
 The example is designed to run on a single node with 8 GPUs, each having at least 40 GB of memory.
 
@@ -14,7 +14,7 @@ The example is designed to run on a single node with 8 GPUs, each having at leas
 | `retrieval_launch.sh` | Launches the retrieval service backed by the processed corpus |
 | `retrieval_server.py` | FastAPI server that powers document retrieval during training |
 | `search_r1_agent.py` | Agent-Lightning rollout script implementing the Search-R1 workflow |
-| `train.sh` | Starts the RL training server that coordinates GRPO optimization |
+| `train_search_r1_agent.py` | RL training script that coordinates GRPO optimization |
 | `qa_em.py` | Exact-match evaluation utilities for validating model predictions |
 
 ---
@@ -54,7 +54,7 @@ The retrieval server implementation is based on `search_r1/search/retrieval_serv
 
 ---
 
-## Run RL Training (GRPO) with Llama-3.2-3b-base
+## Run RL Training (GRPO) with Llama-3.2-3B-Instruct
 
 1. **Start Ray**
 
@@ -65,26 +65,28 @@ The retrieval server implementation is based on `search_r1/search/retrieval_serv
    > If you plan to use WandB for experiment tracking, set the environment variable
    > `WANDB_API_KEY` before starting Ray.
 
-2. **Launch the Agent**
-
-   ```bash
-   python search_r1_agent.py
-   ```
-
-   This script automatically launches **128 agent workers** by default. Each agent follows the Search-R1 workflow, retrieving information from the database and generating answers accordingly.
-
-
-3. **Start the Training Server**
+2. **Start the Training Server**
    In another terminal, run:
 
    ```bash
-   bash train.sh
+   python train_search_r1_agent.py llama
    ```
 
-   This script starts the RL training server.
+   This script starts the RL training. Each agent follows the Search-R1 workflow, retrieving information from the database and generating answers accordingly.
 
 ---
 
-## Evaluation
+## Benchmark Results
 
-Evaluation scripts and benchmark results will be released soon.
+We evaluated Search-R1 across seven diverse question-answering benchmarks, covering both General QA (NQ, TriviaQA, PopQA) and complex multi-hop reasoning tasks (HotpotQA, 2WikiMultiHopQA, Musique, and Bamboogle).
+
+The following tables compare the performance of the original Search-R1 implementation and the Agent-Lightning version across various base models.
+
+| Model | Source | NQ | TriviaQA | PopQA | HotpotQA | 2Wiki | Musique | Bamboogle |
+| :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| **Qwen2.5-3B-Instruct** | **Search-R1 (Original)** | 34.1 | 54.5 | 37.8 | 32.4 | 31.9 | 10.3 | 26.4 |
+| | **Agent-Lightning** | **45.3** | **61.7** | **43.8** | **42.6** | **36.4** | **17.1** | **37.6** |
+| **Qwen2.5-7B-Instruct** | **Search-R1 (Original)** | 39.3 | 61.0 | 39.7 | 37.0 | 41.4 | 14.6 | 36.8 |
+| | **Agent-Lightning** | **46.5** | **65.9** | **46.8** | **43.7** | **46.2** | **20.3** | **47.2** |
+| **Llama-3.2-3B** | **Search-R1 (Reproduced)** | 26.3 | 49.0 | 23.0 | 21.6 | 27.3 | 4.5 | 9.7 |
+| | **Agent-Lightning** | **29.6** | **51.9** | **25.7** | **23.2** | **28.3** | **5.8** | 9.6 |

contrib/recipes/search_r1/qa_em.py

@@ -75,7 +75,7 @@ def extract_solution(solution_str: str) -> Optional[str]:
     matches = list(match_iter)
 
     # If there are 0 or exactly 1 matches, return None
-    if len(matches) <= 1:
+    if len(matches) == 0:
         return None
 
     # If there are 2 or more matches, return the last one

contrib/recipes/search_r1/search_r1_agent.py

@@ -1,16 +1,21 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import os
 import re
+import time
 from typing import Any, Dict, List, Optional, Tuple, TypedDict, cast
 
+import pandas as pd
 import requests
 from openai import OpenAI
 from qa_em import compute_score_em
 
-from agentlightning import LLM, LitAgent, NamedResources, Trainer, reward, setup_logging
+from agentlightning import LLM, LitAgent, NamedResources, Rollout, Trainer, configure_logger, setup_logging
 
 setup_logging()
+logger = configure_logger(name=__name__)
 
 # Copied and adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/scripts/data_process/nq_search.py
 INSTRUCTION_FORMAT = """Answer the given question. You must conduct reasoning inside <think> and </think> first every time you get new information. After reasoning, if you find you lack some knowledge, you can call a search engine by <search> query </search> and it will return the top searched results between <information> and </information>. You can search as many times as your want. If you find no further external knowledge needed, you can directly provide the answer inside <answer> and </answer>, without detailed illustrations. For example, <answer> Beijing </answer>. Question: """
@@ -24,8 +29,7 @@ class RetrievalItem(TypedDict):
     document: Document
 
 
-@reward
-async def eval(prediction: str, ground_truth: List[str]) -> float:
+def eval(prediction: str, ground_truth: List[str]) -> float:
     reward_score = float(compute_score_em(prediction, ground_truth))
     print(f"pred: {prediction} | {type(ground_truth)} gold_answer: {ground_truth} | res: {reward_score}")
     return reward_score
@@ -106,62 +110,109 @@ def call_llm(
     return response.choices[0].message.content or ""
 
 
-class Searchr1Agent(LitAgent[Any]):
-    async def training_rollout_async(
+class SearchR1Agent(LitAgent[Dict[str, Any]]):
+
+    def __init__(
         self,
-        task: Any,
+        val_temperature: Optional[float] = 0.0,
+        max_turns: int = 4,
+    ) -> None:
+        super().__init__()
+        self.val_temperature = val_temperature
+        self.data_dir = os.environ.get("VERL_SEARCHR1_DATA_DIR", "data")
+        self.max_turns = max_turns
+
+    def rollout(
+        self,
+        task: Dict[str, Any],
         resources: NamedResources,
-        rollout: Any,
-        temperature: float = 1.0,
-    ) -> Any:
+        rollout: Rollout,
+    ) -> float | None:
         prompt = INSTRUCTION_FORMAT + task["question"]
         answer_list: List[str] = cast(List[str], task["golden_answers"])
-        llm: LLM = cast(LLM, resources.get("main_llm"))
+        rollout_id = rollout.rollout_id
+        logger.info(f"[Rollout {rollout_id}] Question: {task['question']}")
+        logger.info(f"[Rollout {rollout_id}] Ground Truth: {answer_list}")
+
+        start_time = time.time()
+        llm: LLM = cast(LLM, resources["main_llm"])
         client = OpenAI(
-            base_url=llm.endpoint,
+            base_url=llm.get_base_url(rollout_id, rollout.attempt.attempt_id),  # type: ignore
             api_key=os.environ.get("OPENAI_API_KEY", "token-abc123"),
         )
 
+        if rollout.mode == "train":
+            temperature = llm.sampling_parameters.get("temperature", 1.0)
+        else:
+            temperature = self.val_temperature if self.val_temperature is not None else 0.0
+
         turn_id = 0
         finished_flag = False
         rollout_content: str = ""
 
-        while turn_id < 4 and not finished_flag:
-            turn_id += 1
-            turn_response = call_llm(
-                client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
-            )
-            valid_turn_response = postprocess_response(turn_response)
-            turn_env_feedback = execute_response(valid_turn_response)
-            if len(turn_env_feedback) == 0:
-                finished_flag = True
-            print(f"TURN ID {turn_id} | RESP: {turn_response} | ENV FEEDBACK: {turn_env_feedback}")
-            rollout_content += turn_response + turn_env_feedback
+        try:
+            while turn_id < self.max_turns and not finished_flag:
+                turn_id += 1
+                turn_response = call_llm(
+                    client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
+                )
+                valid_turn_response = postprocess_response(turn_response)
+                rollout_content += valid_turn_response
+                turn_env_feedback = execute_response(valid_turn_response)
+                if len(turn_env_feedback) == 0:
+                    finished_flag = True
+                else:
+                    rollout_content += turn_env_feedback
+                logger.info(f"TURN ID {turn_id} | RESP: {turn_response} | ENV FEEDBACK: {turn_env_feedback}")
 
-        if not finished_flag:
-            turn_response = call_llm(
-                client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
-            )
-            rollout_content += turn_response
-            print(f"LAST TURN GENERATE | RESP: {turn_response}")
+            if not finished_flag:
+                turn_response = call_llm(
+                    client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
+                )
+                rollout_content += turn_response
+                logger.info(f"LAST TURN GENERATE | RESP: {turn_response}")
 
-        reward_score = await eval(rollout_content, answer_list)  # reward is tracked with the decorator
-        print(
+        except Exception as e:
+            logger.exception(f"[Rollout {rollout_id}] Error during rollout: {e}")
+            return None
+
+        end_time_rollout = time.time()
+        reward_score = eval(rollout_content, answer_list)
+        logger.info("[Rollout %s] Reward: %s", rollout_id, reward_score)
+        end_time_eval = time.time()
+
+        logger.info("[Rollout %s] Time taken for rollout: %.2f seconds", rollout_id, end_time_rollout - start_time)
+        logger.info(
+            "[Rollout %s] Time taken for evaluation: %.2f seconds", rollout_id, end_time_eval - end_time_rollout
+        )
+        logger.info(
             "question: {} answer: {} ground_truth: {} reward: {}".format(
                 task["question"], rollout_content, answer_list, reward_score
             )
         )
         return reward_score
 
-    async def validation_rollout_async(
-        self,
-        task: Any,
-        resources: NamedResources,
-        rollout: Any,
-    ) -> Any:
-        # Use the same resources; set temperature to 0.0 for deterministic validation.
-        return await self.training_rollout_async(task, resources, rollout, temperature=0.0)
+
+def debug_search_r1_agent():
+    searchr1_dev_data_path = os.path.join(os.environ.get("VERL_SEARCHR1_DATA_DIR", "data"), "test.parquet")
+    if not os.path.exists(searchr1_dev_data_path):
+        raise FileNotFoundError(f"Search_R1 dev data file {searchr1_dev_data_path} does not exist.")
+    df = pd.read_parquet(searchr1_dev_data_path).head(10)  # type: ignore
+    df = cast(List[Dict[str, Any]], df.to_dict(orient="records"))  # type: ignore
+    print("Debug data:", df)
+
+    trainer = Trainer(
+        n_workers=1,
+        initial_resources={
+            "main_llm": LLM(
+                endpoint=os.environ["OPENAI_API_BASE"],
+                model="gpt-4.1-nano",
+                sampling_parameters={"temperature": 0.0},
+            )
+        },
+    )
+    trainer.dev(SearchR1Agent(), df)
 
 
 if __name__ == "__main__":
-    Trainer(n_workers=128).fit(Searchr1Agent(), "http://localhost:9999/")
+    debug_search_r1_agent()

contrib/recipes/search_r1/train.sh

@@ -1,56 +0,0 @@
-#!/bin/bash
-
-set -e
-
-export N_GPUS=8
-export BASE_MODEL=meta-llama/Llama-3.2-3B
-export ROLLOUT_TP_SIZE=1
-export DATA_DIR=data
-export EXPERIMENT_NAME=searchr1
-export PROJECT_NAME=AgentLightning-searchr1
-echo "Starting training script..."
-
-python -m agentlightning.verl \
-    algorithm.adv_estimator=grpo \
-    data.train_files=${DATA_DIR}/train.parquet \
-    data.val_files=${DATA_DIR}/test.parquet \
-    actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP_SIZE} \
-    trainer.n_gpus_per_node=${N_GPUS} \
-    data.train_batch_size=512 \
-    actor_rollout_ref.rollout.n=5 \
-    actor_rollout_ref.actor.ppo_mini_batch_size=128 \
-    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
-    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
-    actor_rollout_ref.rollout.multi_turn.format=hermes \
-    actor_rollout_ref.model.path=${BASE_MODEL} \
-    data.max_prompt_length=4096 \
-    data.max_response_length=4096 \
-    data.truncation='error' \
-    trainer.val_before_train=True \
-    actor_rollout_ref.actor.optim.lr=1e-6 \
-    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.95 \
-    actor_rollout_ref.model.use_remove_padding=True \
-    actor_rollout_ref.actor.use_kl_loss=true \
-    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
-    actor_rollout_ref.actor.kl_loss_coef=0.001 \
-    actor_rollout_ref.actor.entropy_coeff=0 \
-    actor_rollout_ref.actor.clip_ratio_low=0.2 \
-    actor_rollout_ref.actor.clip_ratio_high=0.3 \
-    actor_rollout_ref.model.enable_gradient_checkpointing=True \
-    actor_rollout_ref.actor.fsdp_config.param_offload=True \
-    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
-    actor_rollout_ref.rollout.name=vllm \
-    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
-    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
-    actor_rollout_ref.ref.fsdp_config.param_offload=True \
-    algorithm.use_kl_in_reward=False \
-    trainer.critic_warmup=0 \
-    trainer.logger=['console','wandb'] \
-    trainer.default_local_dir=checkpoints/searchr1_checkpoints/$EXPERIMENT_NAME \
-    trainer.project_name=${PROJECT_NAME} \
-    trainer.experiment_name=${EXPERIMENT_NAME} \
-    trainer.nnodes=1 \
-    trainer.save_freq=10 \
-    trainer.test_freq=20 \
-    trainer.total_epochs=15 \
-    trainer.total_training_steps=300

contrib/recipes/search_r1/train_search_r1_agent.py

@@ -0,0 +1,171 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from __future__ import annotations
+
+import argparse
+import os
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict
+
+import pandas as pd
+from search_r1_agent import SearchR1Agent
+
+import agentlightning as agl
+
+RL_TRAINING_CONFIG: Dict[str, Any] = {
+    "algorithm": {
+        "adv_estimator": "grpo",
+        "use_kl_in_reward": False,
+    },
+    "data": {
+        "train_files": "data/train.parquet",
+        "val_files": "data/test.parquet",
+        "train_batch_size": 512,
+        "max_prompt_length": 6000,
+        "max_response_length": 4096,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 5,
+            "log_prob_micro_batch_size_per_gpu": 4,
+            "multi_turn": {"format": "hermes"},
+            "name": "vllm",
+            "gpu_memory_utilization": 0.5,
+            "engine_kwargs": {
+                "vllm": {
+                    "enable_auto_tool_choice": True,
+                    "tool_call_parser": "hermes",
+                }
+            },
+        },
+        "actor": {
+            "ppo_mini_batch_size": 256,
+            "ppo_micro_batch_size_per_gpu": 4,
+            "optim": {"lr": 1e-6, "lr_warmup_steps_ratio": 0.95},
+            "use_kl_loss": True,
+            "kl_loss_type": "low_var_kl",
+            "kl_loss_coef": 0.001,
+            "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": 4,
+            "fsdp_config": {"param_offload": True},
+        },
+        "model": {
+            "path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 8,
+        "val_before_train": True,
+        "critic_warmup": 0,
+        "logger": ["console", "wandb"],
+        "project_name": "AgentLightning",
+        "experiment_name": "searchr1",
+        "nnodes": 1,
+        "test_freq": 10,
+        "save_freq": 10,
+        "total_epochs": 15,
+        "total_training_steps": 300,
+        "default_local_dir": "checkpoints/searchr1_checkpoints/",
+    },
+}
+
+
+def config_train_fast() -> Dict[str, Any]:
+    """A fast training run for CI testing purposes."""
+
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    EXPERIMENT_NAME = f"searchr1_{timestamp}"
+    PROJECT_NAME = "AgentLightningCI"
+
+    # Simulate writing to $GITHUB_OUTPUT if it’s set
+    github_output = os.getenv("GITHUB_OUTPUT")
+    if github_output:
+        with open(github_output, "a") as f:
+            f.write(f"project_name={PROJECT_NAME}\n")
+            f.write(f"run_name={EXPERIMENT_NAME}\n")
+
+    print("Set environment variables:")
+    print(f"PROJECT_NAME={PROJECT_NAME}")
+    print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
+    config["actor_rollout_ref"]["model"]["path"] = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    config["data"]["val_files"] = "data/test_dev.parquet"
+    config["trainer"]["total_epochs"] = 1
+    config["trainer"]["total_training_steps"] = 1
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    config["trainer"]["test_freq"] = 1
+    return config
+
+
+def config_train_qwen() -> Dict[str, Any]:
+    """A configuration for training with Qwen-2.5."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    return config
+
+
+def config_train_llama() -> Dict[str, Any]:
+    """A configuration for training with LLaMA-3.2-3B-Instruct.
+
+    You will need a `HF_TOKEN` set to run with this config.
+    """
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["multi_turn"]["format"] = "llama3_json"
+    config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["tool_call_parser"] = "llama3_json"
+    config["actor_rollout_ref"]["model"]["path"] = "meta-llama/Llama-3.2-3B-Instruct"
+    return config
+
+
+def train(config: Dict[str, Any]) -> None:
+
+    agent = SearchR1Agent()
+    algorithm = agl.VERL(config)
+    trainer = agl.Trainer(n_runners=32, algorithm=algorithm)
+
+    train_data = pd.read_parquet(config["data"]["train_files"]).to_dict(orient="records")  # type: ignore
+    val_data = pd.read_parquet(config["data"]["val_files"]).to_dict(orient="records")  # type: ignore
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
+
+
+def main() -> None:
+    """Main function to parse arguments and run training."""
+    parser = argparse.ArgumentParser(description="Train a Search-R1 agent using different model configurations")
+
+    parser.add_argument(
+        "config",
+        choices=["fast", "qwen", "llama"],
+        help="Training configuration: 'fast' (CI testing), 'qwen' (Qwen-2.5-Coder-1.5B), 'llama' (LLaMA-3.2-3B-Instruct)",
+    )
+
+    args = parser.parse_args()
+
+    # Get the appropriate configuration
+    config_functions = {"fast": config_train_fast, "qwen": config_train_qwen, "llama": config_train_llama}
+
+    config = config_functions[args.config]()
+
+    print(f"Starting training with '{args.config}' configuration...")
+
+    train(config)
+
+
+if __name__ == "__main__":
+    main()

contrib/recipes/webshop/.gitignore

@@ -0,0 +1,13 @@
+node_modules/
+.venv/
+.next/
+dist/
+*.tsbuildinfo
+server/webshop/
+
+# Log files from make train
+logs/
+std_log.txt
+
+# Auto-generated by Next.js
+next-env.d.ts