remove debug info

update tests
.
2025-11-21 10:30:47 +08:00 · 2025-11-21 09:58:13 +08:00 · 2025-11-21 08:47:11 +08:00 · 2025-11-21 01:46:01 +08:00 · 2025-11-21 01:45:06 +08:00 · 2025-11-21 00:50:05 +08:00
112 changed files with 12935 additions and 4369 deletions
@@ -0,0 +1,29 @@
+name: Badge - Compatibility
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Backward Compatibility
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-compat.yml', label: 'examples-compat', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,31 @@
+name: Badge - Unit Test
+
+on:
+  workflow_run:
+    workflows:
+      - CPU Test
+      - GPU Test
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'tests-full.yml', label: 'tests-full', variants: ['legacy', 'stable'] },
+              { workflow: 'tests.yml', label: 'tests', variants: ['legacy', 'stable', 'Lint', 'documentation', 'JavaScript'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'APO - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Calc-X - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -22,12 +22,12 @@ run-name: >-
      || format('Calc-X - {0}', github.event_name) }}

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

@@ -116,13 +116,11 @@ jobs:
      # Don't ask why. Don't touch this.
      - name: Calc-X training
        run: |
-          set -ex
          source .venv/bin/activate
          cd examples/calc_x
          ../../scripts/restart_ray.sh
          sleep 5
-          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci
-          sleep 10
+          python train_calc_agent.py --val-file data/test_mini.parquet --ci
        shell: bash
        env:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
@@ -137,14 +135,126 @@ jobs:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}

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

-      - name: Calc-X training with external store
+      - name: Validate training with LLM Proxy
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_llm_proxy.outputs.project_name }} ${{ steps.calc_x_train_llm_proxy.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Training with external store
        run: |
          set -euo pipefail
          source .venv/bin/activate
@@ -182,7 +300,15 @@ jobs:
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
        id: calc_x_train_external_store

-      - name: Calc-X training with role-based environment variables
+      - name: Validate training with external store
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_external_store.outputs.project_name }} ${{ steps.calc_x_train_external_store.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Training with role-based environment variables
        run: |
          set -euo pipefail
          source .venv/bin/activate
@@ -203,3 +329,12 @@ jobs:
        env:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_role_based_env_var
+
+      - name: Validate training with role-based environment variables
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_role_based_env_var.outputs.project_name }} ${{ steps.calc_x_train_role_based_env_var.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Backward Compatibility - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -56,7 +56,7 @@ jobs:
            --group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
      - name: Override VERL (stable)
        run: |
-          uv pip install verl==0.5.0
+          uv pip install verl==0.5.0 vllm==0.10.2
        if: matrix.setup-script == 'stable'
      - name: Freeze dependencies
        run: |
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Spider - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -121,7 +121,7 @@ jobs:
      - name: Validate Spider training
        run: |
          set -ex
-          uv run scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }}
+          uv run scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }} --reward-tolerance 5
        env:
          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'Unsloth - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -14,7 +14,7 @@ on:
 run-name: >-
  ${{ github.event_name == 'repository_dispatch'
      && format(
-        'PR #{0} - Label {1} - {2}',
+        'GPU Test - PR #{0} - {1} - {2}',
        github.event.client_payload.pull_number,
        github.event.client_payload.ci_label,
        github.event.client_payload.correlation_id
@@ -29,6 +29,113 @@ jobs:
      github.event.action == 'ci-all'
    name: GPU Test with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})

+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-tests-full-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Build dashboard
+        run: cd dashboard && npm run build
+
+      - name: Start MongoDB container
+        run: |
+          set -euo pipefail
+          cat /etc/security/limits.conf
+          docker run -d \
+            --name mongodb-test \
+            --ulimit nofile=65535:65535 \
+            -p 27017:27017 \
+            mongo:8.2 \
+            --replSet test-rs
+
+          # Wait for mongod to come up
+          for i in $(seq 1 30); do
+            if docker exec mongodb-test mongosh --quiet --eval 'db.runCommand({ ping: 1 })' >/dev/null 2>&1; then
+              echo "Mongo is up"
+              break
+            fi
+            echo "Waiting for Mongo..."
+            sleep 2
+          done
+
+          # Init replica set (simple single-node)
+          docker exec mongodb-test mongosh --quiet --eval '
+            rs.initiate({
+              _id: "test-rs",
+              members: [{ _id: 0, host: "localhost:27017" }]
+            })
+          '
+        shell: bash
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: Run tests
+        run: |
+          uv run pytest -v --durations=0 tests
+        env:
+          PYTEST_ADDOPTS: "--color=yes"
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+          AGL_TEST_MONGO_URI: mongodb://localhost:27017/?replicaSet=test-rs
+
+
+  minimal-examples:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-gpu' ||
+      github.event.action == 'ci-all'
+    name: Minimal Examples with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
+
    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
    timeout-minutes: 30
    strategy:
@@ -69,18 +176,10 @@ jobs:
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          name: dependencies-minimal-examples-${{ matrix.python-version }}-${{ matrix.setup-script }}
          path: requirements-freeze.txt
          compression-level: 0

-      - uses: actions/setup-node@v6
-        with:
-          node-version: '22'
-      - name: Install JavaScript dependencies
-        run: cd dashboard && npm ci
-      - name: Build dashboard
-        run: cd dashboard && npm run build
-
      - name: Launch LiteLLM Proxy
        run: |
          ./scripts/litellm_run.sh
@@ -88,10 +187,120 @@ jobs:
          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}

-      - name: Run tests
+      - name: Write Traces via Otel Tracer
        run: |
-          uv run pytest -v --durations=0 tests
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_traces.py otel
+          sleep 5
+
+      - name: Write Traces via AgentOps Tracer
        env:
-          PYTEST_ADDOPTS: "--color=yes"
          OPENAI_BASE_URL: http://localhost:12306/
          OPENAI_API_KEY: dummy
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python write_traces.py agentops
+          sleep 5
+
+      - name: Write Traces via Otel Tracer with Client
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          agl store --port 45993 --log-level DEBUG &
+          sleep 5
+          python write_traces.py otel --use-client
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+
+      - name: Write Traces via AgentOps Tracer with Client
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          agl store --port 45993 --log-level DEBUG &
+          sleep 5
+          python write_traces.py agentops --use-client
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+
+      - name: vLLM Server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python vllm_server.py Qwen/Qwen2.5-0.5B-Instruct
+
+      - name: LLM Proxy (OpenAI backend)
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+
+          python llm_proxy.py openai gpt-4.1-mini &
+
+          LLM_PROXY_READY=0
+          for attempt in $(seq 1 30); do
+            if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
+              LLM_PROXY_READY=1
+              break
+            fi
+            sleep 2
+          done
+          if [[ "$LLM_PROXY_READY" != "1" ]]; then
+            echo "LLM proxy failed to become healthy" >&2
+            exit 1
+          fi
+
+          python llm_proxy.py test gpt-4.1-mini
+
+          pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
+          while pgrep -f llm_proxy.py; do
+            echo "Waiting for llm_proxy.py to finish..."
+            sleep 5
+          done
+
+      - name: LLM Proxy (vLLM backend)
+        if: matrix.setup-script != 'legacy'  # Skip if return_token_ids is not supported
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/minimal
+          python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct &
+
+          LLM_PROXY_READY=0
+          for attempt in $(seq 1 30); do
+            if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
+              LLM_PROXY_READY=1
+              break
+            fi
+            sleep 2
+          done
+          if [[ "$LLM_PROXY_READY" != "1" ]]; then
+            echo "LLM proxy failed to become healthy" >&2
+            exit 1
+          fi
+
+          python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
+
+          pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
+          while pgrep -f llm_proxy.py; do
+            echo "Waiting for llm_proxy.py to finish..."
+            sleep 5
+          done
@@ -37,6 +37,7 @@ jobs:
          uv sync --frozen \
            --extra apo \
            --extra verl \
+            --extra mongo \
            --group dev \
            --group torch-cpu \
            --group torch-stable \
@@ -166,7 +167,7 @@ jobs:

      - name: Run tests
        run: |
-          uv run pytest -v --durations=0 tests
+          uv run pytest -v --durations=0 tests -m "not mongo"
        env:
          PYTEST_ADDOPTS: "--color=yes"

@@ -4,7 +4,7 @@

 # Agent Lightning⚡

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

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

 To start using Agent-lightning, check out our [documentation](https://microsoft.github.io/agent-lightning/) and [examples](./examples).
@@ -69,10 +75,11 @@ No rewrites, no lock-in, just a clear path from first rollout to steady improvem
 | Workflow | Status |
 |----------|--------|
 | CPU Tests | [![tests workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml) |
-| GPU Tests | [![tests-full workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml) |
+| Full Tests | [![tests summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
+| UI Tests | [![UI Tests](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml) |
 | Examples Integration | [![examples summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml) |
 | Latest Dependency Compatibility | [![latest summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml) |
-| Legacy Examples Compatibility | [![examples compatibility workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml) |
+| Legacy Examples Compatibility | [![compat summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml) |

 ## ⚡ Citation

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

-__version__ = "0.2.2"
+__version__ = "0.3.0"

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

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

 from opentelemetry.sdk.trace import ReadableSpan

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


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

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


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

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

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

 from pydantic import TypeAdapter

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

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

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

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

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

 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel
@@ -670,7 +670,7 @@ class TracerTraceToTriplet(TraceToTripletBase):
        trace_tree.visualize(filename, interested_span_match=interested_span_match)
        return trace_tree

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

        Args:
@@ -800,7 +800,7 @@ class LlmProxyTraceToTriplet(TraceToTripletBase):
        rid = attrs.get("gen_ai.response.id") or attrs.get("llm.hosted_vllm.id")
        return str(rid) if isinstance(rid, str) and rid else None

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

        Args:
@@ -143,7 +143,7 @@ class Baseline(FastAlgorithm):
        store = self.get_store()

        for index in train_indices + val_indices:
-            queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+            queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
            if len(queuing_rollouts) <= 1:
                # Only enqueue a new rollout when there is at most 1 rollout in the queue.
                sample = dataset[index]
@@ -222,7 +222,7 @@ class Baseline(FastAlgorithm):
                    f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
                )
                while True:
-                    queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+                    queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
                    if len(queuing_rollouts) <= self.max_queue_length:
                        # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
                        sample = concatenated_dataset[index]
@@ -9,7 +9,7 @@ import asyncio
 import logging
 from typing import Iterable

-from agentlightning.logging import configure_logger
+from agentlightning import setup_logging
 from agentlightning.store.client_server import LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore

@@ -25,9 +25,15 @@ def main(argv: Iterable[str] | None = None) -> int:
        action="append",
        help="Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.",
    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the store.",
+    )
    args = parser.parse_args(list(argv) if argv is not None else None)

-    configure_logger()
+    setup_logging(args.log_level)

    store = InMemoryLightningStore()
    server = LightningStoreServer(
@@ -35,6 +41,7 @@ def main(argv: Iterable[str] | None = None) -> int:
        host="0.0.0.0",
        port=args.port,
        cors_allow_origins=args.cors_origins,
+        launch_mode="asyncio",
    )
    try:
        asyncio.run(server.run_forever())
@@ -129,12 +129,13 @@ def reward(fn: FnType) -> FnType:
        return wrapper  # type: ignore


-def emit_reward(reward: float) -> ReadableSpan:
+def emit_reward(reward: float, auto_export: bool = True) -> ReadableSpan:
    """Emit a reward value as an OpenTelemetry span.

    Args:
        reward: Numeric reward to record. Integers and booleans are converted to
            floating point numbers for consistency.
+        auto_export: Whether to export the span automatically.

    Returns:
        Readable span capturing the recorded reward.
@@ -150,7 +151,7 @@ def emit_reward(reward: float) -> ReadableSpan:
        raise ValueError(f"Reward must be a number, got: {type(reward)}")

    # TODO: This should use the tracer from current context by tracer
-    tracer = get_tracer()
+    tracer = get_tracer(use_active_span_processor=auto_export)
    span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
    # Do nothing; it's just a number
    with span:
@@ -2,13 +2,22 @@

 """Utilities shared across emitter implementations."""

+from typing import cast
+from warnings import filterwarnings
+
 import opentelemetry.trace as trace_api
+from opentelemetry.sdk.trace import SpanLimits, SynchronousMultiSpanProcessor, Tracer
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
 from opentelemetry.trace import get_tracer_provider


-def get_tracer() -> trace_api.Tracer:
+def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
    """Resolve the OpenTelemetry tracer configured for Agent Lightning.

+    Args:
+        use_active_span_processor: Whether to use the active span processor.
+
    Returns:
        OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.

@@ -18,5 +27,31 @@ def get_tracer() -> trace_api.Tracer:
    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")

-    tracer_provider = get_tracer_provider()
-    return tracer_provider.get_tracer("agentlightning")
+    tracer_provider = cast(TracerProviderImpl, get_tracer_provider())
+
+    if use_active_span_processor:
+        return tracer_provider.get_tracer("agentlightning")
+
+    else:
+        filterwarnings(
+            "ignore",
+            message=r"You should use InstrumentationScope. Deprecated since version 1.11.1.",
+            category=DeprecationWarning,
+            module="opentelemetry.sdk.trace",
+        )
+
+        return Tracer(
+            tracer_provider.sampler,
+            tracer_provider.resource,
+            # We use an empty span processor to avoid emitting spans to the tracer
+            SynchronousMultiSpanProcessor(),
+            tracer_provider.id_generator,
+            InstrumentationInfo("agentlightning", "", ""),  # type: ignore
+            SpanLimits(),
+            InstrumentationScope(
+                "agentlightning",
+                "",
+                "",
+                {},
+            ),
+        )
@@ -67,10 +67,11 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
        server_host: str | None = None,
        server_port: int | None = None,
        n_runners: int = 1,
-        graceful_timeout: float = 5.0,
-        terminate_timeout: float = 5.0,
+        graceful_timeout: float = 10.0,
+        terminate_timeout: float = 10.0,
        main_process: Literal["algorithm", "runner"] = "algorithm",
        managed_store: bool | None = None,
+        allowed_exit_codes: Iterable[int] = (0, -15),
    ) -> None:
        """Configure the strategy.

@@ -94,6 +95,9 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                LightningStore client/server wrappers automatically. When
                `False` the provided `store` is passed directly to the
                bundles, allowing callers to manage store wrappers manually.
+            allowed_exit_codes: Allowed exit codes for subprocesses.
+                By default, runner can exit gracefully with code 0 or terminated
+                by SIGTERM (-15).
        """
        if role is None:
            role_env = os.getenv("AGL_CURRENT_ROLE")
@@ -133,6 +137,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                raise ValueError("main_process='runner' requires n_runners to be 1")
        self.main_process = main_process
        self.managed_store = resolve_managed_store_flag(managed_store)
+        self.allowed_exit_codes = tuple(allowed_exit_codes)

    async def _execute_algorithm(
        self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent
@@ -338,10 +343,10 @@ class ClientServerExecutionStrategy(ExecutionStrategy):

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

    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
        logger.info(
@@ -13,7 +13,8 @@ from agentops.sdk.exporters import AuthenticatedOTLPExporter
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.metrics.export import MetricExportResult
-from opentelemetry.sdk.trace.export import SpanExportResult
+
+from agentlightning.utils.otlp import LightningStoreOTLPExporter

 logger = logging.getLogger(__name__)

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

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


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

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

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

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

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


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

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


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


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

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


 class BypassableV3Client(V3Client):
@@ -40,12 +40,14 @@ from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
 from litellm.proxy.proxy_server import app, save_worker_config  # pyright: ignore[reportUnknownVariableType]
 from litellm.types.utils import CallTypes
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.types import Scope

-from agentlightning.types import LLM, ProxyLLM
+from agentlightning.types import LLM, ProxyLLM, SpanNames
 from agentlightning.utils.server_launcher import (
    LaunchMode,
    PythonServerLauncher,
@@ -192,7 +194,7 @@ class LightningSpanExporter(SpanExporter):
    def __init__(self, _store: Optional[LightningStore] = None):
        self._store: Optional[LightningStore] = _store  # this is only for testing purposes
        self._buffer: List[ReadableSpan] = []
-        self._lock: Optional[threading.RLock] = None
+        self._lock: Optional[threading.Lock] = None
        self._loop_lock_pid: Optional[int] = None

        # Single dedicated event loop running in a daemon thread.
@@ -201,6 +203,8 @@ class LightningSpanExporter(SpanExporter):
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._loop_thread: Optional[threading.Thread] = None

+        self._otlp_exporter = OTLPSpanExporter()
+
    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
        """Lazily initialize the event loop and thread on first use.

@@ -214,15 +218,15 @@ class LightningSpanExporter(SpanExporter):
            self._loop_thread.start()
        return self._loop

-    def _ensure_lock(self) -> threading.RLock:
+    def _ensure_lock(self) -> threading.Lock:
        """Lazily initialize the lock on first use.

        Returns:
-            threading.RLock: The initialized lock.
+            threading.Lock: The initialized lock.
        """
        self._clear_loop_and_lock()
        if self._lock is None:
-            self._lock = threading.RLock()
+            self._lock = threading.Lock()
        return self._lock

    def _clear_loop_and_lock(self) -> None:
@@ -284,24 +288,18 @@ class LightningSpanExporter(SpanExporter):
        with self._ensure_lock():
            for span in spans:
                self._buffer.append(span)
-
-        # Run the async flush on our private loop, synchronously from caller's POV.
-        async def _locked_flush():
-            # Take the lock inside the coroutine to serialize with other flushes.
-            with self._ensure_lock():
-                return await self._maybe_flush()
-
-        try:
-            loop = self._ensure_loop()
-            fut = asyncio.run_coroutine_threadsafe(_locked_flush(), loop)
-            fut.result()  # Bubble up any exceptions from the coroutine.
-        except Exception as e:
-            logger.exception("Export flush failed: %s", e)
-            return SpanExportResult.FAILURE
+            default_endpoint = self._otlp_exporter._endpoint  # pyright: ignore[reportPrivateUsage]
+            try:
+                self._maybe_flush()
+            except Exception as e:
+                logger.exception("Export flush failed: %s", e)
+                return SpanExportResult.FAILURE
+            finally:
+                self._otlp_exporter._endpoint = default_endpoint  # pyright: ignore[reportPrivateUsage]

        return SpanExportResult.SUCCESS

-    async def _maybe_flush(self):
+    def _maybe_flush(self):
        """Flush ready subtrees from the buffer.

        Strategy:
@@ -323,11 +321,20 @@ class LightningSpanExporter(SpanExporter):
            if not subtree_spans:
                continue

+            # Store is initialized lazily here in most cases.
            store = self._store or get_active_llm_proxy().get_store()
            if store is None:
                logger.warning("Store is not set in LLMProxy. Cannot log spans to store.")
                continue

+            # If the store supports OTLP endpoint, use it.
+            if store.capabilities.get("otlp_traces", False):
+                otlp_traces_endpoint = store.otlp_traces_endpoint()
+                self._otlp_exporter._endpoint = otlp_traces_endpoint  # pyright: ignore[reportPrivateUsage]
+                otlp_enabled = True
+            else:
+                otlp_enabled = False
+
            # Merge all custom headers found in the subtree.
            headers_merged: Dict[str, Any] = {}

@@ -383,10 +390,34 @@ class LightningSpanExporter(SpanExporter):
            sequence_id_decimal = int(sequence_id)

            # Persist each span in the subtree with the resolved identifiers.
-            for span in subtree_spans:
-                await store.add_otel_span(
-                    rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id_decimal, readable_span=span
-                )
+            if otlp_enabled:
+                # If store has OTLP support, directly use OTLP exporter and export in batch
+                for span in subtree_spans:
+                    span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
+                        Resource.create(
+                            {
+                                SpanNames.ROLLOUT_ID: rollout_id,
+                                SpanNames.ATTEMPT_ID: attempt_id,
+                                SpanNames.SPAN_SEQUENCE_ID: sequence_id_decimal,
+                            }
+                        )
+                    )
+                export_result = self._otlp_exporter.export(subtree_spans)
+                if export_result != SpanExportResult.SUCCESS:
+                    raise RuntimeError(f"Failed to export spans via OTLP exporter. Result: {export_result}")
+
+            else:
+                # The old way: store does not support OTLP endpoint
+                for span in subtree_spans:
+                    loop = self._ensure_loop()
+                    add_otel_span_task = store.add_otel_span(
+                        rollout_id=rollout_id,
+                        attempt_id=attempt_id,
+                        sequence_id=sequence_id_decimal,
+                        readable_span=span,
+                    )
+                    fut = asyncio.run_coroutine_threadsafe(add_otel_span_task, loop)
+                    fut.result()  # Bubble up any exceptions from the coroutine.

    def _get_root_span_ids(self) -> Iterable[int]:
        """Yield span_ids for root spans currently in the buffer.
@@ -822,7 +853,6 @@ class StreamConversionMiddleware(BaseHTTPMiddleware):
        )  # e.g., "stop", "length", "tool_calls", "content_filter"

        def sse_chunk(obj: Dict[str, Any]) -> str:
-            print("sse_chunk: ", obj)
            return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"

        # 1) initial chunk with the role
@@ -1209,16 +1239,16 @@ class LLMProxy:
        if self.store is None:
            raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")

-        store_capabilities = self.store.capabilities()
-        if self.server_launcher.args.launch_mode == "mp" and not store_capabilities["zero_copy"]:
+        store_capabilities = self.store.capabilities
+        if self.server_launcher.args.launch_mode == "mp" and not store_capabilities.get("zero_copy", False):
            raise RuntimeError(
                "The store does not support zero-copy. Please use another store, or use asyncio or thread mode to launch the server."
            )
-        elif self.server_launcher.args.launch_mode == "thread" and not store_capabilities["thread_safe"]:
+        elif self.server_launcher.args.launch_mode == "thread" and not store_capabilities.get("thread_safe", False):
            raise RuntimeError(
                "The store is not thread-safe. Please use another store, or use asyncio mode to launch the server."
            )
-        elif self.server_launcher.args.launch_mode == "asyncio" and not store_capabilities["async_safe"]:
+        elif self.server_launcher.args.launch_mode == "asyncio" and not store_capabilities.get("async_safe", False):
            raise RuntimeError("The store is not async-safe. Please use another store.")

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

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

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


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

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

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

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

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

 import asyncio
 import logging
+import random
 import threading
 import time
 from contextlib import suppress
@@ -72,6 +73,7 @@ class LitAgentRunner(Runner[T_task]):
        max_rollouts: Optional[int] = None,
        poll_interval: float = 5.0,
        heartbeat_interval: float = 10.0,
+        interval_jitter: float = 0.1,
        heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
    ) -> None:
        """Initialize the agent runner.
@@ -82,6 +84,9 @@ class LitAgentRunner(Runner[T_task]):
                [`iter`][agentlightning.LitAgentRunner.iter].
            poll_interval: Seconds to wait between store polls when no work is available.
            heartbeat_interval: Seconds to wait between sending heartbeats to the store.
+            interval_jitter: Jitter factor for the poll interval. The actual interval will be between
+                poll_interval - interval_jitter and poll_interval + interval_jitter.
+                This is to avoid the overload caused by the synchronization of the runners.
            heartbeat_launch_mode: Launch mode for the heartbeat loop. Can be "asyncio" or "thread".
                "asyncio" is the default and recommended mode. Use "thread" if you are experiencing blocking coroutines.
        """
@@ -90,7 +95,9 @@ class LitAgentRunner(Runner[T_task]):
        self._max_rollouts = max_rollouts
        self._poll_interval = poll_interval
        self._heartbeat_interval = heartbeat_interval
+        self._interval_jitter = interval_jitter
        self._heartbeat_launch_mode = heartbeat_launch_mode
+        self._random_state = random.Random()

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

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

    def teardown(self, *args: Any, **kwargs: Any) -> None:
        """Teardown the runner and clean up all resources.
@@ -279,8 +286,9 @@ class LitAgentRunner(Runner[T_task]):
        if isinstance(raw_result, float):
            # Preserve the existing spans before another span is emitted
            trace_spans = list(self._tracer.get_last_trace())
-            # This will emit another span to the tracer
-            reward_span = emit_reward(raw_result)
+            # This will NOT emit another span to the tracer
+            reward_span = emit_reward(raw_result, auto_export=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)

@@ -359,7 +367,11 @@ class LitAgentRunner(Runner[T_task]):
                while not stop_event.is_set():
                    await self._emit_heartbeat(store)
                    with suppress(asyncio.TimeoutError):
-                        await asyncio.wait_for(stop_event.wait(), timeout=self._heartbeat_interval)
+                        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")

@@ -378,7 +390,11 @@ class LitAgentRunner(Runner[T_task]):
                asyncio.set_event_loop(loop)
                while not stop_evt.is_set():
                    loop.run_until_complete(self._emit_heartbeat(store))
-                    stop_evt.wait(self._heartbeat_interval)
+                    interval = self._heartbeat_interval + self._random_state.uniform(
+                        -self._interval_jitter, self._interval_jitter
+                    )
+                    interval = max(interval, 0.01)
+                    stop_evt.wait(interval)

            thread = threading.Thread(target=thread_worker, name=f"{self.get_worker_id()}-heartbeat", daemon=True)
            thread.start()
@@ -401,11 +417,13 @@ class LitAgentRunner(Runner[T_task]):
            event: Optional [`ExecutionEvent`][agentlightning.ExecutionEvent] object that can be used to interrupt the sleep.
                If set during the sleep period, the method returns immediately.
        """
+        interval = self._poll_interval + self._random_state.uniform(-self._interval_jitter, self._interval_jitter)
+        interval = max(interval, 0.01)
        if event is None:
-            await asyncio.sleep(self._poll_interval)
+            await asyncio.sleep(interval)
            return
        current_time = time.time()
-        next_time = current_time + self._poll_interval
+        next_time = current_time + interval
        while time.time() < next_time:
            await asyncio.sleep(0.1)
            if event.is_set():
@@ -451,7 +469,7 @@ class LitAgentRunner(Runner[T_task]):

            start_time = time.time()
            async with self._tracer.trace_context(
-                name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
+                name=rollout_id, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
            ):
                await self._trigger_hooks(
                    hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
@@ -2,6 +2,7 @@

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

@@ -11,5 +12,6 @@ __all__ = [
    "LightningStoreClient",
    "LightningStoreServer",
    "InMemoryLightningStore",
+    "CollectionBasedLightningStore",
    "LightningStoreThreaded",
 ]
@@ -18,6 +18,7 @@ from agentlightning.types import (
    Span,
    TaskInput,
    Worker,
+    WorkerStatus,
 )


@@ -53,8 +54,11 @@ UNSET = _UnsetType()
 Unset = _UnsetType  # Alias for convenience


-class LightningStoreCapabilities(TypedDict):
-    """Capability of a LightningStore implementation."""
+class LightningStoreCapabilities(TypedDict, total=False):
+    """Capability of a LightningStore implementation.
+
+    All keys are optional and false by default.
+    """

    thread_safe: bool
    """Whether the store is thread-safe."""
@@ -62,6 +66,8 @@ class LightningStoreCapabilities(TypedDict):
    """Whether the store is async-safe."""
    zero_copy: bool
    """Whether the store has only one copy across all threads/processes."""
+    otlp_traces: bool
+    """Whether the store supports OTLP/HTTP traces."""


 class LightningStore:
@@ -86,14 +92,32 @@ class LightningStore:
    Unless stated otherwise, missing identifiers should result in a `ValueError`.
    """

+    @property
    def capabilities(self) -> LightningStoreCapabilities:
        """Return the capabilities of the store."""
        return LightningStoreCapabilities(
            thread_safe=False,
            async_safe=False,
            zero_copy=False,
+            otlp_traces=False,
        )

+    def otlp_traces_endpoint(self) -> str:
+        """Return the OTLP/HTTP traces endpoint of the store.
+
+        The traces can have rollout ID and attempt ID (and optionally sequence ID)
+        saved in the "resource" of the spans.
+        The store, if it supports OTLP, should be able to receive the traces and save them
+        via [`add_span`][agentlightning.LightningStore.add_span] or
+        [`add_otel_span`][agentlightning.LightningStore.add_otel_span].
+
+        The endpoint should be compatible with [OTLP HTTP protocol](https://opentelemetry.io/docs/specs/otlp/).
+        It's not necessarily compatible with OTLP gRPC protocol.
+
+        The returned endpoint will usually ends with `/v1/traces`.
+        """
+        raise NotImplementedError()
+
    async def start_rollout(
        self,
        input: TaskInput,
@@ -269,30 +293,77 @@ class LightningStore:
        raise NotImplementedError()

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

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

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

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

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

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

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

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

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

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

        Raises:
            NotImplementedError: Subclasses must implement retrieval.
@@ -416,19 +511,61 @@ class LightningStore:
        """
        raise NotImplementedError()

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

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

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

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

        Raises:
            NotImplementedError: Subclasses must implement the query.
@@ -555,11 +692,29 @@ class LightningStore:

    async def query_workers(
        self,
-    ) -> List[Worker]:
+        *,
+        status_in: Optional[Sequence[WorkerStatus]] = None,
+        worker_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Worker]:
        """Query all workers in the system.

+        Args:
+            status_in: Optional whitelist of [`WorkerStatus`][agentlightning.WorkerStatus] values.
+            worker_id_contains: Optional substring match for worker identifiers.
+            filter_logic: Logical operator to combine the optional filters above.
+            sort_by: Field to sort by. Must be a numeric or string field of [`Worker`][agentlightning.Worker].
+            sort_order: Order to sort by.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
+
        Returns:
-            A list of all workers.
+            Sequence of Workers. Returns an empty sequence when none exist.
+            The return value is not guaranteed to be a list.
        """
        raise NotImplementedError()

@@ -0,0 +1,18 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import Collection, FilterOptions, KeyValue, LightningCollections, PaginatedResult, Queue, SortOptions
+from .memory import DequeBasedQueue, DictBasedKeyValue, InMemoryLightningCollections, ListBasedCollection
+
+__all__ = [
+    "Collection",
+    "Queue",
+    "KeyValue",
+    "FilterOptions",
+    "SortOptions",
+    "PaginatedResult",
+    "LightningCollections",
+    "ListBasedCollection",
+    "DequeBasedQueue",
+    "DictBasedKeyValue",
+    "InMemoryLightningCollections",
+]
@@ -0,0 +1,356 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    cast,
+)
+
+if TYPE_CHECKING:
+    from typing import Self
+
+from agentlightning.types import (
+    Attempt,
+    FilterField,
+    FilterOptions,
+    PaginatedResult,
+    ResourcesUpdate,
+    Rollout,
+    SortOptions,
+    Span,
+    Worker,
+)
+
+T = TypeVar("T")  # Recommended to be a BaseModel
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class Collection(Generic[T]):
+    """Behaves like a list of items. Supporting addition, updating, and deletion of items."""
+
+    def primary_keys(self) -> Sequence[str]:
+        """Get the primary keys of the collection."""
+        raise NotImplementedError()
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}]>"
+
+    def item_type(self) -> Type[T]:
+        """Get the type of the items in the collection."""
+        raise NotImplementedError()
+
+    async def size(self) -> int:
+        """Get the number of items in the collection."""
+        raise NotImplementedError()
+
+    async def query(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+        limit: int = -1,
+        offset: int = 0,
+    ) -> PaginatedResult[T]:
+        """Query the collection with the given filters, sort order, and pagination.
+
+        Args:
+            filter:
+                The filters to apply to the collection. See [`FilterOptions`][agentlightning.FilterOptions].
+
+            sort:
+                The options for sorting the collection. See [`SortOptions`][agentlightning.SortOptions].
+                The field must exist in the model. If field might contain null values, in which case the behavior is undefined
+                (i.e., depending on the implementation).
+
+            limit:
+                Max number of items to return. Use -1 for "no limit".
+
+            offset:
+                Number of items to skip from the start of the *matching* items.
+
+        Returns:
+            PaginatedResult with items, limit, offset, and total matched items.
+        """
+        raise NotImplementedError()
+
+    async def get(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+    ) -> Optional[T]:
+        """Get the first item that matches the given filters.
+
+        Args:
+            filter: The filters to apply to the collection.
+                See [`FilterOptions`][agentlightning.store.collection.FilterOptions].
+            sort: Sort options. See [`SortOptions`][agentlightning.store.collection.SortOptions].
+
+        Returns:
+            The first item that matches the given filters, or None if no item matches.
+        """
+        raise NotImplementedError()
+
+    async def insert(self, items: Sequence[T]) -> None:
+        """Add the given items to the collection.
+
+        Raises:
+            ValueError: If an item with the same primary key already exists.
+        """
+        raise NotImplementedError()
+
+    async def update(self, items: Sequence[T]) -> None:
+        """Update the given items in the collection.
+
+        Raises:
+            ValueError: If an item with the primary keys does not exist.
+        """
+        raise NotImplementedError()
+
+    async def upsert(self, items: Sequence[T]) -> None:
+        """Upsert the given items into the collection.
+
+        If the items with the same primary keys already exist, they will be updated.
+        Otherwise, they will be inserted.
+        """
+        raise NotImplementedError()
+
+    async def delete(self, items: Sequence[T]) -> None:
+        """Delete the given items from the collection.
+
+        Args:
+            items: The items to delete from the collection.
+
+        Raises:
+            ValueError: If the items with the primary keys to be deleted do not exist.
+        """
+        raise NotImplementedError()
+
+
+class Queue(Generic[T]):
+    """Behaves like a deque. Supporting appending items to the end and popping items from the front."""
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}]>"
+
+    def item_type(self) -> Type[T]:
+        """Get the type of the items in the queue."""
+        raise NotImplementedError()
+
+    async def has(self, item: T) -> bool:
+        """Check if the given item is in the queue."""
+        raise NotImplementedError()
+
+    async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
+        """Append the given items to the end of the queue.
+
+        Args:
+            items: The items to append to the end of the queue.
+
+        Returns:
+            The items that were appended to the end of the queue.
+        """
+        raise NotImplementedError()
+
+    async def dequeue(self, limit: int = 1) -> Sequence[T]:
+        """Pop the given number of items from the front of the queue.
+
+        Args:
+            limit: The number of items to pop from the front of the queue.
+
+        Returns:
+            The items that were popped from the front of the queue.
+            If there are less than `limit` items in the queue, the remaining items will be returned.
+        """
+        raise NotImplementedError()
+
+    async def peek(self, limit: int = 1) -> Sequence[T]:
+        """Peek the given number of items from the front of the queue.
+
+        Args:
+            limit: The number of items to peek from the front of the queue.
+
+        Returns:
+            The items that were peeked from the front of the queue.
+            If there are less than `limit` items in the queue, the remaining items will be returned.
+        """
+        raise NotImplementedError()
+
+    async def size(self) -> int:
+        """Get the number of items in the queue."""
+        raise NotImplementedError()
+
+
+class KeyValue(Generic[K, V]):
+    """Behaves like a dictionary. Supporting addition, updating, and deletion of items."""
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}>"
+
+    async def has(self, key: K) -> bool:
+        """Check if the given key is in the dictionary."""
+        raise NotImplementedError()
+
+    async def get(self, key: K, default: V | None = None) -> V | None:
+        """Get the value for the given key, or the default value if the key is not found."""
+        raise NotImplementedError()
+
+    async def set(self, key: K, value: V) -> None:
+        """Set the value for the given key."""
+        raise NotImplementedError()
+
+    async def pop(self, key: K, default: V | None = None) -> V | None:
+        """Pop the value for the given key, or the default value if the key is not found."""
+        raise NotImplementedError()
+
+    async def size(self) -> int:
+        """Get the number of items in the dictionary."""
+        raise NotImplementedError()
+
+
+class LightningCollections:
+    """Collections of rollouts, attempts, spans, resources, and workers.
+
+    [LightningStore][agentlightning.LightningStore] implementations can use this as a storage base
+    to implement the store API.
+    """
+
+    @property
+    def rollouts(self) -> Collection[Rollout]:
+        """Collections of rollouts."""
+        raise NotImplementedError()
+
+    @property
+    def attempts(self) -> Collection[Attempt]:
+        """Collections of attempts."""
+        raise NotImplementedError()
+
+    @property
+    def spans(self) -> Collection[Span]:
+        """Collections of spans."""
+        raise NotImplementedError()
+
+    @property
+    def resources(self) -> Collection[ResourcesUpdate]:
+        """Collections of resources."""
+        raise NotImplementedError()
+
+    @property
+    def workers(self) -> Collection[Worker]:
+        """Collections of workers."""
+        raise NotImplementedError()
+
+    @property
+    def rollout_queue(self) -> Queue[str]:
+        """Queue of rollouts (tasks)."""
+        raise NotImplementedError()
+
+    @property
+    def span_sequence_ids(self) -> KeyValue[str, int]:
+        """Dictionary (counter) of span sequence IDs."""
+        raise NotImplementedError()
+
+    def atomic(self, *args: Any, **kwargs: Any) -> AsyncContextManager[Self]:
+        """Perform a atomic operation on the collections.
+
+        Subclass may use args and kwargs to support multiple levels of atomicity.
+
+        Args:
+            *args: Arguments to pass to the operation.
+            **kwargs: Keyword arguments to pass to the operation.
+        """
+        raise NotImplementedError()
+
+    async def execute(self, callback: Callable[[Self], Awaitable[T]]) -> T:
+        """Execute the given callback within an atomic operation."""
+        async with self.atomic() as collections:
+            return await callback(collections)
+
+
+FilterMap = Mapping[str, FilterField]
+
+
+def merge_must_filters(target: MutableMapping[str, FilterField], definition: Any) -> None:
+    """Normalize a `_must` filter group into the provided mapping.
+
+    Mainly for validation purposes.
+    """
+    if definition is None:
+        return
+
+    entries: List[Mapping[str, FilterField]] = []
+    if isinstance(definition, Mapping):
+        entries.append(cast(Mapping[str, FilterField], definition))
+    elif isinstance(definition, Sequence) and not isinstance(definition, (str, bytes)):
+        for entry in definition:  # type: ignore
+            if not isinstance(entry, Mapping):
+                raise TypeError("Each `_must` entry must be a mapping of field names to operators")
+            entries.append(cast(Mapping[str, FilterField], entry))
+    else:
+        raise TypeError("`_must` filters must be provided as a mapping or sequence of mappings")
+
+    for entry in entries:
+        for field_name, ops in entry.items():
+            existing = target.get(field_name, {})
+            merged_ops: Dict[str, Any] = dict(existing)
+            for op_name, expected in ops.items():
+                if op_name in merged_ops:
+                    raise ValueError(f"Duplicate operator '{op_name}' for field '{field_name}' in must filters")
+                merged_ops[op_name] = expected
+            target[field_name] = cast(FilterField, merged_ops)
+
+
+def normalize_filter_options(
+    filter_options: Optional[FilterOptions],
+) -> Tuple[Optional[FilterMap], Optional[FilterMap], Literal["and", "or"]]:
+    """Convert FilterOptions to the internal structure and resolve aggregate logic."""
+    if not filter_options:
+        return None, None, "and"
+
+    aggregate = cast(Literal["and", "or"], filter_options.get("_aggregate", "and"))
+    if aggregate not in ("and", "or"):
+        raise ValueError(f"Unsupported filter aggregate '{aggregate}'")
+
+    # Extract normalized filters and must filters from the filter options.
+    normalized: Dict[str, FilterField] = {}
+    must_filters: Dict[str, FilterField] = {}
+    for field_name, ops in filter_options.items():
+        if field_name == "_aggregate":
+            continue
+        if field_name == "_must":
+            merge_must_filters(must_filters, ops)
+            continue
+        normalized[field_name] = cast(FilterField, dict(ops))  # type: ignore
+
+    return (normalized or None, must_filters or None, aggregate)
+
+
+def resolve_sort_options(sort: Optional[SortOptions]) -> Tuple[Optional[str], Literal["asc", "desc"]]:
+    """Extract sort field/order from the caller-provided SortOptions."""
+    if not sort:
+        return None, "asc"
+
+    sort_name = sort.get("name")
+    if not sort_name:
+        raise ValueError("Sort options must include a 'name' field")
+
+    sort_order = sort.get("order", "asc")
+    if sort_order not in ("asc", "desc"):
+        raise ValueError(f"Unsupported sort order '{sort_order}'")
+
+    return sort_name, sort_order
@@ -0,0 +1,744 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import weakref
+from collections import deque
+from contextlib import asynccontextmanager
+from typing import (
+    Any,
+    Deque,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
+
+from agentlightning.types import (
+    Attempt,
+    FilterField,
+    FilterOptions,
+    PaginatedResult,
+    ResourcesUpdate,
+    Rollout,
+    SortOptions,
+    Span,
+    Worker,
+)
+
+from .base import (
+    Collection,
+    FilterMap,
+    KeyValue,
+    LightningCollections,
+    Queue,
+    normalize_filter_options,
+    resolve_sort_options,
+)
+
+T = TypeVar("T")  # Recommended to be a BaseModel, not a dict
+K = TypeVar("K")
+V = TypeVar("V")
+
+logger = logging.getLogger(__name__)
+
+# Nested structure type:
+# dict[pk1] -> dict[pk2] -> ... -> item
+ListBasedCollectionItemType = Union[
+    Dict[Any, "ListBasedCollectionItemType[T]"],  # intermediate node
+    Dict[Any, T],  # leaf node dictionary
+]
+
+MutationMode = Literal["insert", "update", "upsert", "delete"]
+
+
+def _item_matches_filters(
+    item: object,
+    filters: Optional[FilterMap],
+    filter_logic: Literal["and", "or"],
+    must_filters: Optional[FilterMap] = None,
+) -> bool:
+    """Check whether an item matches the provided filter definition.
+
+    Filter format:
+
+    ```json
+    {
+        "_aggregate": "or",
+        "field_name": {
+            "exact": <value>,
+            "within": <iterable_of_allowed_values>,
+            "contains": <substring_or_element>,
+        },
+        ...
+    }
+    ```
+
+    Operators within the same field are stored in a unified pool and combined using
+    a universal logical operator.
+    """
+    if must_filters and not _item_matches_filters(item, must_filters, "and"):
+        return False
+
+    if not filters:
+        return True
+
+    all_conditions_match: List[bool] = []
+
+    for field_name, ops in filters.items():
+        item_value = getattr(item, field_name, None)
+
+        for op_name, expected in ops.items():
+            # Ignore no-op filters
+            if expected is None:
+                continue
+
+            if op_name == "exact":
+                all_conditions_match.append(item_value == expected)
+
+            elif op_name == "within":
+                try:
+                    all_conditions_match.append(item_value in expected)  # type: ignore[arg-type]
+                except TypeError:
+                    all_conditions_match.append(False)
+
+            elif op_name == "contains":
+                if item_value is None:
+                    all_conditions_match.append(False)
+                elif isinstance(item_value, str) and isinstance(expected, str):
+                    all_conditions_match.append(expected in item_value)
+                else:
+                    # Fallback: treat as generic iterable containment.
+                    try:
+                        all_conditions_match.append(expected in item_value)  # type: ignore[arg-type]
+                    except TypeError:
+                        all_conditions_match.append(False)
+            else:
+                raise ValueError(f"Unsupported filter operator '{op_name}' for field '{field_name}'")
+
+    return all(all_conditions_match) if filter_logic == "and" else any(all_conditions_match)
+
+
+def _get_sort_value(item: object, sort_by: str) -> Any:
+    """Get a sort key for the given item/field.
+
+    - If the field name ends with '_time', values are treated as comparable timestamps.
+    - For other fields we try to infer a safe default from the Pydantic model annotation.
+    """
+    value = getattr(item, sort_by, None)
+
+    if sort_by.endswith("_time"):
+        # For *_time fields, push missing values to the end.
+        return float("inf") if value is None else value
+
+    if value is None:
+        # Introspect model field type to choose a reasonable default for None.
+        model_fields = getattr(item.__class__, "model_fields", {})
+        if sort_by not in model_fields:
+            raise ValueError(
+                f"Failed to sort items by '{sort_by}': field does not exist " f"on {item.__class__.__name__}"
+            )
+
+        field_type_str = str(model_fields[sort_by].annotation)
+        if "str" in field_type_str or "Literal" in field_type_str:
+            return ""
+        if "int" in field_type_str:
+            return 0
+        if "float" in field_type_str:
+            return 0.0
+        raise ValueError(f"Failed to sort items by '{sort_by}': unsupported field type {field_type_str!r}")
+
+    return value
+
+
+class ListBasedCollection(Collection[T]):
+    """In-memory implementation of Collection using a nested dict for O(1) primary-key lookup.
+
+    The internal structure is:
+
+        {
+            pk1_value: {
+                pk2_value: {
+                    ...
+                        pkN_value: item
+                }
+            }
+        }
+
+    where the nesting depth equals the number of primary keys.
+
+    Sorting behavior:
+
+    1. If no sort_by is provided, the items are returned in the order of insertion.
+    2. If sort_by is provided, the items are sorted by the value of the sort_by field.
+    3. If the sort_by field is a timestamp, the null values are treated as infinity.
+    4. If the sort_by field is not a timestamp, the null values are treated as empty string
+       if the field is str-like, 0 if the field is int-like, 0.0 if the field is float-like.
+    """
+
+    def __init__(self, items: List[T], item_type: Type[T], primary_keys: Sequence[str]):
+        if not primary_keys:
+            raise ValueError("primary_keys must be non-empty")
+
+        self._items: Dict[Any, Any] = {}
+        self._size: int = 0
+        if issubclass(item_type, dict):
+            raise TypeError(f"Expect item to be not a dict, got {item_type.__name__}")
+        self._item_type: Type[T] = item_type
+        self._primary_keys: Tuple[str, ...] = tuple(primary_keys)
+
+        # Pre-populate the collection with the given items.
+        for item in items or []:
+            self._mutate_single(item, mode="insert")
+
+    def primary_keys(self) -> Sequence[str]:
+        """Return the primary key field names for this collection."""
+        return self._primary_keys
+
+    def item_type(self) -> Type[T]:
+        """Return the Pydantic model type of items stored in this collection."""
+        return self._item_type
+
+    async def size(self) -> int:
+        """Return the number of items stored in the collection."""
+        return self._size
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self._size})>"
+
+    # -------------------------------------------------------------------------
+    # Internal helpers
+    # -------------------------------------------------------------------------
+
+    def _ensure_item_type(self, item: T) -> None:
+        """Validate that the item matches the declared item_type."""
+        if not isinstance(item, self._item_type):
+            raise TypeError(f"Expected item of type {self._item_type.__name__}, " f"got {type(item).__name__}")
+
+    def _extract_primary_key_values(self, item: T) -> Tuple[Any, ...]:
+        """Extract the primary key values from an item.
+
+        Raises:
+            ValueError: If any primary key is missing on the item.
+        """
+        values: List[Any] = []
+        for key in self._primary_keys:
+            if not hasattr(item, key):
+                raise ValueError(f"Item {item} does not have primary key field '{key}'")
+            values.append(getattr(item, key))
+        return tuple(values)
+
+    def _render_key_values(self, key_values: Sequence[Any]) -> str:
+        return ", ".join(f"{name}={value!r}" for name, value in zip(self._primary_keys, key_values))
+
+    def _locate_node(
+        self,
+        key_values: Sequence[Any],
+        create_missing: bool,
+    ) -> Tuple[MutableMapping[Any, Any], Any]:
+        """Locate the parent mapping and final key for an item path.
+
+        Args:
+            key_values: The sequence of primary key values.
+            create_missing: Whether to create intermediate dictionaries as needed.
+
+        Returns:
+            (parent_mapping, final_key)
+
+        Raises:
+            KeyError: If the path does not exist and create_missing is False.
+            ValueError: If the internal structure is corrupted (non-dict where dict is expected).
+        """
+        if not key_values:
+            raise ValueError("key_values must be non-empty")
+
+        current: MutableMapping[Any, Any] = self._items
+        for idx, value in enumerate(key_values):
+            is_last = idx == len(key_values) - 1
+            if is_last:
+                # At the final level, current[value] is the item (or will be).
+                return current, value  # type: ignore
+
+            # Intermediate level: current[value] must be a dict.
+            if value not in current:
+                if not create_missing:
+                    raise KeyError(f"Path does not exist for given primary keys: {self._render_key_values(key_values)}")
+                current[value] = {}
+            next_node = current[value]  # type: ignore
+            if not isinstance(next_node, dict):
+                raise ValueError(f"Internal structure corrupted: expected dict, got {type(next_node)!r}")  # type: ignore
+            current = next_node  # type: ignore
+
+        # We should always return inside the loop.
+        raise RuntimeError("Unreachable")
+
+    def _mutate_single(self, item: T, mode: MutationMode) -> None:
+        """Core mutation logic shared by insert, update, upsert, and delete."""
+        self._ensure_item_type(item)
+        key_values = self._extract_primary_key_values(item)
+
+        if mode in ("insert", "upsert"):
+            parent, final_key = self._locate_node(key_values, create_missing=True)
+            exists = final_key in parent
+
+            if mode == "insert":
+                if exists:
+                    raise ValueError(f"Item already exists with primary key(s): {self._render_key_values(key_values)}")
+                parent[final_key] = item
+                self._size += 1
+            else:  # upsert
+                if not exists:
+                    self._size += 1
+                parent[final_key] = item
+
+        elif mode in ("update", "delete"):
+            # For update/delete we must not create missing paths.
+            try:
+                parent, final_key = self._locate_node(key_values, create_missing=False)
+            except KeyError:
+                raise ValueError(
+                    f"Item does not exist with primary key(s): {self._render_key_values(key_values)}"
+                ) from None
+
+            if final_key not in parent:
+                raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
+
+            if mode == "update":
+                parent[final_key] = item
+            else:  # delete
+                del parent[final_key]
+                self._size -= 1
+        else:
+            raise ValueError(f"Unknown mutation mode: {mode}")
+
+    def _iter_items(
+        self,
+        root: Optional[Mapping[Any, Any]] = None,
+        filters: Optional[FilterMap] = None,
+        must_filters: Optional[FilterMap] = None,
+        filter_logic: Literal["and", "or"] = "and",
+    ) -> Iterable[T]:
+        """Iterate over all items in the nested dictionary structure, optionally applying filters."""
+        if root is None:
+            root = self._items
+        if not root:
+            return
+        stack: List[Mapping[Any, Any]] = [root]
+        while stack:
+            node = stack.pop()
+            for value in node.values():
+                # Leaf nodes contain items; intermediate nodes are dicts.
+                if isinstance(value, self._item_type):
+                    if _item_matches_filters(value, filters, filter_logic, must_filters):
+                        yield value
+                elif isinstance(value, dict):
+                    stack.append(value)  # type: ignore
+                else:
+                    raise ValueError(
+                        f"Internal structure corrupted: expected dict or {self._item_type.__name__}, "
+                        f"got {type(value)!r}"
+                    )
+
+    def _iter_matching_items(
+        self,
+        filters: Optional[FilterMap],
+        must_filters: Optional[FilterMap],
+        filter_logic: Literal["and", "or"],
+    ) -> Iterable[T]:
+        """Efficiently iterate over items matching filters, using primary-key prefix when possible."""
+        # Fast path: when optional filters can't form a prefix, fall back to scanning.
+        if filter_logic != "and" and must_filters is None:
+            return self._iter_items(filters=filters, must_filters=must_filters, filter_logic=filter_logic)
+
+        # Try to derive a primary-key prefix from exact filters.
+        pk_values_prefix: List[Any] = []
+        prefix_sources: List[FilterMap] = []
+        if must_filters:
+            prefix_sources.append(must_filters)
+        if filter_logic == "and" and filters:
+            prefix_sources.append(filters)
+
+        for pk in self._primary_keys:
+            # combined_ops are: [{"exact": value}, {"within": [...]}, ...]
+            combined_ops: List[FilterField] = []
+            for source in prefix_sources:
+                field_ops = source.get(pk)  # type: ignore[union-attr]
+                if field_ops:
+                    combined_ops.append(field_ops)
+            if not combined_ops:
+                break
+            # Only allow a pure {"exact": value} constraint.
+            exact_value: Any | None = None
+            allow_prefix = True
+            for ops in combined_ops:
+                if set(ops.keys()) != {"exact"}:
+                    allow_prefix = False
+                    break
+                candidate = ops.get("exact")
+                if candidate is None:
+                    allow_prefix = False
+                    break
+                if exact_value is not None and candidate != exact_value:
+                    # Contradictory exact filters mean no items can match.
+                    logger.warning(f"Contradictory exact filters for field '{pk}': {exact_value} != {candidate}")
+                    return ()
+                exact_value = candidate
+
+            if not allow_prefix:
+                break
+
+            value = exact_value
+            if value is None:
+                break
+            pk_values_prefix.append(value)
+
+        if not pk_values_prefix:
+            return self._iter_items(filters=filters, must_filters=must_filters, filter_logic=filter_logic)
+
+        try:
+            if len(pk_values_prefix) == len(self._primary_keys):
+                # All primary keys specified -> at most a single item.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                single_item = parent.get(final_key)
+                if isinstance(single_item, self._item_type) and _item_matches_filters(
+                    single_item,
+                    filters,
+                    filter_logic,
+                    must_filters,
+                ):
+                    return (single_item,)
+                return ()
+            else:
+                # Prefix of primary keys specified -> iterate only the subtree below that prefix.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                subtree = parent.get(final_key)
+                if isinstance(subtree, dict):
+                    return self._iter_items(
+                        subtree,  # type: ignore
+                        filters=filters,
+                        must_filters=must_filters,
+                        filter_logic=filter_logic,
+                    )
+                return ()
+        except KeyError:
+            # No items exist for this primary-key prefix.
+            return ()
+
+    async def query(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+        limit: int = -1,
+        offset: int = 0,
+    ) -> PaginatedResult[T]:
+        """Query the collection with filters, sort order, and pagination.
+
+        Args:
+            filter: Mapping of field name to operator dict along with the optional `_aggregate` logic.
+            sort: Options describing which field to sort by and in which order.
+            limit: Max number of items to return. Use -1 for "no limit".
+            offset: Number of items to skip from the start of the *matching* items.
+        """
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)
+
+        # No sorting: stream through items and apply pagination on the fly.
+        if not sort_by:
+            matched_items: List[T] = []
+            total_matched = 0
+
+            for item in items_iter:
+                # Count every match for 'total'
+                total_matched += 1
+
+                # Apply offset/limit window
+                if total_matched <= offset:
+                    continue
+                if limit != -1 and len(matched_items) >= limit:
+                    # Still need to finish iteration to get accurate total_matched.
+                    continue
+
+                matched_items.append(item)
+
+            return PaginatedResult(
+                items=matched_items,
+                limit=limit,
+                offset=offset,
+                total=total_matched,
+            )
+
+        # With sorting: we must materialize all matching items to sort them.
+        all_matches: List[T] = list(items_iter)
+
+        total_matched = len(all_matches)
+        reverse = sort_order == "desc"
+        all_matches.sort(key=lambda x: _get_sort_value(x, sort_by), reverse=reverse)
+
+        if limit == -1:
+            paginated_items = all_matches[offset:]
+        else:
+            paginated_items = all_matches[offset : offset + limit]
+
+        return PaginatedResult(
+            items=paginated_items,
+            limit=limit,
+            offset=offset,
+            total=total_matched,
+        )
+
+    async def get(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+    ) -> Optional[T]:
+        """Return the first (or best-sorted) item that matches the given filters, or None."""
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)
+
+        if not sort_by:
+            # Just return the first matching item, if any.
+            for item in items_iter:
+                return item
+            return None
+
+        # Single-pass min/max according to sort_order.
+        best_item: Optional[T] = None
+        best_key: Any = None
+
+        for item in items_iter:
+            key = _get_sort_value(item, sort_by)
+            if best_item is None:
+                best_item = item
+                best_key = key
+                continue
+
+            if sort_order == "asc":
+                if key < best_key:
+                    best_item, best_key = item, key
+            else:
+                if key > best_key:
+                    best_item, best_key = item, key
+
+        return best_item
+
+    async def insert(self, items: Sequence[T]) -> None:
+        """Insert the given items.
+
+        Raises:
+            ValueError: If any item with the same primary keys already exists.
+        """
+        for item in items:
+            self._mutate_single(item, mode="insert")
+
+    async def update(self, items: Sequence[T]) -> None:
+        """Update the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        for item in items:
+            self._mutate_single(item, mode="update")
+
+    async def upsert(self, items: Sequence[T]) -> None:
+        """Upsert the given items (insert if missing, otherwise update)."""
+        for item in items:
+            self._mutate_single(item, mode="upsert")
+
+    async def delete(self, items: Sequence[T]) -> None:
+        """Delete the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        # We use a two-phase approach to avoid partial deletion if one fails:
+        # first compute key_values to validate, then perform deletions.
+        for item in items:
+            # _mutate_single will validate existence and update size.
+            self._mutate_single(item, mode="delete")
+
+
+class DequeBasedQueue(Queue[T]):
+    """Queue implementation backed by collections.deque.
+
+    Provides O(1) amortized enqueue (append) and dequeue (popleft).
+    """
+
+    def __init__(self, item_type: Type[T], items: Optional[Sequence[T]] = None):
+        self._items: Deque[T] = deque()
+        self._item_type: Type[T] = item_type
+        if items:
+            self._items.extend(items)
+
+    def item_type(self) -> Type[T]:
+        return self._item_type
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({len(self._items)})>"
+
+    async def has(self, item: T) -> bool:
+        if not isinstance(item, self._item_type):
+            raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
+        return item in self._items
+
+    async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
+        for item in items:
+            if not isinstance(item, self._item_type):
+                raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
+            self._items.append(item)
+        return items
+
+    async def dequeue(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        out: List[T] = []
+        for _ in range(min(limit, len(self._items))):
+            out.append(self._items.popleft())
+        return out
+
+    async def peek(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        result: List[T] = []
+        count = min(limit, len(self._items))
+        for idx, item in enumerate(self._items):
+            if idx >= count:
+                break
+            result.append(item)
+        return result
+
+    async def size(self) -> int:
+        return len(self._items)
+
+
+class DictBasedKeyValue(KeyValue[K, V]):
+    """KeyValue implementation backed by a plain dictionary."""
+
+    def __init__(self, data: Optional[Mapping[K, V]] = None):
+        self._values: Dict[K, V] = dict(data) if data else {}
+
+    async def has(self, key: K) -> bool:
+        return key in self._values
+
+    async def get(self, key: K, default: V | None = None) -> V | None:
+        return self._values.get(key, default)
+
+    async def set(self, key: K, value: V) -> None:
+        self._values[key] = value
+
+    async def pop(self, key: K, default: V | None = None) -> V | None:
+        return self._values.pop(key, default)
+
+    async def size(self) -> int:
+        return len(self._values)
+
+
+class InMemoryLightningCollections(LightningCollections):
+    """In-memory implementation of LightningCollections using Python data structures.
+
+    Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
+    """
+
+    def __init__(self):
+        self._lock = _LoopAwareAsyncLock()
+        self._rollouts = ListBasedCollection(items=[], item_type=Rollout, primary_keys=["rollout_id"])
+        self._attempts = ListBasedCollection(items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"])
+        self._spans = ListBasedCollection(
+            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"]
+        )
+        self._resources = ListBasedCollection(items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"])
+        self._workers = ListBasedCollection(items=[], item_type=Worker, primary_keys=["worker_id"])
+        self._rollout_queue = DequeBasedQueue(items=[], item_type=str)
+        self._span_sequence_ids = DictBasedKeyValue[str, int](data={})  # rollout_id -> sequence_id
+
+    @property
+    def rollouts(self) -> ListBasedCollection[Rollout]:
+        return self._rollouts
+
+    @property
+    def attempts(self) -> ListBasedCollection[Attempt]:
+        return self._attempts
+
+    @property
+    def spans(self) -> ListBasedCollection[Span]:
+        return self._spans
+
+    @property
+    def resources(self) -> ListBasedCollection[ResourcesUpdate]:
+        return self._resources
+
+    @property
+    def workers(self) -> ListBasedCollection[Worker]:
+        return self._workers
+
+    @property
+    def rollout_queue(self) -> DequeBasedQueue[str]:
+        return self._rollout_queue
+
+    @property
+    def span_sequence_ids(self) -> DictBasedKeyValue[str, int]:
+        return self._span_sequence_ids
+
+    @asynccontextmanager
+    async def atomic(self, *args: Any, **kwargs: Any):
+        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside."""
+        async with self._lock:
+            yield self
+
+    async def evict_spans_for_rollout(self, rollout_id: str) -> None:
+        """Evict all spans for a given rollout ID.
+
+        Uses private API for efficiency.
+        """
+        self._spans._items.pop(rollout_id, [])  # pyright: ignore[reportPrivateUsage]
+
+
+class _LoopAwareAsyncLock:
+    """Async lock that transparently rebinds to the current event loop.
+
+    The lock intentionally remains *thread-unsafe*: callers must only use it from
+    one thread at a time. If multiple threads interact with the store, each
+    thread gets its own event loop specific lock.
+    """
+
+    def __init__(self) -> None:
+        self._locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = weakref.WeakKeyDictionary()
+
+    # When serializing and deserializing, we don't need to serialize the locks.
+    # Because another process will have its own set of event loops and its own lock.
+    def __getstate__(self) -> dict[str, Any]:
+        return {}
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self._locks = weakref.WeakKeyDictionary()
+
+    def _get_lock_for_current_loop(self) -> asyncio.Lock:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._locks[loop] = lock
+        return lock
+
+    async def __aenter__(self) -> asyncio.Lock:
+        lock = self._get_lock_for_current_loop()
+        await lock.acquire()
+        return lock
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> None:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None or not lock.locked():
+            raise RuntimeError("Lock released without being acquired")
+        lock.release()
@@ -0,0 +1,82 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import uuid
+from typing import (
+    Any,
+    Callable,
+    Mapping,
+    TypeVar,
+)
+
+from pymongo import AsyncMongoClient
+
+from .base import LightningStoreCapabilities
+from .collection.mongo import MongoClientPool, MongoLightningCollections
+from .collection_based import CollectionBasedLightningStore
+
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+
+def _generate_partition_id() -> str:
+    return "pt-" + hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+
+
+class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollections]):
+    """
+    MongoDB implementation of LightningStore using MongoDB collections.
+    Data is persistent and can be shared between multiple processes.
+
+    Args:
+        client: The MongoDB client. Could be a string URI or an instance of AsyncMongoClient.
+        database: The MongoDB database. Could be a string name or an instance of AsyncDatabase.
+            You must provide at least one of client or database.
+        partition_id: The partition id. Useful when sharing the database among multiple Agent-lightning trainers.
+    """
+
+    def __init__(
+        self,
+        *,
+        client: AsyncMongoClient[Mapping[str, Any]] | str,
+        database_name: str | None = None,
+        partition_id: str | None = None,
+    ) -> None:
+        self._auto_created_client = False
+        if isinstance(client, str):
+            self._client = AsyncMongoClient[Mapping[str, Any]](client)
+            self._auto_created_client = True
+        else:
+            self._client = client
+        if database_name is None:
+            database_name = "agentlightning"
+            logger.info("No database name provided, using default 'agentlightning'")
+
+        if partition_id is None:
+            partition_id = _generate_partition_id()
+            logger.info("No partition id provided, generated a new one: %s", partition_id)
+
+        self._client_pool = MongoClientPool(self._client)
+
+        super().__init__(collections=MongoLightningCollections(self._client_pool, database_name, partition_id))
+
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        return LightningStoreCapabilities(
+            thread_safe=True,
+            async_safe=True,
+            zero_copy=True,
+            otlp_traces=False,
+        )
+
+    async def close(self) -> None:
+        """Close the store by closing the client pool."""
+        await self._client_pool.close()
+        # If I created the client, I should close it too.
+        if self._auto_created_client:
+            await self._client.close()
@@ -19,6 +19,7 @@ from agentlightning.types import (
    Span,
    TaskInput,
    Worker,
+    WorkerStatus,
 )

 from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
@@ -36,9 +37,10 @@ class LightningStoreThreaded(LightningStore):
        self.store = store
        self._lock = threading.Lock()

+    @property
    def capabilities(self) -> LightningStoreCapabilities:
        """Return the capabilities of the store."""
-        capabilities = self.store.capabilities()
+        capabilities = self.store.capabilities
        return {
            **capabilities,
            "async_safe": True,
@@ -78,15 +80,48 @@ class LightningStoreThreaded(LightningStore):
    async def query_rollouts(
        self,
        *,
+        status_in: Optional[Sequence[RolloutStatus]] = None,
+        rollout_id_in: Optional[Sequence[str]] = None,
+        rollout_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
        status: Optional[Sequence[RolloutStatus]] = None,
        rollout_ids: Optional[Sequence[str]] = None,
-    ) -> List[Rollout]:
+    ) -> Sequence[Rollout]:
        with self._lock:
-            return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
+            return await self.store.query_rollouts(
+                status_in=status_in,
+                rollout_id_in=rollout_id_in,
+                rollout_id_contains=rollout_id_contains,
+                filter_logic=filter_logic,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+                status=status,
+                rollout_ids=rollout_ids,
+            )

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

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

+    async def query_resources(
+        self,
+        *,
+        resources_id: Optional[str] = None,
+        resources_id_contains: Optional[str] = None,
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.query_resources(
+                resources_id=resources_id,
+                resources_id_contains=resources_id_contains,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )
+
    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
        with self._lock:
            return await self.store.add_resources(resources)
@@ -138,9 +193,39 @@ class LightningStoreThreaded(LightningStore):
        self,
        rollout_id: str,
        attempt_id: str | Literal["latest"] | None = None,
-    ) -> List[Span]:
+        *,
+        trace_id: Optional[str] = None,
+        trace_id_contains: Optional[str] = None,
+        span_id: Optional[str] = None,
+        span_id_contains: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        parent_id_contains: Optional[str] = None,
+        name: Optional[str] = None,
+        name_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        limit: int = -1,
+        offset: int = 0,
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+    ) -> Sequence[Span]:
        with self._lock:
-            return await self.store.query_spans(rollout_id, attempt_id)
+            return await self.store.query_spans(
+                rollout_id,
+                attempt_id,
+                trace_id=trace_id,
+                trace_id_contains=trace_id_contains,
+                span_id=span_id,
+                span_id_contains=span_id_contains,
+                parent_id=parent_id,
+                parent_id_contains=parent_id_contains,
+                name=name,
+                name_contains=name_contains,
+                filter_logic=filter_logic,
+                limit=limit,
+                offset=offset,
+                sort_by=sort_by,
+                sort_order=sort_order,
+            )

    async def update_rollout(
        self,
@@ -182,9 +267,26 @@ class LightningStoreThreaded(LightningStore):
                metadata=metadata,
            )

-    async def query_workers(self) -> List[Worker]:
+    async def query_workers(
+        self,
+        *,
+        status_in: Optional[Sequence[WorkerStatus]] = None,
+        worker_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Worker]:
        with self._lock:
-            return await self.store.query_workers()
+            return await self.store.query_workers(
+                status_in=status_in,
+                worker_id_contains=worker_id_contains,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )

    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
        with self._lock:
@@ -2,25 +2,24 @@

 from __future__ import annotations

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

 import agentops
 import agentops.sdk.core
+import opentelemetry.trace as trace_api
 from agentops.sdk.core import TracingCore
-from agentops.sdk.processors import SpanProcessor
-from opentelemetry.instrumentation.utils import suppress_instrumentation
-from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.trace 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 .base import Tracer
+from .otel import LightningSpanProcessor, OtelTracer

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


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

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

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

        if self.instrument_managed:
            self.instrument(worker_id)
@@ -85,16 +83,9 @@ class AgentOpsTracer(Tracer):

        self._lightning_span_processor = LightningSpanProcessor()

-        try:
-            # new versions
-            instance = agentops.sdk.core.tracer
-            # TODO: The span processor cannot be deleted once added.
-            # This might be a problem if the tracer is entered and exited multiple times.
-            instance.provider.add_span_processor(self._lightning_span_processor)  # type: ignore
-        except AttributeError:
-            # old versions
-            instance = TracingCore.get_instance()  # type: ignore
-            instance._provider.add_span_processor(self._lightning_span_processor)  # type: ignore
+        # 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

    def teardown_worker(self, worker_id: int) -> None:
        super().teardown_worker(worker_id)
@@ -111,7 +102,7 @@ class AgentOpsTracer(Tracer):
        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
-    ) -> AsyncGenerator[LightningSpanProcessor, None]:
+    ) -> AsyncGenerator[trace_api.Tracer, None]:
        """
        Starts a new tracing context. This should be used as a context manager.

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

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

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

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

-    def get_last_trace(self) -> List[ReadableSpan]:
-        """
-        Retrieves the raw list of captured spans from the most recent trace.
-
-        Returns:
-            A list of OpenTelemetry `ReadableSpan` objects.
-        """
-        if not self._lightning_span_processor:
-            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
-        return self._lightning_span_processor.spans()
-
    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
        """
        Get the Langchain callback handler for integrating with Langchain.
@@ -204,135 +206,26 @@ class AgentOpsTracer(Tracer):

    get_langchain_callback_handler = get_langchain_handler  # alias

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

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

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

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

 import logging
+from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional

 from opentelemetry.sdk.trace import ReadableSpan
@@ -51,6 +52,18 @@ class Tracer(ParallelWorkerBase):
    ```
    """

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

-        If a store is provided, the spans will be added to the store when tracing.
-
        Args:
            name: The name for the root span of this trace context.
-            store: The store to add the spans to.
+            store: The store to add the spans to. Deprecated in favor of passing store to init_worker().
            rollout_id: The rollout ID to add the spans to.
            attempt_id: The attempt ID to add the spans to.
        """
@@ -81,7 +92,6 @@ class Tracer(ParallelWorkerBase):
        self,
        name: Optional[str] = None,
        *,
-        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
    ) -> ContextManager[Any]:
@@ -138,3 +148,30 @@ class Tracer(ParallelWorkerBase):
        """
        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
        return None
+
+    @contextmanager
+    def lifespan(self, store: Optional[LightningStore] = None):
+        """A context manager to manage the lifespan of the tracer.
+
+        This can be used to set up and tear down any necessary resources
+        for the tracer, useful for debugging purposes.
+
+        Args:
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
+        """
+        has_init = False
+        has_init_worker = False
+        try:
+            self.init()
+            has_init = True
+
+            self.init_worker(0, store)
+            has_init_worker = True
+
+            yield
+
+        finally:
+            if has_init_worker:
+                self.teardown_worker(0)
+            if has_init:
+                self.teardown()
@@ -19,6 +19,8 @@ from opentelemetry.trace.span import (
    TraceState,
 )

+from agentlightning.store import LightningStore
+
 from .base import Tracer

 logger = logging.getLogger(__name__)
@@ -68,14 +70,15 @@ class HttpTracer(Tracer):
        self.subprocess_mode = subprocess_mode
        self.subprocess_timeout = subprocess_timeout

-    def init_worker(self, worker_id: int) -> None:
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
        """
        Initialize the tracer in a worker process.

        Args:
            worker_id: The ID of the worker process.
+            store: The store to add the spans to.
        """
-        super().init_worker(worker_id)
+        super().init_worker(worker_id, store)
        logger.info(f"[Worker {worker_id}] HttpTracer initialized.")

    @asynccontextmanager
@@ -2,16 +2,26 @@

 from __future__ import annotations

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

 import opentelemetry.trace as trace_api
-from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
+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 agentlightning.store.base import LightningStore
+from agentlightning.types.tracer import SpanNames
+from agentlightning.utils.otlp import LightningStoreOTLPExporter

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

 logger = logging.getLogger(__name__)
@@ -29,21 +39,31 @@ class OtelTracer(Tracer):
        # This provider is only initialized when the worker is initialized.
        self._tracer_provider: Optional[TracerProvider] = None
        self._lightning_span_processor: Optional[LightningSpanProcessor] = None
+        self._simple_span_processor: Optional[SimpleSpanProcessor] = None
+        self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
        self._initialized: bool = False

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

        if self._initialized:
            logger.error("Tracer provider is already initialized. OpenTelemetry may not work as expected.")

-        tracer_provider = TracerProvider()
-        trace_api.set_tracer_provider(tracer_provider)
+        self._tracer_provider = TracerProvider()
+        trace_api.set_tracer_provider(self._tracer_provider)
        self._lightning_span_processor = LightningSpanProcessor()
-        tracer_provider.add_span_processor(self._lightning_span_processor)
+        self._tracer_provider.add_span_processor(self._lightning_span_processor)
+        self._otlp_span_exporter = LightningStoreOTLPExporter()
+        self._simple_span_processor = SimpleSpanProcessor(self._otlp_span_exporter)
+        self._tracer_provider.add_span_processor(self._simple_span_processor)
        self._initialized = True

+        logger.info(f"[Worker {worker_id}] OpenTelemetry tracer provider initialized.")
+
    def teardown_worker(self, worker_id: int):
        super().teardown_worker(worker_id)
        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
@@ -57,7 +77,7 @@ class OtelTracer(Tracer):
        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
-    ) -> AsyncGenerator[LightningSpanProcessor, None]:
+    ) -> AsyncGenerator[trace_api.Tracer, None]:
        """
        Starts a new tracing context. This should be used as a context manager.

@@ -68,20 +88,37 @@ class OtelTracer(Tracer):
            attempt_id: Optional attempt ID to add the spans to.

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

-        if store is not None and rollout_id is not None and attempt_id is not None:
-            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
-            with ctx as processor:
-                yield processor
-        elif store is None and rollout_id is None and attempt_id is None:
-            with self._lightning_span_processor:
-                yield self._lightning_span_processor
+        if store is not None:
+            warnings.warn(
+                "store is deprecated in favor of init_worker(). It will be removed in the future.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
        else:
-            raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
+            store = self._store
+
+        if rollout_id is not None and attempt_id is not None:
+            if store is None:
+                raise ValueError("store is required to be initialized when rollout_id and attempt_id are provided")
+            if store.capabilities.get("otlp_traces", False) is True:
+                logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
+                self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
+            else:
+                self._disable_native_otlp_exporter()
+            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
+            with ctx:
+                yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
+        elif rollout_id is None and attempt_id is None:
+            self._disable_native_otlp_exporter()
+            with self._lightning_span_processor:
+                yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
+        else:
+            raise ValueError("rollout_id and attempt_id must be either all provided or all None")

    def get_last_trace(self) -> List[ReadableSpan]:
        """
@@ -93,3 +130,222 @@ class OtelTracer(Tracer):
        if not self._lightning_span_processor:
            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
        return self._lightning_span_processor.spans()
+
+    def _get_tracer_provider(self) -> TracerProviderImpl:
+        if self._tracer_provider is None:
+            raise RuntimeError("TracerProvider is not initialized. Call init_worker() first.")
+        return self._tracer_provider
+
+    def _enable_native_otlp_exporter(self, store: LightningStore, rollout_id: str, attempt_id: str):
+        tracer_provider = self._get_tracer_provider()
+        active_span_processor = tracer_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
+
+        # Override the resources so that the server knows where the request comes from.
+        tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
+            Resource.create(
+                {
+                    SpanNames.ROLLOUT_ID: rollout_id,
+                    SpanNames.ATTEMPT_ID: attempt_id,
+                }
+            )
+        )
+        instrumented = False
+        candidates: List[str] = []
+        for processor in active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+            if isinstance(processor, LightningSpanProcessor):
+                # We don't need the LightningSpanProcessor any more.
+                logger.debug("LightningSpanProcessor already present in TracerProvider, disabling it.")
+                processor.disable_store_submission = True
+            elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
+                # Instead, we rely on the OTLPSpanExporter to send spans to the store.
+                if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
+                    processor.span_exporter.enable_store_otlp(store.otlp_traces_endpoint(), rollout_id, attempt_id)
+                    logger.debug(f"Set LightningStoreOTLPExporter endpoint to {store.otlp_traces_endpoint()}")
+                    instrumented = True
+                else:
+                    candidates.append(
+                        f"{processor.__class__.__name__} with {processor.span_exporter.__class__.__name__}"
+                    )
+            else:
+                candidates.append(f"{processor.__class__.__name__}")
+
+        if not instrumented:
+            raise RuntimeError(
+                "Failed to enable native OTLP exporter: no BatchSpanProcessor or SimpleSpanProcessor with "
+                "LightningStoreOTLPExporter found in TracerProvider. Please try using a non-OTLP store."
+                "Candidates are: " + ", ".join(candidates)
+            )
+
+    def _disable_native_otlp_exporter(self):
+        tracer_provider = self._get_tracer_provider()
+        active_span_processor = tracer_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
+        tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
+            Resource.create(
+                {
+                    SpanNames.ROLLOUT_ID: "",
+                    SpanNames.ATTEMPT_ID: "",
+                }
+            )
+        )  # reset resource
+        for processor in active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+            if isinstance(processor, LightningSpanProcessor):
+                # We will be in need of the LightningSpanProcessor again.
+                logger.debug("Enabling LightningSpanProcessor in TracerProvider.")
+                processor.disable_store_submission = False
+
+
+class LightningSpanProcessor(SpanProcessor):
+    """Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
+    to a [`LightningStore`][agentlightning.LightningStore].
+
+    It serves two purposes:
+
+    1. Records all the spans in a local buffer.
+    2. Submits the spans to the event loop to be added to the store.
+    """
+
+    def __init__(self, disable_store_submission: bool = False):
+        self._disable_store_submission: bool = disable_store_submission
+        self._spans: List[ReadableSpan] = []
+
+        # Store related context and states
+        self._store: Optional[LightningStore] = None
+        self._rollout_id: Optional[str] = None
+        self._attempt_id: Optional[str] = None
+        self._lock = threading.Lock()
+
+        # private asyncio loop running in a daemon thread
+        self._loop_ready = threading.Event()
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[threading.Thread] = None
+
+    @property
+    def disable_store_submission(self) -> bool:
+        """Whether to disable submitting spans to the store."""
+        return self._disable_store_submission
+
+    @disable_store_submission.setter
+    def disable_store_submission(self, value: bool) -> None:
+        self._disable_store_submission = value
+
+    def _ensure_loop(self) -> None:
+        if self._loop_thread is None or self._loop is None:
+            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
+
+    def _loop_runner(self):
+        loop = asyncio.new_event_loop()
+        self._loop = loop
+        asyncio.set_event_loop(loop)
+        self._loop_ready.set()
+        loop.run_forever()
+        loop.close()
+
+    def __enter__(self):
+        self._last_trace = None
+        self._spans = []
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
+        self._store = None
+        self._rollout_id = None
+        self._attempt_id = None
+
+    def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
+        # submit to the dedicated loop and wait synchronously
+        self._ensure_loop()
+        if self._loop is None:
+            raise RuntimeError("Loop is not initialized. This should not happen.")
+
+        # If already on the exporter loop thread, schedule and return immediately.
+        # ---------------------------------------------------------------------------
+        # WHY THIS CONDITIONAL EXISTS:
+        # In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
+        # (or another finalizer) while the Python garbage collector is running on the
+        # *same thread* that owns our exporter event loop ("otel-loop").
+        #
+        # When that happens, on_end() executes on the exporter loop thread itself.
+        # If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
+        # it would deadlock immediately — because the loop cannot both wait on and run
+        # the same coroutine. The Future stays pending forever and the loop stops
+        # processing scheduled callbacks.
+        #
+        # To avoid that self-deadlock, we detect when on_end() runs on the exporter
+        # loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
+        # instead of blocking with .result().
+        #
+        # This situation can occur because Python calls __del__ in whatever thread
+        # releases the last reference, which can easily be our loop thread if the
+        # object is dereferenced during loop._run_once().
+        # ---------------------------------------------------------------------------
+        if threading.current_thread() is self._loop_thread:
+            self._loop.call_soon_threadsafe(asyncio.create_task, coro)  # type: ignore
+            return None
+
+        fut = asyncio.run_coroutine_threadsafe(coro, self._loop)  # type: ignore
+        return fut.result(timeout=timeout)  # raises on error  # type: ignore
+
+    def shutdown(self) -> None:
+        if self._loop:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._loop = None
+        if self._loop_thread:
+            self._loop_thread.join(timeout=5)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def spans(self) -> List[ReadableSpan]:
+        """
+        Get the list of spans collected by this processor.
+        This is useful for debugging and testing purposes.
+
+        Returns:
+            List of ReadableSpan objects collected during tracing.
+        """
+        return self._spans
+
+    def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
+        # simple context manager without nesting into asyncio
+        class _Ctx:
+            def __enter__(_):  # type: ignore
+                # Use _ instead of self to avoid shadowing the instance method.
+                with self._lock:
+                    self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
+                    self._last_trace = None
+                    self._spans = []
+                return self
+
+            def __exit__(_, exc_type, exc, tb):  # type: ignore
+                with self._lock:
+                    self._store = self._rollout_id = self._attempt_id = None
+
+        return _Ctx()
+
+    def on_end(self, span: ReadableSpan) -> None:
+        """
+        Process a span when it ends.
+
+        Args:
+            span: The span that has ended.
+        """
+        # Skip if span is not sampled
+        if not span.context or not span.context.trace_flags.sampled:
+            return
+
+        if not self._disable_store_submission and self._store and self._rollout_id and self._attempt_id:
+            try:
+                # Submit add_otel_span to the event loop and wait for it to complete
+                with suppress_instrumentation():
+                    self._ensure_loop()
+                    self._await_in_loop(
+                        self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
+                        timeout=60.0,
+                    )
+            except Exception:
+                # log; on_end MUST NOT raise
+                logger.exception(f"Error adding span to store: {span.name}")
+
+        self._spans.append(span)
@@ -10,14 +10,19 @@ from typing import (
    Callable,
    Dict,
    Generic,
+    Iterator,
    List,
    Literal,
+    Mapping,
    Optional,
    Protocol,
+    Sequence,
    SupportsIndex,
+    TypedDict,
    TypeVar,
    Union,
    cast,
+    overload,
 )

 from opentelemetry.sdk.trace import ReadableSpan
@@ -51,6 +56,10 @@ __all__ = [
    "Hook",
    "Worker",
    "WorkerStatus",
+    "PaginatedResult",
+    "FilterOptions",
+    "SortOptions",
+    "FilterField",
 ]

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


 class SpanAttributeNames(str, Enum):
@@ -0,0 +1,428 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import gzip
+import logging
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar
+
+from fastapi import Request, Response
+from google.protobuf import json_format
+from google.rpc.status_pb2 import Status
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.proto.collector.logs.v1.logs_service_pb2 import (
+    ExportLogsServiceRequest,
+    ExportLogsServiceResponse,
+)
+from opentelemetry.proto.collector.metrics.v1.metrics_service_pb2 import (
+    ExportMetricsServiceRequest,
+    ExportMetricsServiceResponse,
+)
+from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
+    ExportTraceServiceRequest,
+    ExportTraceServiceResponse,
+)
+from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
+from opentelemetry.proto.resource.v1.resource_pb2 import Resource as ProtoResource
+from opentelemetry.proto.trace.v1.trace_pb2 import Span as ProtoSpan
+from opentelemetry.proto.trace.v1.trace_pb2 import Status as ProtoStatus
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExportResult
+from opentelemetry.util.types import AttributeValue
+
+from agentlightning.store.base import LightningStore
+from agentlightning.types.tracer import (
+    Attributes,
+    Event,
+    Link,
+    OtelResource,
+    Span,
+    SpanContext,
+    SpanNames,
+    TraceStatus,
+    convert_timestamp,
+)
+
+PROTOBUF_CT = "application/x-protobuf"
+
+logger = logging.getLogger(__name__)
+
+
+T_request = TypeVar("T_request", ExportLogsServiceRequest, ExportMetricsServiceRequest, ExportTraceServiceRequest)
+T_response = TypeVar("T_response", ExportLogsServiceResponse, ExportMetricsServiceResponse, ExportTraceServiceResponse)
+
+
+async def handle_otlp_export(
+    request: Request,
+    request_message_cls: Type[T_request],
+    response_message_cls: Type[T_response],
+    message_callback: Optional[Callable[[T_request], Awaitable[None]]],
+    signal_name: str,
+) -> Response:
+    """
+    Generic handler for /v1/traces, /v1/metrics, /v1/logs.
+
+    Convert the OTLP Protobuf request to a JSON-like object.
+    """
+    content_type = request.headers.get("Content-Type", "").split(";")[0].strip()
+
+    if content_type != PROTOBUF_CT:
+        # For brevity we only support binary protobuf here.
+        return _bad_request_response(
+            request,
+            f"Unsupported Content-Type '{content_type}', expected '{PROTOBUF_CT}'",
+            content_type=PROTOBUF_CT,
+        )
+
+    raw_body = await request.body()
+    body = _read_body_maybe_gzip(request, raw_body)
+
+    # Empty request is allowed and should still succeed.
+    if not body:
+        req_msg = request_message_cls()
+    else:
+        req_msg = request_message_cls()
+        try:
+            req_msg.ParseFromString(body)
+        except Exception as exc:
+            return _bad_request_response(request, f"Unable to parse OTLP {signal_name} payload: {exc}")
+
+    if message_callback is not None:
+        await message_callback(req_msg)
+
+    # Build success response. Partial success field is left unset.
+    resp_msg = response_message_cls()
+
+    # Encode response in the same Content-Type as request.
+    if content_type == PROTOBUF_CT:
+        resp_bytes = resp_msg.SerializeToString()
+    else:
+        resp_bytes = json_format.MessageToJson(resp_msg).encode("utf-8")
+
+    resp_bytes, headers = _maybe_gzip_response(request, resp_bytes)
+
+    return Response(
+        content=resp_bytes,
+        media_type=content_type,
+        status_code=200,
+        headers=headers,
+    )
+
+
+async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningStore) -> List[Span]:
+    """Parse an OTLP proto payload into List[Span].
+
+    A store is needed here for generating a sequence ID for each span.
+    """
+    output_spans: List[Span] = []
+
+    for resource_spans in request.resource_spans:
+        # Resource-level attributes & IDs
+        resource_attrs = _kv_list_to_dict(resource_spans.resource.attributes)
+        # rollout_id, attempt_id from resource attributes when present.
+        rollout_id_resource = resource_attrs.get(SpanNames.ROLLOUT_ID)
+        attempt_id_resource = resource_attrs.get(SpanNames.ATTEMPT_ID)
+        # If sequence id is provided, all the spans will share the same sequence ID.
+        # unless otherwise overridden by span-level attributes.
+        sequence_id_resource = resource_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
+
+        otel_resource = _resource_from_proto(resource_spans.resource, getattr(resource_spans, "schema_url", ""))
+
+        # Each ScopeSpans contains multiple spans
+        for scope_spans in resource_spans.scope_spans:
+            for proto_span in scope_spans.spans:
+                trace_id_hex = _bytes_to_trace_id_hex(proto_span.trace_id)
+                span_id_hex = _bytes_to_span_id_hex(proto_span.span_id)
+                parent_id_hex = _bytes_to_span_id_hex(proto_span.parent_span_id) if proto_span.parent_span_id else None
+
+                # Status
+                status_code_str = _STATUS_CODE_MAP.get(proto_span.status.code, "UNSET")
+                status = TraceStatus(
+                    status_code=status_code_str,
+                    description=proto_span.status.message or None,
+                )
+
+                # Attributes
+                span_attrs = _kv_list_to_dict(proto_span.attributes)
+
+                # Context
+                context = SpanContext(
+                    trace_id=trace_id_hex,
+                    span_id=span_id_hex,
+                    is_remote=False,
+                    trace_state={},
+                )
+
+                # Try to get if span attributes contain something like rollout_id or attempt_id
+                # Override the resource-level attributes with the span-level attributes if present.
+                rollout_id_span = span_attrs.get(SpanNames.ROLLOUT_ID)
+                attempt_id_span = span_attrs.get(SpanNames.ATTEMPT_ID)
+                sequence_id_span = span_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
+
+                # Normalize to regular strings and ints
+                rollout_id_raw = rollout_id_span if rollout_id_span is not None else rollout_id_resource
+                attempt_id_raw = attempt_id_span if attempt_id_span is not None else attempt_id_resource
+                sequence_id_raw = sequence_id_span if sequence_id_span is not None else sequence_id_resource
+
+                rollout_id, attempt_id = _normalize_rollout_attempt_id(rollout_id_raw, attempt_id_raw)
+                sequence_id = _normalize_sequence_id(sequence_id_raw)
+
+                if rollout_id is None or attempt_id is None:
+                    logger.warning(
+                        "Both rollout_id and attempt_id must be present in resource attributes. "
+                        "Spans will not be able to log to the store because of missing IDs: rollout_id=%s, attempt_id=%s, sequence_id=%s",
+                        rollout_id,
+                        attempt_id,
+                        sequence_id,
+                    )
+                    continue
+
+                # Generate a new sequence ID if not provided
+                if sequence_id is None:
+                    current_sequence_id = await store.get_next_span_sequence_id(
+                        rollout_id=rollout_id, attempt_id=attempt_id
+                    )
+                else:
+                    current_sequence_id = sequence_id
+
+                # Build Span
+                span = Span(
+                    rollout_id=rollout_id,
+                    attempt_id=attempt_id,
+                    sequence_id=current_sequence_id,
+                    trace_id=trace_id_hex,
+                    span_id=span_id_hex,
+                    parent_id=parent_id_hex,
+                    name=proto_span.name,
+                    status=status,
+                    attributes=span_attrs,
+                    events=_events_from_proto(proto_span),
+                    links=_links_from_proto(proto_span),
+                    start_time=convert_timestamp(proto_span.start_time_unix_nano),
+                    end_time=convert_timestamp(proto_span.end_time_unix_nano),
+                    context=context,
+                    parent=None,  # OTLP only has parent_span_id; we don't have full SpanContext
+                    resource=otel_resource,
+                )
+
+                output_spans.append(span)
+
+    return output_spans
+
+
+class LightningStoreOTLPExporter(OTLPSpanExporter):
+    """OTLP Exporter that write to a LightningStore-compatible backend.
+
+    The backend requires two special attributes on each span:
+
+    - `agentlightning.rollout_id`: The rollout ID to associate the span with.
+    - `agentlightning.attempt_id`: The attempt ID to associate the span with.
+
+    It can optionally use the following attribute to sequence spans:
+
+    - `agentlightning.span_sequence_id`: A decimal string representing the sequence ID of the span.
+    """
+
+    _default_endpoint: Optional[str] = None
+    _rollout_id: Optional[str] = None
+    _attempt_id: Optional[str] = None
+
+    def enable_store_otlp(self, endpoint: str, rollout_id: str, attempt_id: str) -> None:
+        """Enable storing OTLP data to a specific LightningStore rollout/attempt."""
+        self._rollout_id = rollout_id
+        self._attempt_id = attempt_id
+
+        self._default_endpoint = self._endpoint
+        self._endpoint = endpoint
+
+    def disable_store_otlp(self) -> None:
+        """Disable storing OTLP data to LightningStore."""
+        self._rollout_id = None
+        self._attempt_id = None
+        if self._default_endpoint is not None:
+            self._endpoint = self._default_endpoint
+
+    def should_bypass(self) -> bool:
+        """Check if the exporter should bypass the default export if rollout_id and attempt_id are not set."""
+        return True
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        if self._rollout_id is not None and self._attempt_id is not None:
+            # rollout_id and attempt_id are present in resource attributes
+            # It means that the server supports OTLP endpoint.
+            for span in spans:
+                # Override the resources so that the server knows where the request comes from.
+                span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
+                    Resource.create(
+                        {
+                            SpanNames.ROLLOUT_ID: self._rollout_id,
+                            SpanNames.ATTEMPT_ID: self._attempt_id,
+                        }
+                    )
+                )
+            return super().export(spans)
+        elif not self.should_bypass():
+            logger.debug("Rollout ID and Attempt ID not set; using default OTLP exporter behavior.")
+            return super().export(spans)
+        else:
+            logger.debug("Rollout ID and Attempt ID not set; bypassing export.")
+            return SpanExportResult.SUCCESS
+
+
+def _read_body_maybe_gzip(request: Request, raw_body: bytes) -> bytes:
+    """
+    Decompress body if Content-Encoding: gzip; otherwise return as is.
+    """
+    encoding = request.headers.get("Content-Encoding", "").lower()
+    if encoding == "gzip":
+        return gzip.decompress(raw_body)
+    return raw_body
+
+
+def _maybe_gzip_response(request: Request, payload: bytes) -> Tuple[bytes, Dict[str, str]]:
+    """
+    If Accept-Encoding includes gzip, gzip the payload and set Content-Encoding header.
+    """
+    ae = request.headers.get("Accept-Encoding", "")
+    tokens = [token.split(";")[0].strip().lower() for token in ae.split(",") if token.strip()]
+    headers: Dict[str, str] = {}
+    if "gzip" in tokens:
+        payload = gzip.compress(payload)
+        headers["Content-Encoding"] = "gzip"
+    return payload, headers
+
+
+def _bad_request_response(request: Request, message: str, content_type: str = PROTOBUF_CT) -> Response:
+    """
+    Build a 400 response whose body is a protobuf Status message, encoded
+    in the same Content-Type as the request (OTLP/HTTP requirement).
+    """
+    status_msg = Status(message=message)
+
+    if content_type == PROTOBUF_CT:
+        body = status_msg.SerializeToString()
+    else:
+        # Fallback: JSON representation of Status.
+        body = json_format.MessageToJson(status_msg).encode("utf-8")
+
+    body, headers = _maybe_gzip_response(request, body)
+
+    return Response(
+        content=body,
+        status_code=400,
+        media_type=content_type,
+        headers=headers,
+    )
+
+
+def _normalize_rollout_attempt_id(
+    rollout_id: Optional[AttributeValue], attempt_id: Optional[AttributeValue]
+) -> Tuple[Optional[str], Optional[str]]:
+    """Normalize a rollout or attempt ID to a string."""
+    rollout_id_str = str(rollout_id) if rollout_id is not None else None
+    attempt_id_str = str(attempt_id) if attempt_id is not None else None
+    return rollout_id_str, attempt_id_str
+
+
+def _normalize_sequence_id(sequence_id: Optional[AttributeValue]) -> Optional[int]:
+    """Normalize a sequence ID to an integer."""
+    if sequence_id is None:
+        return None
+    try:
+        sequence_id_int = int(str(sequence_id))
+    except (ValueError, TypeError):
+        logger.warning(
+            "Invalid sequence_id value in resource attributes: %r. Must be an integer or string representing an integer. Assuming None.",
+            sequence_id,
+        )
+        sequence_id_int = None
+    return sequence_id_int
+
+
+def _any_value_to_python(value: AnyValue) -> Any:
+    """Convert OTLP AnyValue -> plain Python value."""
+    kind = value.WhichOneof("value")
+    if kind is None:
+        return None
+    if kind == "string_value":
+        return value.string_value
+    if kind == "bool_value":
+        return value.bool_value
+    if kind == "int_value":
+        return int(value.int_value)
+    if kind == "double_value":
+        return float(value.double_value)
+    if kind == "array_value":
+        return [_any_value_to_python(v) for v in value.array_value.values]
+    if kind == "kvlist_value":
+        # Map<string, AnyValue> -> dict
+        return {kv.key: _any_value_to_python(kv.value) for kv in value.kvlist_value.values}
+    if kind == "bytes_value":
+        # Serialize bytes as hex string to stay JSON-friendly
+        return value.bytes_value.hex()
+    return None
+
+
+def _kv_list_to_dict(kvs: Sequence[KeyValue]) -> Attributes:
+    """Convert repeated KeyValue -> Attributes dict."""
+    return {kv.key: _any_value_to_python(kv.value) for kv in kvs}
+
+
+_STATUS_CODE_MAP = {
+    ProtoStatus.STATUS_CODE_UNSET: "UNSET",
+    ProtoStatus.STATUS_CODE_OK: "OK",
+    ProtoStatus.STATUS_CODE_ERROR: "ERROR",
+}
+
+
+def _bytes_to_trace_id_hex(b: bytes) -> str:
+    # OTLP uses 16-byte trace IDs; format as 32-char hex
+    if not b:
+        return "0" * 32
+    return b.hex().rjust(32, "0")
+
+
+def _bytes_to_span_id_hex(b: bytes) -> str:
+    # OTLP uses 8-byte span IDs; format as 16-char hex
+    if not b:
+        return "0" * 16
+    return b.hex().rjust(16, "0")
+
+
+def _events_from_proto(span: ProtoSpan) -> List[Event]:
+    """Event converter from OTLP ProtoSpan to List[Event]."""
+    return [
+        Event(
+            name=e.name,
+            attributes=_kv_list_to_dict(e.attributes),
+            timestamp=convert_timestamp(e.time_unix_nano),
+        )
+        for e in span.events
+    ]
+
+
+def _links_from_proto(span: ProtoSpan) -> List[Link]:
+    """Link converter from OTLP ProtoSpan to List[Link]."""
+    links: List[Link] = []
+    for link in span.links:
+        trace_id_hex = _bytes_to_trace_id_hex(link.trace_id)
+        span_id_hex = _bytes_to_span_id_hex(link.span_id)
+        ctx = SpanContext(
+            trace_id=trace_id_hex,
+            span_id=span_id_hex,
+            is_remote=False,
+            trace_state={},  # OTLP trace_state is currently a string; you can parse if needed
+        )
+        links.append(
+            Link(
+                context=ctx,
+                attributes=_kv_list_to_dict(link.attributes) or None,
+            )
+        )
+    return links
+
+
+def _resource_from_proto(resource: ProtoResource, schema_url: str = "") -> OtelResource:
+    return OtelResource(
+        attributes=_kv_list_to_dict(resource.attributes),
+        schema_url=schema_url or "",
+    )
@@ -15,7 +15,7 @@ import traceback
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass
 from multiprocessing.process import BaseProcess
-from typing import Any, AsyncContextManager, AsyncIterator, Dict, Literal, Optional
+from typing import Any, AsyncContextManager, AsyncIterator, Dict, Literal, Optional, cast

 import aiohttp
 import requests
@@ -53,6 +53,8 @@ class PythonServerLauncherArgs:
    """
    log_level: int = logging.INFO
    """The log level to use."""
+    access_log: bool = False
+    """Whether to turn on access logs."""
    startup_timeout: float = 60.0
    """The timeout to wait for the server to start up."""
    kill_unhealthy_server: bool = True
@@ -63,6 +65,8 @@ class PythonServerLauncherArgs:
    """The timeout to wait for the thread to join."""
    process_join_timeout: float = 10.0
    """The timeout to wait for the process to join."""
+    timeout_keep_alive: int = 30
+    """The timeout to keep the connection alive."""


@dataclass
@@ -156,7 +160,9 @@ async def run_uvicorn_asyncio(

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

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

@@ -608,6 +614,13 @@ class PythonServerLauncher:
        self._host: Optional[str] = self.args.host
        self._port: Optional[int] = self.args.port
        self._access_host: Optional[str] = self.args.access_host
+        self.initialize()
+
+    def initialize(self):
+        # ensure the host/port/access_host are set
+        self._ensure_host()
+        self._ensure_port()
+        self._ensure_access_host()

        # uvicorn (in-proc asyncio)
        self._uvicorn_server: Optional[uvicorn.Server] = None
@@ -626,6 +639,26 @@ class PythonServerLauncher:
        # is_running flag
        self._is_running: bool = False

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

    def _ensure_access_host(self) -> str:
-        if self.args.access_host is None:
-            if self._ensure_host() in ("0.0.0.0", "::"):
-                # Probe host normalization for 0.0.0.0
-                logger.warning("No access host provided, using default outbound IPv4 address for this machine.")
-                self._access_host = _get_default_ipv4_address()
+        if self._access_host is None:
+            if self.args.access_host is None:
+                if self._ensure_host() in ("0.0.0.0", "::"):
+                    # Probe host normalization for 0.0.0.0
+                    logger.warning("No access host provided, using default outbound IPv4 address for this machine.")
+                    self._access_host = _get_default_ipv4_address()
+                else:
+                    logger.warning("No access host provided, using the host provided.")
+                    self._access_host = self._ensure_host()
            else:
-                logger.warning("No access host provided, using the host provided.")
-                self._access_host = self._ensure_host()
-        else:
-            self._access_host = self.args.access_host
-        return self._access_host
+                self._access_host = self.args.access_host
+        return self._access_host  # type: ignore

    def _create_uvicorn_server(self) -> uvicorn.Server:
        config = uvicorn.Config(
@@ -762,7 +796,9 @@ class PythonServerLauncher:
            host=self._ensure_host(),
            port=self._ensure_port(),
            log_level=self.args.log_level,
+            access_log=self.args.access_log,
            loop="asyncio",
+            timeout_keep_alive=self.args.timeout_keep_alive,
        )
        return uvicorn.Server(config)

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

        if evt.kind == "error":
            logger.error("Threaded server failed to start (%s): %s\n%s", evt.exc_type, evt.message, evt.traceback)
            await asyncio.to_thread(self._thread.join, self.args.thread_join_timeout)
            if self._thread.is_alive():
-                raise RuntimeError(evt.message or "Threaded server failed to start and refused to shut down.")
+                logger.error("Threaded server failed to start and refused to shut down.")
+            raise RuntimeError(evt.message)
        else:
            logger.info("Threaded server started successfully.")
            self._is_running = True
@@ -893,7 +931,7 @@ class PythonServerLauncher:
                "workers": int(self.args.n_workers),
                "worker_class": "uvicorn_worker.UvicornWorker",
                "loglevel": logging.getLevelName(self.args.log_level).lower(),
-                "accesslog": None,
+                "accesslog": "-" if self.args.access_log else None,
                "errorlog": "-",
                "preload_app": True,
                "graceful_timeout": int(
@@ -939,11 +977,12 @@ class PythonServerLauncher:
            evt: ChildEvent = await asyncio.to_thread(self._mp_event_queue.get, True, timeout)
        except queue.Empty:
            if not self._proc.is_alive():
-                logger.error("Server process failed to start and is not alive. No error event was received.")
-                return
-            logger.error("Server process failed to start and sends no event. This should not happen.")
+                raise RuntimeError("Server process failed to start and is not alive. No error event was received.")
+            logger.error(
+                "Server process failed to start and sends no event. This should not happen. Shutting down server."
+            )
            await self._stop_serving_process()
-            return
+            raise RuntimeError("Server process failed to start and sends no event. This should not happen.")

        if evt.kind == "error":
            logger.error(
@@ -955,7 +994,8 @@ class PythonServerLauncher:
            )
            await asyncio.to_thread(self._proc.join, self.args.process_join_timeout)
            if self._proc.is_alive():
-                raise RuntimeError(evt.message or "Server process failed to start and refused to shut down.")
+                logger.error("Server process failed to start and refused to shut down.")
+            raise RuntimeError(evt.message)
        else:
            logger.info("Subprocess server started successfully.")
            self._is_running = True
@@ -18,14 +18,12 @@ from flask import Flask, Response, abort, request
 from tensordict import TensorDict
 from verl import DataProto

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

-configure_logger()
-
 __all__ = [
    "AgentModeDaemon",
    "get_left_padded_ids_and_attention_mask",
@@ -559,14 +557,17 @@ class AgentModeDaemon:
        )  # FIXME: Evaluate whether grouping stats by source is actually needed.

        for rollout_id, rollout in self._completed_rollouts_v0.items():
+            final_reward_raw: Optional[float] = rollout.final_reward
            final_reward = self._fillna_reward(rollout)
            if not rollout.triplets:
                print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
                sample_stat_list.append({"reward": final_reward})
                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(
                    {
@@ -574,6 +575,7 @@ class AgentModeDaemon:
                        "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
                        "turn_count": len(rollout.triplets),
                        "reward": final_reward,
+                        "has_reward": final_reward_raw is not None,
                    }
                )
            sample_stat_list.append(
@@ -582,6 +584,7 @@ class AgentModeDaemon:
                    "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
                    "turn_count": len(rollout.triplets),
                    "reward": final_reward,
+                    "has_reward": final_reward_raw is not None,
                }
            )
        metric_dict: Dict[str, Any] = {}
@@ -596,6 +599,9 @@ class AgentModeDaemon:
                {
                    f"val/{data_source}/n_rollouts": len(sample_stats),
                    f"val/{data_source}/n_rollouts_w_trace": len(stats_w_trace_by_source[data_source]),
+                    f"val/{data_source}/n_rollouts_w_reward": len(
+                        [stat for stat in sample_stats if stat["has_reward"]]
+                    ),
                    f"val/{data_source}/reward": np.mean(
                        [stat["reward"] for stat in sample_stats]
                    ),  # each rollout must have a reward (fillna if missing)
@@ -614,6 +620,7 @@ class AgentModeDaemon:
            {
                "val/n_rollouts": len(sample_stat_list),
                "val/n_rollouts_w_trace": len(stats_w_trace),
+                "val/n_rollouts_w_reward": len([stat for stat in sample_stat_list if stat["has_reward"]]),
                "val/reward": np.mean(
                    [stat["reward"] for stat in sample_stat_list]
                ),  # each rollout must have a reward (fillna if missing)
@@ -638,9 +645,10 @@ class AgentModeDaemon:
        # 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts
        finished_id_to_sample_info: Dict[str, Dict[str, Any]] = {}
        finished_id_to_final_reward: Dict[str, float] = {}
+        sample_with_reward_count = 0
        for rollout_id, rollout in self._completed_rollouts_v0.items():
            original_sample = self._task_id_to_original_sample[rollout_id]
-
+            sample_with_reward_count += int(rollout.final_reward is not None)
            final_reward = self._fillna_reward(rollout)

            if not rollout.triplets:
@@ -759,6 +767,7 @@ class AgentModeDaemon:
            "training/reward": np.mean(list(finished_id_to_final_reward.values())),
            "training/n_rollouts": len(finished_id_to_final_reward),
            "training/n_rollouts_w_trace": len(finished_id_to_sample_info),
+            "training/n_rollouts_w_reward": sample_with_reward_count,
            "training/n_truncated_triplets": n_trunc_sample_because_of_response,
            "training/n_triplets": n_transition,
        }
@@ -12,6 +12,7 @@ from typing import Dict, Tuple

 import numpy as np
 import torch
+import verl
 from codetiming import Timer
 from omegaconf import OmegaConf
 from tqdm import tqdm
@@ -403,14 +404,20 @@ class AgentLightningTrainer(RayPPOTrainer):
        assert self.async_rollout_mode, "If agent mode is enabled, async server must be enabled"
        if self.adapter is not None and not isinstance(self.adapter, TraceToTripletBase):
            raise ValueError("Adapter must be a TraceToTripletBase for currently VERL implementation.")
+        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 = AgentModeDaemon(
            self.config.agentlightning.port,
            self.config.actor_rollout_ref.rollout.n,
            train_information={
-                # Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
-                # However, it is possible that verl updates the naming and causes incompatibility.
-                # Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
-                "model": "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:]),
+                "model": model,
                "temperature": self.config.actor_rollout_ref.rollout.temperature,
            },
            tokenizer=self.tokenizer,
@@ -1,12 +1,12 @@
 {
  "name": "agent-lightning-dashboard",
-  "version": "0.2.2",
+  "version": "0.3.0",
  "lockfileVersion": 3,
  "requires": true,
  "packages": {
    "": {
      "name": "agent-lightning-dashboard",
-      "version": "0.2.2",
+      "version": "0.3.0",
      "dependencies": {
        "@mantine/core": "8.3.5",
        "@mantine/hooks": "8.3.5",
@@ -1,7 +1,7 @@
 {
  "name": "agent-lightning-dashboard",
  "type": "module",
-  "version": "0.2.2",
+  "version": "0.3.0",
  "scripts": {
    "dev": "vite",
    "build": "tsc && vite build",
@@ -37,7 +37,7 @@ from agentlightning.types import (
 )


-def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) -> None:
+async def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) -> None:
    """
    Inject mock data directly into the InMemoryLightningStore.

@@ -217,20 +217,12 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
    )

    # Inject rollouts directly into store
-    store._rollouts["ro-story-001"] = rollout1
-    store._rollouts["ro-story-002"] = rollout2
-    store._rollouts["ro-story-003"] = rollout3
-    store._rollouts["ro-story-004"] = rollout4
-    store._rollouts["ro-story-005"] = rollout5
-    store._rollouts["ro-story-006"] = rollout6
+    await store.collections.rollouts.insert([rollout1, rollout2, rollout3, rollout4, rollout5, rollout6])

    # Inject attempts directly into store
-    store._attempts["ro-story-001"] = [attempt1]
-    store._attempts["ro-story-002"] = [attempt2_1, attempt2_2]
-    store._attempts["ro-story-003"] = [attempt3_1, attempt3_2, attempt3_3]
-    store._attempts["ro-story-004"] = []  # No attempt for preparing rollout
-    store._attempts["ro-story-005"] = [attempt5]
-    store._attempts["ro-story-006"] = [attempt6]
+    await store.collections.attempts.insert(
+        [attempt1, attempt2_1, attempt2_2, attempt3_1, attempt3_2, attempt3_3, attempt5, attempt6]
+    )

    # Create and inject spans with diverse data
    # Spans for ro-story-001 (Running) - Multiple nested spans with ongoing execution
@@ -545,11 +537,7 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
        ),
    ]

-    store._spans["ro-story-001"] = spans_ro1
-    store._spans["ro-story-002"] = spans_ro2_a1 + spans_ro2_a2
-    store._spans["ro-story-003"] = spans_ro3_a3
-    store._spans["ro-story-005"] = spans_ro5
-    store._spans["ro-story-006"] = spans_ro6
+    await store.collections.spans.insert(spans_ro1 + spans_ro2_a1 + spans_ro2_a2 + spans_ro3_a3 + spans_ro5 + spans_ro6)

    # Create and inject resources with diverse types
    resource1 = ResourcesUpdate(
@@ -628,11 +616,7 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
        },
    )

-    store._resources["rs-story-001"] = resource1
-    store._resources["rs-story-002"] = resource2
-    store._resources["rs-story-003"] = resource3
-    store._resources["rs-story-004"] = resource4
-    store._resources["rs-story-005"] = resource5
+    await store.collections.resources.insert([resource1, resource2, resource3, resource4, resource5])
    store._latest_resources_id = "rs-story-005"

    # Register workers with diverse states and activity windows.
@@ -694,8 +678,7 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
        ),
    ]

-    for worker in workers:
-        store._workers[worker.worker_id] = worker
+    await store.collections.workers.insert(workers)


 async def main():
@@ -704,7 +687,7 @@ async def main():
    args = parser.parse_args()

    store = InMemoryLightningStore()
-    inject_mock_data(store, now=args.now)
+    await inject_mock_data(store, now=args.now)

    # Start server
    server = LightningStoreServer(store, "127.0.0.1", 8765, "*")
@@ -1,5 +1,18 @@
 # Changelog

+## Agent-lightning v0.2.2 (11/12/2025)
+
+Agent-lightning v0.2.2 is a stabilization release for v0.2.1. It introduces several bug fixes.
+
+* Fix compatibility issues with VERL 0.6.0.
+* Fix model name for pre-downloaded models in VERL.
+* Fix preparing status transition on rollout when creating attempts.
+* Fix OpenAI Agents SDK compatibility issues.
+
+**Full Changelog**: https://github.com/microsoft/agent-lightning/compare/v0.2.1...v0.2.2
+
+---
+
 ## Agent-lightning v0.2.1 (10/30/2025)

 Agent-lightning v0.2.1 is a stabilization release for v0.2.0. It introduces several bug fixes and new features, plus a number of unlisted CI improvements.
@@ -152,6 +152,13 @@ Programmatically this is encapsulated by [`Span.from_opentelemetry(readable_span

    [`add_span`][agentlightning.LightningStore.add_span] or [`add_otel_span`][agentlightning.LightningStore.add_otel_span] both appends a span *and* acts as a heartbeat that can revive `unresponsive` → `running`.

+## OTLP Compatibility
+
+Some of the LightningStore implementations support exporting traces via the [OTLP/HTTP specification](https://opentelemetry.io/docs/specs/otlp/). For example, [`LightningStoreServer`][agentlightning.LightningStoreServer] exposes `/v1/traces` endpoint, it implements the binary Protobuf variant defined by the spec, including the required `Content-Type: application/x-protobuf`, optional `Content-Encoding: gzip`, and status responses encoded as `google.rpc.Status`. Agent-lightning helps parsing `ExportTraceServiceRequest` messages, validate identifiers, normalize resource metadata, and allocate sequence
+numbers so store implementations only need to persist [`Span`][agentlightning.Span] objects in order.
+
+Because the interface speaks standard OTLP, any OpenTelemetry-compatible SDK or collector can emit spans directly to a LightningStore OTLP endpoint without custom shims. The server responds according to the OTLP contract (status code, encoding, and error payloads), which keeps Agent-lightning interoperable with existing observability tooling. This compatibility serves as a strong complement to the OpenTelemetry conversion discussed above.
+
 ## Store Implementations

 Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore]:
@@ -162,6 +169,8 @@ Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][

 For production you will likely want persistence. We’re actively building a SQLite-backed store that keeps the same API surface while adding durability, crash recovery, and better historical span queries. If you need something sooner, implement your own store by subclassing [`LightningStore`][agentlightning.LightningStore] and providing concrete storage for the small set of abstract methods (`enqueue_rollout`, `dequeue_rollout`, `update_attempt`, `add_span`, etc.). This document plus the tests in `tests/store/` illustrate the expected behavior.

+Different store implementations may have different capabilities. For example, [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] does not support exporting traces via OTLP. Try to distinguish the capabilities of a store implementation by checking the [`capabilities`][agentlightning.LightningStore.capabilities] property.
+
 ## Thread Safety

 **[`LightningStoreThreaded`][agentlightning.LightningStoreThreaded]** is a subclass of [`LightningStore`][agentlightning.LightningStore] that wraps another underlying store to make a store instance safe for multi-threaded callers. It wraps every state-mutating call in a mutex. Specifically:
@@ -196,7 +205,7 @@ await server.start()      # starts uvicorn in a daemon thread and waits for /hea
 # Client (same or different process)
 client = agl.LightningStoreClient("http://localhost:4747")

-print(await client.query_rollouts(status=["queuing"]))
+print(await client.query_rollouts(status_in=["queuing"]))

 await client.close()
 await server.stop()
@@ -321,6 +321,44 @@ For the LLaMA profile, export an `HF_TOKEN` before running so VERL can download
    env RAY_DEBUG=legacy HYDRA_FULL_ERROR=1 VLLM_USE_V1=1 ray start --head --dashboard-host=0.0.0.0
    ```

+!!! note "Launching Training with NPUs"
+
+    The example also supports running with **Huawei Ascend NPUs**. This feature is contributed by [Teams from Huawei](https://github.com/microsoft/agent-lightning/pull/272). To use it, resort to the function `config_train_npu` in the script.
+
+    **Hardware Supported:** Atlas 200T A2 Box16, Atlas 900 A2 PODc, Atlas 800T A3. At least **a single 40GB NPU** is required to run the **Qwen2.5-Coder-1.5B-Instruct** model.
+
+    **Environment Setup:** Python 3.11.13, CANN 8.2.RC1, torch 2.7.1+cpu, torch_npu 2.7.1.dev20250724. For basic environment preparation, please refer to this [document](https://gitcode.com/Ascend/pytorch).
+
+    Before installing dependencies, configure the following pip mirrors:
+
+    ```bash
+    pip config set global.index-url http://repo.huaweicloud.com/repository/pypi/simple
+    pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/ https://mirrors.huaweicloud.com/ascend/repos/pypi"
+    ```
+
+    Then install vLLM, vLLM-Ascend and VERL:
+
+    ```bash
+    pip install vllm==0.10.0 --trusted-host repo.huaweicloud.com
+    pip install vllm-Ascend==0.10.0rc1 --trusted-host repo.huaweicloud.com
+    pip install verl==0.5.0
+    ```
+
+    To ensure the VERL framework runs correctly on NPU, add the following lines to `verl/utils/vllm_utils.py`:
+
+    ```python
+    from vllm_ascend.patch import platform
+    from vllm_ascend.patch import worker
+    ```
+
+    See the following reference for more details: [https://github.com/vllm-project/vllm-ascend/issues/1776](https://github.com/vllm-project/vllm-ascend/issues/1776).
+
+    After the above dependencies have been installed, from [`examples/spider`]({{ src("examples/spider") }}) run the following script command:
+
+    ```bash
+    python train_sql_agent.py npu
+    ```
+
 ### Debugging the Agent without VERL

 [`sql_agent.py`]({{ src("examples/spider/sql_agent.py") }}) also provides a `debug_sql_agent()` helper to run the LangGraph workflow directly against a local or hosted OpenAI-compatible endpoint before using VERL.
@@ -99,7 +99,7 @@ async def find_best_prompt(store, prompts_to_test, task_input):
        await store.wait_for_rollouts([rollout.rollout_id])

        # 4. Query the completed rollout and its spans
-        completed_rollout = (await store.query_rollouts([rollout.rollout_id]))[0]
+        completed_rollout = await store.get_rollout_by_id(rollout.rollout_id)
        print(f"[Algo] Received Result: {completed_rollout.model_dump_json(indent=None)}")

        spans = await store.query_spans(rollout.rollout_id)
@@ -38,7 +38,7 @@

 ::: agentlightning.store.utils.propagate_status

-::: agentlightning.tracer.agentops.LightningSpanProcessor
+::: agentlightning.tracer.otel.LightningSpanProcessor

 ::: agentlightning.utils.server_launcher.PythonServerLauncher

@@ -46,6 +46,10 @@

 ::: agentlightning.utils.server_launcher.LaunchMode

+::: agentlightning.utils.otlp.handle_otlp_export
+
+::: agentlightning.utils.otlp.spans_from_proto
+
 ## Deprecated APIs

 ::: agentlightning.server.AgentLightningServer
@@ -2,10 +2,14 @@

 ::: agentlightning.LightningStore

+::: agentlightning.LightningStoreCapabilities
+
 ## Store Implementations

 ::: agentlightning.InMemoryLightningStore

+::: agentlightning.CollectionBasedLightningStore
+
 ## Client-Server and Thread-safe Wrappers

 ::: agentlightning.LightningStoreServer
@@ -13,3 +17,21 @@
 ::: agentlightning.LightningStoreClient

 ::: agentlightning.LightningStoreThreaded
+
+## Collections and Collection Implementations
+
+::: agentlightning.store.collection.Collection
+
+::: agentlightning.store.collection.Queue
+
+::: agentlightning.store.collection.KeyValue
+
+::: agentlightning.store.collection.LightningCollections
+
+::: agentlightning.store.collection.ListBasedCollection
+
+::: agentlightning.store.collection.DequeBasedQueue
+
+::: agentlightning.store.collection.DictBasedKeyValue
+
+::: agentlightning.store.collection.InMemoryLightningCollections
@@ -23,3 +23,11 @@
 ## CLI Builder

 ::: agentlightning.lightning_cli
+
+## Logging
+
+::: agentlightning.configure_logger
+
+::: agentlightning.setup_module_logging
+
+::: agentlightning.setup_logging
@@ -32,6 +32,14 @@

 ::: agentlightning.Hook

+::: agentlightning.PaginatedResult
+
+::: agentlightning.FilterOptions
+
+::: agentlightning.SortOptions
+
+::: agentlightning.FilterField
+
 ## Resources

 ::: agentlightning.Resource
@@ -7,10 +7,11 @@ This catalog highlights the examples shipped with Agent-lightning.
 | [apo](./apo) | Automatic Prompt Optimization tutorials covering built-in, custom, and debugging workflows. | [![apo workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-apo.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-apo.yml) |
 | [azure](./azure) | Supervised fine-tuning with Azure OpenAI. | **Unmaintained** — last verified with Agent-lightning v0.2.1 |
 | [calc_x](./calc_x) | VERL-powered math reasoning agent training that uses AutoGen with an MCP calculator tool. | [![calc_x workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-calc-x.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-calc-x.yml) |
+| [minimal](./minimal) | Bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation. | [![minimal workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
 | [rag](./rag) | Retrieval-Augmented Generation pipeline targeting the MuSiQue dataset with Wikipedia retrieval. | **Unmaintained** — last verified with Agent-lightning v0.1.1 |
 | [search_r1](./search_r1) | Framework-free Search-R1 reinforcement learning training workflow with a retrieval backend. | **Unmaintained** — last verified with Agent-lightning v0.1.2 |
 | [spider](./spider) | Text-to-SQL reinforcement learning training on the Spider dataset using LangGraph. | [![spider workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-spider.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-spider.yml) |
+| [tinker](./tinker) | Reinforcement learning with Tinker as the backend training service. | **Unmaintained** — last verified with Agent-lightning v0.2.2 |
 | [unsloth](./unsloth) | Supervised fine-tuning example powered by Unsloth with 4-bit quantization and LoRA. | [![unsloth workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unsloth.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unsloth.yml) |
-| [tinker](./tinker) | Reinforcement learning with Tinker as the backend training service. | **Unmaintained** — last verified with Agent-lightning v0.2.1 |

 *NOTE: CI status avoid taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. Each example's own README also displays its `examples-*` workflow status whenever the project is maintained by CI.*
@@ -29,7 +29,7 @@ python apo_custom_algorithm_trainer.py

 import argparse
 import asyncio
-from typing import List, Optional
+from typing import Optional, Sequence

 from openai import AsyncOpenAI
 from rich.console import Console
@@ -117,7 +117,7 @@ async def apo_rollout(task: str, prompt_template: agl.PromptTemplate) -> float:
    return await llm_judge(task, text)


-async def log_llm_span(spans: List[agl.Span]) -> None:
+async def log_llm_span(spans: Sequence[agl.Span]) -> None:
    """Logs the LLM related spans that records prompts and responses."""
    for span in spans:
        if "chat.completion" in span.name:
@@ -181,5 +181,5 @@ async def main():


 if __name__ == "__main__":
-    agl.configure_logger()
+    agl.setup_logging()
    asyncio.run(main())
@@ -19,7 +19,7 @@ python apo_custom_algorithm.py runner
 from apo_custom_algorithm import apo_algorithm, apo_rollout
 from rich.console import Console

-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.algorithm import algo
 from agentlightning.store import LightningStore

@@ -39,6 +39,6 @@ async def apo_algorithm_usable_in_trainer(*, store: LightningStore):


 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
    trainer = Trainer(n_workers=1, algorithm=apo_algorithm_usable_in_trainer)
    trainer.fit(apo_rollout)
@@ -8,7 +8,7 @@ from typing import cast

 from apo_custom_algorithm import apo_rollout

-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.litagent import LitAgent
 from agentlightning.runner import LitAgentRunner
 from agentlightning.store import InMemoryLightningStore
@@ -105,7 +105,7 @@ def debug_with_trainer():


 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()

    parser = argparse.ArgumentParser(description="Debug APO with runner or trainer approach.")
    parser.add_argument(
@@ -12,7 +12,7 @@ from typing import Any
 import dotenv
 from openai import OpenAI

-from agentlightning import configure_logger
+from agentlightning import setup_logging
 from agentlightning.litagent import LitAgent
 from agentlightning.trainer import Trainer

@@ -40,7 +40,7 @@ class SimpleAgent(LitAgent[Any]):


 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
    dotenv.load_dotenv()
    agent = SimpleAgent()
    # Use 2 workers to simulate multiple clients
@@ -8,7 +8,7 @@ from typing import Tuple, cast
 from openai import AsyncOpenAI
 from room_selector import RoomSelectionTask, load_room_tasks, prompt_template_baseline, room_selector

-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.adapter import TraceToMessages
 from agentlightning.algorithm.apo import APO
 from agentlightning.types import Dataset
@@ -33,7 +33,7 @@ def setup_apo_logger(file_path: str = "apo.log") -> None:


 def main() -> None:
-    configure_logger()
+    setup_logging()
    setup_apo_logger()

    openai_client = AsyncOpenAI()
@@ -2,7 +2,7 @@

 from aoai_finetune import AzureOpenAIFinetune

-from agentlightning import configure_logger
+from agentlightning import setup_logging

 finetune_algo = AzureOpenAIFinetune(
    base_deployment_name="gpt-4.1-mini",
@@ -12,7 +12,7 @@ finetune_algo = AzureOpenAIFinetune(
    data_filter_ratio=0.6,
 )

-configure_logger()
+setup_logging()


 def test_deployment():
@@ -5,13 +5,13 @@ from aoai_finetune import AzureOpenAIFinetune
 from capital_agent import capital_agent
 from rich.console import Console

-from agentlightning import TraceToMessages, Trainer, configure_logger
+from agentlightning import TraceToMessages, Trainer, setup_logging

 console = Console()


 def main():
-    configure_logger()
+    setup_logging()
    finetune_algo = AzureOpenAIFinetune(
        base_deployment_name="gpt-4.1-mini",
        finetuned_deployment_name="gpt-4.1-mini-ft",
@@ -84,7 +84,9 @@ async def calc_agent(task: MathProblem, llm: agl.LLM) -> None:
        try:
            output_format = "Output the answer when you are ready. The answer should be surrounded by three sharps (`###`), in the form of ### ANSWER: <answer> ###."
            prompt = task["question"] + " " + output_format
-            result = await calc_agent.run(task=prompt)
+            # Sometimes MCP tools can timeout. In that case, the whole agent will block.
+            # We thus set a timeout of 5 minutes so that the agent will not block indefinitely.
+            result = await asyncio.wait_for(calc_agent.run(task=prompt), timeout=300.0)
            # evaluate
            last_message = cast(str, result.messages[-1].content)  # type: ignore
            answer = re.search(r"###\s*ANSWER:\s*(.+?)(\s*###|$)", last_message)
@@ -92,6 +94,9 @@ async def calc_agent(task: MathProblem, llm: agl.LLM) -> None:
                answer = answer.group(1)
            else:
                answer = last_message
+        except asyncio.TimeoutError as e:
+            print("Timeout occurred. Error:", str(e))
+            answer = "None"
        except Exception as e:
            print("Failure:", str(e))
            answer = "None"
@@ -19,9 +19,9 @@ from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
 from eval_utils import evaluate_v0_1

-from agentlightning import LLM, LitAgent, NamedResources, Trainer, configure_logger
+from agentlightning import LLM, LitAgent, NamedResources, Trainer, setup_logging

-configure_logger()
+setup_logging()

 calculator_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-calculator"])

@@ -30,6 +30,7 @@ AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=runner python train_calc_agent.py --externa

 import argparse
 import os
+import uuid
 from datetime import datetime
 from typing import Any, Dict, Optional, cast

@@ -146,32 +147,38 @@ def train(
    if ci or ci_fast:
        # Config the experiment name and project name so that they are available to CI
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-        EXPERIMENT_NAME = f"calc_x_{timestamp}"
+        random_suffix = uuid.uuid4().hex[:8]
+        EXPERIMENT_NAME = f"calc_x_{timestamp}_{random_suffix}"

        PROJECT_NAME = "AgentLightningCI"

-        # Simulate writing to $GITHUB_OUTPUT if it’s set
-        github_output = os.getenv("GITHUB_OUTPUT")
-        if github_output:
-            with open(github_output, "a") as f:
-                f.write(f"project_name={PROJECT_NAME}\n")
-                f.write(f"run_name={EXPERIMENT_NAME}\n")
+        # Skip this step if AGL_CURRENT_ROLE is runner
+        agl_current_role = os.getenv("AGL_CURRENT_ROLE")

-        print("Set environment variables:")
-        print(f"PROJECT_NAME={PROJECT_NAME}")
-        print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
+        if agl_current_role != "runner":
+            # 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}")

        # Keep it tiny/light without adding new knobs
-        config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
+        config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.8
        config["trainer"]["total_epochs"] = 1
-        config["trainer"]["total_training_steps"] = 6
-        config["trainer"]["test_freq"] = 6
+        config["trainer"]["total_training_steps"] = 20
+        config["trainer"]["test_freq"] = 20
        config["trainer"]["experiment_name"] = EXPERIMENT_NAME
        config["trainer"]["project_name"] = PROJECT_NAME
        config["trainer"].pop("save_freq", None)

        if ci_fast:
            # Extra fast CI toggle for testing purposes.
+            config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
            config["trainer"]["total_training_steps"] = 1
            config["trainer"]["test_freq"] = 1

@@ -209,6 +216,7 @@ def main():
        default="",
        help="Connect to an external store instead of creating a new one in memory",
    )
+    parser.add_argument("--debug", action="store_true", help="Enable debug logging")

    args = parser.parse_args()

@@ -223,6 +231,8 @@ def main():
    if args.ci_fast:
        args.ci = True

+    agl.setup_logging("DEBUG" if args.debug else "INFO")
+
    train(
        train_file=args.train_file,
        val_file=args.val_file,
@@ -0,0 +1,17 @@
+# Minimal Component Showcase
+
+`examples/minimal` provides bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation.
+
+Each module have been documented with its own CLI usage in the module-level docstring. Use this directory as a reference when wiring the same pieces into a larger system.
+
+## What’s Included?
+
+| Component | Demonstrated In | Highlights |
+| --- | --- | --- |
+| LightningStore + OTLP ingestion | `write_traces.py` | Shows how `OtelTracer` and `AgentOpsTracer` open rollouts, emit spans, and optionally forward them to a remote store client. |
+| LLM proxying | `llm_proxy.py` | Guards either OpenAI or a local vLLM deployment with `LLMProxy`, proving how requests are routed through `/rollout/<id>/attempt/<id>` namespaces and captured in the store. |
+| vLLM lifecycle | `vllm_server.py` | Minimal context manager that shells out to `vllm serve`, monitors readiness, and tears down the process safely. |
+
+All runtime instructions (CLI arguments, required environment variables, etc.) are embedded directly in each script’s top-level docstring so the source stays self-documenting.
+
+For full-fledged training workflows or multi-component experiments, browse the other subdirectories under `examples/`. This `minimal` folder deliberately keeps each demonstration focused on a single component so you can understand and test them independently.
@@ -0,0 +1,254 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Examples to serve an LLM proxy for a vLLM server or an OpenAI service.
+
+Usage: run one of the following commands to start a server.
+
+```bash
+python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct
+```
+
+Use the following command to test the LLM proxy.
+
+```bash
+python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
+```
+
+You can also test the OpenAI Proxy path (`OPENAI_API_KEY` environment variable is required).
+
+```bash
+dotenv run python llm_proxy.py openai gpt-4.1-mini
+```
+"""
+
+import argparse
+import asyncio
+import os
+from typing import Sequence, no_type_check
+
+import aiohttp
+from portpicker import pick_unused_port
+from rich.console import Console
+from vllm_server import vllm_server
+
+import agentlightning as agl
+
+console = Console()
+
+
+async def serve_llm_proxy_with_vllm(model_name: str, store_port: int = 43887):
+    """Serve an LLM proxy for a vLLM server."""
+    # Create a store to store the traces
+    store = agl.InMemoryLightningStore()
+    store_server = agl.LightningStoreServer(store, "127.0.0.1", store_port)
+    await store_server.start()
+
+    # Create a vLLM server
+    vllm_port = pick_unused_port()
+    with vllm_server(model_name, vllm_port) as vllm_endpoint:
+        # Server is up.
+
+        # Create an LLM proxy to guard the vLLM server and catch the traces
+        llm_proxy = agl.LLMProxy(
+            port=43886,
+            model_list=[
+                {
+                    "model_name": model_name,
+                    "litellm_params": {
+                        "model": f"hosted_vllm/{model_name}",
+                        "api_base": vllm_endpoint,
+                    },
+                }
+            ],
+            store=store_server,
+        )
+
+        try:
+            await llm_proxy.start()
+
+            # Wait forever
+            await asyncio.sleep(float("inf"))
+
+        finally:
+            # Stop the LLM proxy and the store server
+            await llm_proxy.stop()
+            await store_server.stop()
+
+
+async def serve_llm_proxy_with_openai(model_name: str, store_port: int = 43887):
+    """Serve an LLM proxy for an OpenAI server."""
+    # Create a store to store the traces
+    store = agl.InMemoryLightningStore()
+    store_server = agl.LightningStoreServer(store, "127.0.0.1", store_port)
+    await store_server.start()
+
+    if not os.getenv("OPENAI_API_KEY"):
+        raise ValueError("OPENAI_API_KEY environment variable is not set")
+
+    # Create an LLM proxy to guard the OpenAI server and catch the traces
+    llm_proxy = agl.LLMProxy(
+        port=43886,
+        model_list=[
+            {
+                "model_name": model_name,
+                "litellm_params": {
+                    "model": "openai/" + model_name,
+                    # Must have OpenAI API key set in the environment variable
+                },
+            }
+        ],
+        store=store_server,
+        callbacks=["opentelemetry"],
+    )
+
+    try:
+        await llm_proxy.start()
+        # Wait forever
+        await asyncio.sleep(float("inf"))
+    finally:
+        # Stop the LLM proxy and the store server
+        await llm_proxy.stop()
+        await store_server.stop()
+
+
+async def test_llm_proxy(model_name: str, store_port: int = 43887):
+    """Test the LLM proxy by sending a request to the proxy and checking the response.
+
+    We do it via aiohttp here. This can also be done with OpenAI client.
+    """
+    # We first connect to the store server and start a rollout.
+    store = agl.LightningStoreClient(f"http://localhost:{store_port}")
+    rollout = await store.start_rollout(input={"origin": "test_llm_proxy"})
+
+    # The chat completion URL is simply /v1/chat/completions under the namespace of current rollout and attempt.
+    # This ensures the traces are properly put into the correct bucket.
+    chat_completion_url = (
+        f"http://localhost:43886/rollout/{rollout.rollout_id}/attempt/{rollout.attempt.attempt_id}/v1/chat/completions"
+    )
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(
+            chat_completion_url,
+            json={
+                "model": model_name,
+                "messages": [{"role": "user", "content": "Hello, what's your name?"}],
+            },
+        ) as response:
+            response_body = await response.json()
+            console.print("Response body:", response_body)
+            _verify_response_body(response_body, model_name)
+
+    spans = await store.query_spans(rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id)
+    for span in spans:
+        console.print("Span:", span)
+    _verify_span(spans)
+
+    await store.close()
+
+
+@no_type_check
+def _verify_response_body(response_body: dict, model_name: str):
+    """Expect Response body to be something like this:
+
+    ```python
+    {
+        'id': 'chatcmpl-996a90a8678e4ed0a0d2724df2c0bba5',
+        'created': 1763178218,
+        'model': 'hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct',
+        'object': 'chat.completion',
+        'choices': [
+            {
+                'finish_reason': 'stop',
+                'index': 0,
+                'message': {
+                    'content': 'Hello! I am Qwen, an AI language model created by Alibaba Cloud. My name is Qwen, and I can assist you with
+    various tasks and provide information on a wide range of topics. How may I help you today?',
+                    'role': 'assistant'
+                },
+                'provider_specific_fields': {
+                    'stop_reason': None,
+                    'token_ids': [9707, 0, ...],
+                }
+            }
+        ],
+        'usage': {'completion_tokens': 48, 'prompt_tokens': 36, 'total_tokens': 84},
+        'prompt_token_ids': [151644, 8948, ...],
+    }
+    ```
+    """
+    if "qwen" in model_name.lower():
+        assert "qwen" in response_body["choices"][0]["message"]["content"].lower()
+        assert (
+            "provider_specific_fields" in response_body["choices"][0]
+        ), "provider_specific_fields not found in response body"
+        assert (
+            "token_ids" in response_body["choices"][0]["provider_specific_fields"]
+        ), "token_ids not found in response body"
+        assert "prompt_token_ids" in response_body, "prompt_token_ids not found in response body"
+    else:
+        assert "chatgpt" in response_body["choices"][0]["message"]["content"].lower()
+
+
+def _verify_span(spans: Sequence[agl.Span]):
+    """Only a few spans are checked here.
+
+    `raw_gen_ai_request` span:
+
+    ```python
+    Span(
+        rollout_id='ro-4c68a7e686a1',
+        attempt_id='at-308eb814',
+        sequence_id=1,
+        name='raw_gen_ai_request',
+        attributes={
+            'llm.hosted_vllm.messages': '[{\'role\': \'user\', \'content\': "Hello, what\'s your name?"}]',
+            'llm.hosted_vllm.extra_body': "{'return_token_ids': True}",
+            'llm.hosted_vllm.choices': '... \'token_ids\': [40, 1079, 1207, 16948, ...',
+            'llm.hosted_vllm.model': 'Qwen/Qwen2.5-0.5B-Instruct',
+            'llm.hosted_vllm.prompt_token_ids': '[151644, 8948, ...]',
+        },
+        resource=OtelResource(
+            attributes={
+                'agentlightning.rollout_id': 'ro-4c68a7e686a1',
+                'agentlightning.attempt_id': 'at-308eb814',
+                'agentlightning.span_sequence_id': 1
+            },
+        )
+    )
+    ```
+    """
+
+    assert len(spans) > 1
+    has_raw_gen_ai_request = False
+    for span in spans:
+        if span.name == "raw_gen_ai_request":
+            has_raw_gen_ai_request = True
+            if "llm.hosted_vllm.messages" in span.attributes:
+                assert "return_token_ids" in span.attributes["llm.hosted_vllm.extra_body"]  # type: ignore
+                assert "token_ids" in span.attributes["llm.hosted_vllm.choices"]  # type: ignore
+                assert span.attributes["llm.hosted_vllm.prompt_token_ids"]
+        assert "agentlightning.rollout_id" in span.resource.attributes
+        assert "agentlightning.attempt_id" in span.resource.attributes
+        assert "agentlightning.span_sequence_id" in span.resource.attributes
+
+    assert has_raw_gen_ai_request, "raw_gen_ai_request span not found"
+
+
+if __name__ == "__main__":
+    agl.setup_logging()
+    parser = argparse.ArgumentParser(description="LLM Proxy runner")
+    parser.add_argument(
+        "mode",
+        choices=["vllm", "openai", "test"],
+        help="Which function to run",
+    )
+    parser.add_argument("model", type=str, help="Model name to serve.")
+
+    args = parser.parse_args()
+
+    if args.mode == "vllm":
+        asyncio.run(serve_llm_proxy_with_vllm(args.model))
+    elif args.mode == "openai":
+        asyncio.run(serve_llm_proxy_with_openai(args.model))
+    elif args.mode == "test":
+        asyncio.run(test_llm_proxy(args.model))
@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Programmatically launch and stop an vLLM server."""
+
+import subprocess
+import time
+from contextlib import contextmanager
+from typing import Optional
+
+import httpx
+from openai import OpenAI
+from rich.console import Console
+
+console = Console()
+
+
+@contextmanager
+def vllm_server(
+    model_path: str,
+    port: int,
+    startup_timeout: float = 300.0,
+    terminate_timeout: float = 10.0,
+    gpu_memory_utilization: float = 0.7,
+    auto_tool_choice: bool = True,
+    tool_call_parser: Optional[str] = "hermes",
+):
+    """Serves a vLLM model from command line.
+
+    Args:
+        model_path: The path to the vLLM model. It can be either a local path or a Hugging Face model ID.
+        port: The port to serve the model on.
+        startup_timeout: The timeout for the server to start.
+        terminate_timeout: The timeout for the server to terminate.
+        gpu_memory_utilization: The GPU memory utilization for the server. Set it lower to avoid OOM.
+        auto_tool_choice: Whether to enable auto tool choice.
+        tool_call_parser: The tool call parser to use.
+    """
+    proc: Optional[subprocess.Popen[bytes]] = None
+    try:
+        vllm_serve_args = [
+            "--gpu-memory-utilization",
+            str(gpu_memory_utilization),
+            "--port",
+            str(port),
+        ]
+        if auto_tool_choice:
+            vllm_serve_args.append("--enable-auto-tool-choice")
+        if tool_call_parser is not None:
+            vllm_serve_args.append("--tool-call-parser")
+            vllm_serve_args.append(tool_call_parser)
+
+        proc = subprocess.Popen(["vllm", "serve", model_path, *vllm_serve_args])
+
+        # Wait for the server to be ready
+        url = f"http://localhost:{port}/health"
+        start = time.time()
+        client = httpx.Client()
+
+        while True:
+            try:
+                if client.get(url).status_code == 200:
+                    break
+            except Exception:
+                result = proc.poll()
+                if result is not None and result != 0:
+                    raise RuntimeError("Server exited unexpectedly.") from None
+                time.sleep(0.5)
+                if time.time() - start > startup_timeout:
+                    raise RuntimeError(f"Server failed to start in {startup_timeout} seconds.") from None
+
+        yield f"http://localhost:{port}/v1"
+    finally:
+        # Terminate the server
+        if proc is None:
+            return
+        proc.terminate()
+        try:
+            proc.wait(terminate_timeout)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+
+
+if __name__ == "__main__":
+    with vllm_server("Qwen/Qwen2.5-0.5B-Instruct", 8080) as endpoint:
+        client = OpenAI(base_url=endpoint, api_key="dummy")
+        response = client.chat.completions.create(
+            model="Qwen/Qwen2.5-0.5B-Instruct",
+            messages=[{"role": "user", "content": "Hello, what's your name?"}],
+        )
+        console.print(response)
+        assert "qwen" in response.choices[0].message.content.lower()  # type: ignore
@@ -0,0 +1,187 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example to write traces to a LightningStore via raw OpenTelemetry or AgentOpsTracer.
+
+The example can be run with or without using a Lightning Store server.
+When running this server, the traces will be written to the server via OTLP endpoint.
+
+Prior to running this example with `--use-client` flag, please start a LightningStore server with OTLP enabled first:
+
+```bash
+agl store --port 45993 --log-level DEBUG
+```
+"""
+
+import argparse
+import asyncio
+import time
+from typing import Sequence
+
+from openai import AsyncOpenAI
+from rich.console import Console
+
+from agentlightning import AgentOpsTracer, LightningStoreClient, OtelTracer, Span, emit_reward, setup_logging
+from agentlightning.store import InMemoryLightningStore
+
+console = Console()
+
+
+async def send_traces_via_otel(use_client: bool = False):
+    tracer = OtelTracer()
+    if not use_client:
+        store = InMemoryLightningStore()
+    else:
+        store = LightningStoreClient("http://localhost:45993")
+    rollout = await store.start_rollout(input={"origin": "write_traces_example"})
+
+    with tracer.lifespan(store):
+        # Initialize the capture of one single trace for one single rollout
+        async with tracer.trace_context(
+            "trace-manual", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+        ) as tracer:
+            with tracer.start_as_current_span("grpc-span-1"):
+                time.sleep(0.01)
+
+                # Nested Span
+                with tracer.start_as_current_span("grpc-span-2"):
+                    time.sleep(0.01)
+
+            with tracer.start_as_current_span("grpc-span-3"):
+                time.sleep(0.01)
+
+            # This creates a reward span
+            emit_reward(1.0)
+
+    traces = await store.query_spans(rollout_id=rollout.rollout_id)
+    console.print(traces)
+
+    # Quickly validate the traces
+    assert len(traces) == 4
+    span_names = [span.name for span in traces]
+    assert "grpc-span-1" in span_names
+    assert "grpc-span-2" in span_names
+    assert "grpc-span-3" in span_names
+    assert "agentlightning.reward" in span_names
+
+    last_span = traces[-1]
+    assert last_span.name == "agentlightning.reward"
+    # NOTE: Try not to rely on this attribute. It may change in the future.
+    # Use utils from agentlightning.emitter to get the reward value.
+    assert last_span.attributes["reward"] == 1.0
+
+    if use_client:
+        # When using client, the resource should have rollout_id and attempt_id set
+        for span in traces:
+            assert "agentlightning.rollout_id" in span.resource.attributes
+            assert "agentlightning.attempt_id" in span.resource.attributes
+
+    if isinstance(store, LightningStoreClient):
+        await store.close()
+
+
+async def send_traces_via_agentops(use_client: bool = False):
+    tracer = AgentOpsTracer()
+    if not use_client:
+        store = InMemoryLightningStore()
+    else:
+        store = LightningStoreClient("http://localhost:45993")
+    rollout = await store.start_rollout(input={"origin": "write_traces_example"})
+
+    # Initialize the tracer lifespan
+    # One lifespan can contain multiple traces
+    with tracer.lifespan(store):
+        # Initialize the capture of one single trace for one single rollout
+        async with tracer.trace_context(
+            "trace-1", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+        ):
+            openai_client = AsyncOpenAI()
+            response = await openai_client.chat.completions.create(
+                model="gpt-4.1-mini",
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "Hello, what's your name?"},
+                ],
+            )
+            assert response.choices[0].message.content is not None
+            assert "chatgpt" in response.choices[0].message.content.lower()
+
+    traces = await store.query_spans(rollout_id=rollout.rollout_id)
+    console.print(traces)
+    await _verify_agentops_traces(traces, use_client=use_client)
+    if isinstance(store, LightningStoreClient):
+        await store.close()
+
+
+async def _verify_agentops_traces(spans: Sequence[Span], use_client: bool = False):
+    """Expected traces to something like:
+
+    ```python
+    Span(
+        rollout_id='ro-ef9ff8a429d1',
+        attempt_id='at-37cc5f24',
+        sequence_id=1,
+        trace_id='b3a16b603f7805934215d467e717c9e7',
+        span_id='2782d5d750f49b2d',
+        parent_id='2fb97c818363bce3',
+        name='openai.chat.completion',
+        status=TraceStatus(status_code='OK', description=None),
+        attributes={
+            'gen_ai.request.type': 'chat',
+            'gen_ai.system': 'OpenAI',
+            'gen_ai.request.model': 'gpt-4.1-mini',
+            'gen_ai.request.streaming': False,
+            'gen_ai.prompt.0.role': 'system',
+            'gen_ai.prompt.0.content': 'You are a helpful assistant.',
+            'gen_ai.prompt.1.role': 'user',
+            'gen_ai.prompt.1.content': "Hello, what's your name?",
+            'gen_ai.response.id': 'chatcmpl-Cc1osPWiArOwCS8nUkp0kZuZPkpY4',
+            'gen_ai.response.model': 'gpt-4.1-mini-2025-04-14',
+            'gen_ai.completion.0.role': 'assistant',
+            'gen_ai.completion.0.content': "Hello! I'm ChatGPT, your AI assistant. How can I help you today?",
+        },
+        resource=OtelResource(
+            attributes={
+                'agentops.project.id': 'temporary',
+                'agentlightning.rollout_id': 'ro-ef9ff8a429d1',
+                'agentlightning.attempt_id': 'at-37cc5f24'
+            },
+            schema_url=''
+        )
+    )
+    ```
+    """
+    assert len(spans) == 2
+    for span in spans:
+        if span.name == "openai.chat.completion":
+            assert span.attributes["gen_ai.request.model"] == "gpt-4.1-mini"
+            assert span.attributes["gen_ai.request.streaming"] == False
+            assert span.attributes["gen_ai.prompt.0.role"] == "system"
+            assert span.attributes["gen_ai.prompt.0.content"] == "You are a helpful assistant."
+            assert span.attributes["gen_ai.prompt.1.role"] == "user"
+            assert span.attributes["gen_ai.prompt.1.content"] == "Hello, what's your name?"
+            assert "chatgpt" in span.attributes["gen_ai.completion.0.content"].lower()  # type: ignore
+            if use_client:
+                assert "agentlightning.rollout_id" in span.resource.attributes
+                assert "agentlightning.attempt_id" in span.resource.attributes
+        else:
+            assert "trace-1" in span.name
+            assert span.attributes["agentops.span.kind"] == "session"
+
+
+def main():
+    setup_logging("DEBUG")
+    parser = argparse.ArgumentParser()
+    parser.add_argument("mode", choices=["otel", "agentops"])
+    parser.add_argument("--use-client", action="store_true")
+    args = parser.parse_args()
+
+    if args.mode == "otel":
+        asyncio.run(send_traces_via_otel(use_client=args.use_client))
+    elif args.mode == "agentops":
+        asyncio.run(send_traces_via_agentops(use_client=args.use_client))
+    else:
+        raise ValueError(f"Invalid mode: {args.mode}")
+
+
+if __name__ == "__main__":
+    main()
@@ -15,10 +15,10 @@ from agentlightning import (
    LitAgent,
    NamedResources,
    Trainer,
-    configure_logger,
+    setup_logging,
 )

-configure_logger()
+setup_logging()

 agent_prompt = """You are an assistant who answers questions using Wikipedia retriever. Answer the question using only the retrieved passages. Verify your answer directly against the text.

@@ -8,9 +8,9 @@ import requests
 from openai import OpenAI
 from qa_em import compute_score_em

-from agentlightning import LLM, LitAgent, NamedResources, Trainer, configure_logger, reward
+from agentlightning import LLM, LitAgent, NamedResources, Trainer, reward, setup_logging

-configure_logger()
+setup_logging()

 # 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: """
@@ -37,6 +37,8 @@ Train a SQL agent using the Qwen2.5-Coder-1.5B-Instruct model with the following
 python train_sql_agent.py qwen
 ```

+If you want to use an NPU for training, please refer to the **Launch Training with NPUS** section in [How to Train a SQL Agent](../../docs/how-to/train-sql-agent.md).
+
 ### Debugging

 To test and debug the SQL agent interactively:
@@ -9,6 +9,7 @@ as well as https://langchain-ai.github.io/langgraph/tutorials/sql-agent/

 from __future__ import annotations

+import logging
 import os
 import re
 import shutil
@@ -29,9 +30,9 @@ from spider_eval.exec_eval import eval_exec_match

 import agentlightning as agl

-agl.configure_logger()
+agl.setup_logging(apply_to=[__name__])

-logger = agl.configure_logger(name=__name__)
+logger = logging.getLogger(__name__)


 WRITE_QUERY_PROMPT = ChatPromptTemplate(
@@ -137,6 +137,20 @@ def config_train_qwen() -> Dict[str, Any]:
    return config


+def config_train_npu() -> Dict[str, Any]:
+    """A configuration for training with NPU."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    del config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["enable_auto_tool_choice"]
+    del config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["tool_call_parser"]
+    del config["trainer"]["logger"][1]
+    config["actor_rollout_ref"]["actor"]["use_torch_compile"] = False
+    config["trainer"]["val_before_train"] = False
+    config["trainer"]["save_freq"] = 256
+    config["trainer"]["device"] = "npu"
+    return config
+
+
 def config_train_llama() -> Dict[str, Any]:
    """A configuration for training with LLaMA-3.2-1B-Instruct.

@@ -171,8 +185,8 @@ def main() -> None:

    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)",
+        choices=["fast", "qwen", "llama", "npu"],
+        help="Training configuration: 'fast' (CI testing), 'qwen' (Qwen-2.5-Coder-1.5B), 'llama' (LLaMA-3.2-3B),'npu' (Train with NPU)",
    )

    parser.add_argument(
@@ -182,8 +196,12 @@ def main() -> None:
    args = parser.parse_args()

    # Get the appropriate configuration
-    config_functions = {"fast": config_train_fast, "qwen": config_train_qwen, "llama": config_train_llama}
-
+    config_functions = {
+        "fast": config_train_fast,
+        "qwen": config_train_qwen,
+        "llama": config_train_llama,
+        "npu": config_train_npu,
+    }
    config = config_functions[args.config]()

    # Set active agent - use provided value or default based on config choice
@@ -40,7 +40,7 @@ T_task = TypeVar("T_task")
 WAIT_FOR_ROLLOUTS_INTERVAL = 5.0


-def reconstruct_transitions(spans: List[Span], adapter: TraceToTripletBase, rollout_id: str) -> Trajectory:
+def reconstruct_transitions(spans: Sequence[Span], adapter: TraceToTripletBase, rollout_id: str) -> Trajectory:
    """Convert Agent-lightning spans into a Tinker `Trajectory`.

    This function infers observations, actions, and rewards from the trace triplets emitted by Agent-lightning's
@@ -242,13 +242,13 @@ async def do_group_of_group_rollouts(

    # 4) Await all groups, but still allow interleaving via the shared semaphore.
    trajectory_groups: List[TrajectoryGroup] = []
-    for group_idx, (builder, tasks) in enumerate(zip(env_group_builders_P, per_group_tasks)):
+    for group_idx, (builder, group_envs, tasks) in enumerate(zip(env_group_builders_P, groups_envs, per_group_tasks)):
        rollouts_and_trajectories_G = await asyncio.gather(*tasks)
        rollouts_G, trajectories_G = cast(
            Tuple[List[Rollout], List[Trajectory]], zip(*rollouts_and_trajectories_G, strict=True)
        )
        # Compute rewards/metrics for this group.
-        rewards_and_metrics_G = await builder.compute_group_rewards(trajectories_G)
+        rewards_and_metrics_G = await builder.compute_group_rewards(trajectories_G, group_envs)
        rewards_G, metrics_G = zip(*rewards_and_metrics_G, strict=True)

        # Attach AGL-specific metrics for error handling.
@@ -192,7 +192,7 @@ def main():

    args = parser.parse_args()

-    agl.configure_logger()
+    agl.setup_logging()
    if args.mode == "algo":
        run_algo()
    elif args.mode == "runner":
@@ -367,7 +367,7 @@ def main() -> None:
    runner_parser.set_defaults(func=_run_runner)

    args = parser.parse_args()
-    agl.configure_logger()
+    agl.setup_logging()
    args.func(args)


@@ -7,7 +7,6 @@ It should be included in CI in future if we decided to maintain this example.
 """

 import asyncio
-import logging
 from typing import cast

 import openai
@@ -24,13 +23,12 @@ from agentlightning import (
    LLMProxy,
    LlmProxyTraceToTriplet,
    TracerTraceToTriplet,
-    configure_logger,
    emit_reward,
+    setup_logging,
 )
 from agentlightning.store import LightningStoreThreaded

-configure_logger(name="agentlightning")
-configure_logger(name="agl_tinker", level=logging.INFO)
+setup_logging(apply_to=["agl_tinker"])


 async def test_tracer():
@@ -58,7 +56,7 @@ async def test_tracer():
    try:
        tracer = AgentOpsTracer()
        tracer.init()
-        tracer.init_worker(0)
+        tracer.init_worker(worker_id=0, store=store)

        # init tracer before llm_proxy to avoid tracer provider being not active.
        console.print("Starting LLM proxy...")
@@ -72,7 +70,7 @@ async def test_tracer():
        client = openai.OpenAI(base_url="http://localhost:4000/v1", api_key="dummy")

        async with tracer.trace_context(
-            name="test_llm", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+            name="test_llm", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
        ):
            response = client.chat.completions.create(
                model=model_name,
@@ -29,7 +29,7 @@ from openai import AsyncOpenAI
 from rich.console import Console
 from trl import SFTConfig, SFTTrainer  # type: ignore

-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.litagent import rollout
 from agentlightning.types import LLM, Dataset

@@ -173,5 +173,5 @@ def math_agent_dry_run() -> None:


 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
    math_agent_dry_run()
@@ -32,7 +32,7 @@ from math_agent import GsmProblem, load_math_dataset
 from rich.console import Console
 from unsloth_helper import unsloth_training

-from agentlightning import configure_logger
+from agentlightning import setup_logging
 from agentlightning.adapter import LlmProxyTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store import LightningStore, LightningStoreClient
@@ -380,7 +380,7 @@ async def sft_algorithm(*, store: LightningStore) -> None:


 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()

    store = LightningStoreClient("http://localhost:4747")

@@ -17,7 +17,7 @@ from math_agent import GsmProblem, load_math_dataset, math_agent
 from rich.console import Console
 from sft_algorithm import sft_one_iter

-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.adapter import TraceToTripletBase
 from agentlightning.algorithm import Algorithm
 from agentlightning.llm_proxy import LLMProxy
@@ -94,7 +94,7 @@ class UnslothSupervisedFinetuning(Algorithm):


 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()

    algo = UnslothSupervisedFinetuning(
        max_iterations=2,
@@ -17,7 +17,7 @@ import multiprocessing
 from math_agent import GsmProblem, math_agent
 from rich.console import Console

-from agentlightning import configure_logger
+from agentlightning import setup_logging
 from agentlightning.runner import LitAgentRunner
 from agentlightning.store import LightningStore, LightningStoreClient
 from agentlightning.tracer import OtelTracer
@@ -67,6 +67,6 @@ def spawn_runners(*, store: LightningStore, n_runners: int) -> None:


 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
    store = LightningStoreClient("http://localhost:4747")
    spawn_runners(store=store, n_runners=4)
@@ -1,6 +1,6 @@
 [project]
 name = "agentlightning"
-version = "0.2.2"
+version = "0.3.0"
 description = "Agent-lightning is the absolute trainer to light up AI agents."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -35,7 +35,12 @@ apo = [
 # though it's listed here for completeness.
 verl = [
  "verl>=0.5.0",
-  "vllm>=0.8.4,<0.11.0", # Due to interface change of ExternalZeroMQDistributedExecutor
+  "vllm>=0.8.4", # Due to interface change of ExternalZeroMQDistributedExecutor
+]
+
+# Store-related dependencies.
+mongo = [
+  "pymongo",
 ]

 [project.scripts]
@@ -93,7 +98,10 @@ torch-stable = [
  "torch>=2.8.0",
  "torchvision>=0.23.0",
  "transformers>=4.55.0",
-  "vllm>=0.10.2",
+  # vLLM 0.11.1 requires PyTorch 2.9.0, which is incompatible with flash-attn
+  # https://github.com/Dao-AILab/flash-attention/issues/1967
+  # Similar issues with vLLM 0.11.2
+  "vllm>=0.10.2,!=0.11.1,!=0.11.2",
  # LiteLLM can then be upgraded with new vLLM
  "litellm[proxy]>=1.78",
 ]
@@ -161,7 +169,7 @@ autogen = [
 ]
 openai-agents = [
  "openai-agents",
-  "openai<2.7.0",
+  "openai",
  # Compatibility issues with the latest version of OpenAI, temporarily fixed version
  # Issue: https://github.com/openai/openai-agents-python/issues/2038
  "mcp",
@@ -196,7 +204,7 @@ agents = [
 ]

 [tool.uv]
-required-version = ">=0.9.5"
+required-version = ">=0.9.9"
 conflicts = [
  [
    { group = "core-legacy" },
@@ -233,7 +241,7 @@ torch = [
  { index = "pytorch-cu128", group = "torch-cu128" },
  { index = "pytorch-cpu", group = "torch-cpu" },
 ]
-tinker_cookbook = { git = "https://github.com/thinking-machines-lab/tinker-cookbook", rev = "72ba5e6a1f52c0887e2674615e318ce21a39cc2a" }
+tinker_cookbook = { git = "https://github.com/thinking-machines-lab/tinker-cookbook", rev = "20e26a629797188aa8c6f34474b0d4757b20b90d" }

 [[tool.uv.index]]
 name = "pypi"
@@ -282,6 +290,13 @@ exclude = [

 [tool.pytest.ini_options]
 testpaths = ["tests"]
+markers = [
+  "openai: tests that require OpenAI API",
+  "gpu: tests that require GPU",
+  "agentops: tests that require AgentOps",
+  "llmproxy: tests that require LiteLLM",
+  "mongo: tests that require MongoDB",
+]

 [tool.black]
 line-length = 120
@@ -7,7 +7,9 @@
    "agentlightning/instrumentation",
    "agentlightning/algorithm/apo",
    "agentlightning/algorithm/verl",
-    "agentlightning/cli/vllm.py"
+    "agentlightning/cli/vllm.py",
+    "agentlightning/store/collection/mongo.py",
+    "agentlightning/store/mongo.py"
  ],

  "pythonVersion": "3.12",
@@ -65,10 +65,14 @@ module.exports = async function badgeAggregation({ github, context, core, depend
      workflow_id: dep.workflow,
      branch: 'main', // Always check the main branch status no matter what
      status: 'completed', // only completed runs
-      per_page: 1, // latest only
+      per_page: 50, // retrieve latest 50 so we can filter
+      sort: 'created',
+      direction: 'desc',
    });

-    const run = runsData?.workflow_runs?.[0];
+    const filteredRuns = runsData?.workflow_runs?.filter(run => ['schedule', 'workflow_dispatch'].includes(run.event));
+
+    const run = filteredRuns?.[0];
    if (!run) {
      failures.push(`No completed run found for ${dep.label} on branch "${branch}"`);
      continue;
@@ -84,18 +88,22 @@ module.exports = async function badgeAggregation({ github, context, core, depend

    // Match each required variant to a job. We look for the variant in parentheses, e.g. "(latest)".
    for (const variant of dep.variants || []) {
-      const job = jobs.find(j => typeof j.name === 'string' && j.name.includes(variant));
+      const matchingJobs = jobs.filter(
+        j => typeof j.name === 'string' && j.name.includes(variant)
+      );

-      if (!job) {
+      if (matchingJobs.length === 0) {
        failures.push(`Missing job for ${dep.label} (variant: ${variant})`);
        continue;
      }

-      core.info(`[${dep.label}] ${job.name} => ${job.conclusion}`);
+      for (const job of matchingJobs) {
+        core.info(`[${dep.label}] ${job.name} => ${job.conclusion}`);

-      // Accept only a strict "success".
-      if (job.conclusion !== 'success') {
-        failures.push(`${dep.label} (${job.name}) concluded ${job.conclusion}`);
+        // Accept only a strict "success".
+        if (job.conclusion !== 'success') {
+          failures.push(`${dep.label} (${job.name}) concluded ${job.conclusion}`);
+        }
      }
    }
  }
@@ -25,7 +25,12 @@ sudo apt-get install -y \
    tmux \
    vim \
    git-lfs \
-    nodejs
+    nodejs \
+    gnupg2 \
+    apt-transport-https \
+    ca-certificates \
+    gnupg \
+    lsb-release

 git lfs install

@@ -37,9 +42,9 @@ sudo reboot now

 # Install CUDA Toolkit
 wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
-sudo dpkg -i cuda-keyring_1.1-1_all.deb
+sudo dpkg -i cuda-keyring_1.1-1_all.deb && rm cuda-keyring_1.1-1_all.deb
 sudo apt-get update
-sudo apt-get -y install cuda-toolkit-12-8
+sudo apt-get -y install cuda-toolkit
 sudo reboot now

 # Add paths globally
@@ -49,6 +54,67 @@ export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
 EOF
 sudo chmod +x /etc/profile.d/cuda.sh

+# Add Docker's official GPG key
+sudo install -m 0755 -d /etc/apt/keyrings
+sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
+sudo chmod a+r /etc/apt/keyrings/docker.asc
+
+# Add the repository to Apt sources:
+sudo tee /etc/apt/sources.list.d/docker.sources <<EOF
+Types: deb
+URIs: https://download.docker.com/linux/ubuntu
+Suites: $(. /etc/os-release && echo "${UBUNTU_CODENAME:-$VERSION_CODENAME}")
+Components: stable
+Signed-By: /etc/apt/keyrings/docker.asc
+EOF
+
+sudo apt -y update
+
+# Install the Docker packages
+sudo apt -y install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
+
+# Create docker group only if it doesn't exist
+# sudo groupadd docker
+
+# Add current user to docker group if not already a member
+sudo usermod -aG docker "$USER"
+# A hack to add cloudtest user to docker group as well
+sudo sed -i '/^docker:/ s/$/,cloudtest/' /etc/group
+# This shouldn't be run on CI
+# newgrp docker
+
+# Install NVIDIA Container Toolkit
+curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
+  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
+    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
+    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
+
+sudo apt-get update
+export NVIDIA_CONTAINER_TOOLKIT_VERSION=1.18.0-1
+sudo apt-get install -y \
+    nvidia-container-toolkit=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
+    nvidia-container-toolkit-base=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
+    libnvidia-container-tools=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
+    libnvidia-container1=${NVIDIA_CONTAINER_TOOLKIT_VERSION}
+
+# Configure the NVIDIA Container Toolkit
+sudo nvidia-ctk runtime configure --runtime=docker
+sudo systemctl restart docker
+
+# Install Azure CLI
+curl -sLS https://packages.microsoft.com/keys/microsoft.asc |
+  gpg --dearmor | sudo tee /etc/apt/keyrings/microsoft.gpg > /dev/null
+sudo chmod go+r /etc/apt/keyrings/microsoft.gpg
+AZ_DIST=$(lsb_release -cs)
+echo "Types: deb
+URIs: https://packages.microsoft.com/repos/azure-cli/
+Suites: ${AZ_DIST}
+Components: main
+Architectures: $(dpkg --print-architecture)
+Signed-by: /etc/apt/keyrings/microsoft.gpg" | sudo tee /etc/apt/sources.list.d/azure-cli.sources
+sudo apt-get update
+sudo apt-get install -y azure-cli
+
 # Disable the periodical apt-get upgrade.
 # Sometimes, unattended upgrade blocks apt-get install
 sudo sed -i -e "s/Update-Package-Lists \"1\"/Update-Package-Lists \"0\"/g" /etc/apt/apt.conf.d/10periodic
@@ -8,7 +8,7 @@ import openai


 def main() -> None:
-    client = openai.OpenAI()
+    client = openai.OpenAI(timeout=30.0)
    models = client.models.list()
    print("Available models:", models)

@@ -0,0 +1,9 @@
+# MongoDB Development Setup
+
+This script is used to setup MongoDB for development.
+
+## Usage
+
+```bash
+docker compose up -d
+```
@@ -0,0 +1,10 @@
+services:
+  mongo:
+    image: mongo:latest
+    container_name: mongo-dev
+    ports:
+      - "27017:27017"
+    command: ["mongod", "--bind_ip_all", "--replSet", "rs0"]
+    volumes:
+      - ./data:/data/db
+      - ./init-rs.js:/docker-entrypoint-initdb.d/init-rs.js:ro
@@ -0,0 +1,6 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+rs.initiate({
+  _id: "rs0",
+  members: [{ _id: 0, host: "localhost:27017" }],
+});
@@ -1,14 +1,34 @@
 # Copyright (c) Microsoft. All rights reserved.

+import argparse
 import sys

 import wandb

-if len(sys.argv) != 3:
-    print("Usage: python validate_example_wandb.py <project> <run_name>")

-project = sys.argv[1]
-run_name = sys.argv[2]
+def parse_args():
+    parser = argparse.ArgumentParser(description="Validate a Weights & Biases run for reward/trace rollouts.")
+    parser.add_argument("project", help="W&B project name")
+    parser.add_argument("run_name", help="W&B run display name")
+    parser.add_argument(
+        "--reward-tolerance",
+        type=int,
+        default=0,
+        help="Allowed difference between first and last val/n_rollouts_w_reward",
+    )
+    parser.add_argument(
+        "--trace-tolerance",
+        type=int,
+        default=0,
+        help="Allowed difference between first and last val/n_rollouts_w_trace",
+    )
+    return parser.parse_args()
+
+
+args = parse_args()
+
+project = args.project
+run_name = args.run_name
 api = wandb.Api()
 entity_name = api.default_entity
 print("Default entity:", entity_name)
@@ -24,18 +44,56 @@ else:
    print(f"::error::Run with name '{run_name}' not found in project '{project}'.")
    sys.exit(1)

-hist = run.history(keys=["val/reward"], pandas=True)
+hist = run.history(keys=["val/reward", "val/n_rollouts_w_reward", "val/n_rollouts_w_trace"], pandas=True)
 print("History:", hist)
 if hist.empty:
    print("::error::No history found for the run.")
    sys.exit(1)
 else:
-    first, last = hist["val/reward"].iloc[0], hist["val/reward"].iloc[-1]
-    if last <= first:
+    # Check whether all rollouts have (approximately) succeeded
+    first_row = hist.iloc[0]
+    last_row = hist.iloc[-1]
+
+    first_reward_rollouts = first_row["val/n_rollouts_w_reward"]
+    last_reward_rollouts = last_row["val/n_rollouts_w_reward"]
+    reward_diff = abs(first_reward_rollouts - last_reward_rollouts)
+
+    if reward_diff > args.reward_tolerance or (first_reward_rollouts == 0 and last_reward_rollouts == 0):
        print(
-            f"::warning title=Training no improvement::No improvement (run_name={run_name} start={first:.4f}, end={last:.4f})"
+            "::error::Some rollouts have failed to produce rewards: "
+            f"{first_reward_rollouts} -> {last_reward_rollouts} "
+            f"(tolerance={args.reward_tolerance})"
+        )
+        sys.exit(1)
+    elif first_reward_rollouts != last_reward_rollouts:
+        print(
+            "::warning::First and last val/n_rollouts_w_reward are different: "
+            f"{first_reward_rollouts} -> {last_reward_rollouts}"
+        )
+
+    first_trace_rollouts = first_row["val/n_rollouts_w_trace"]
+    last_trace_rollouts = last_row["val/n_rollouts_w_trace"]
+    trace_diff = abs(first_trace_rollouts - last_trace_rollouts)
+
+    if trace_diff > args.trace_tolerance or (first_trace_rollouts == 0 and last_trace_rollouts == 0):
+        print(
+            "::error::Some rollouts have failed to produce traces: "
+            f"{first_trace_rollouts} -> {last_trace_rollouts} "
+            f"(tolerance={args.trace_tolerance})"
+        )
+        sys.exit(1)
+    elif first_trace_rollouts != last_trace_rollouts:
+        print(
+            "::warning::First and last val/n_rollouts_w_trace are different: "
+            f"{first_trace_rollouts} -> {last_trace_rollouts}"
+        )
+
+    first_reward, last_reward = first_row["val/reward"], last_row["val/reward"]
+    if last_reward <= first_reward:
+        print(
+            f"::warning title=Training no improvement::No improvement (run_name={run_name} start={first_reward:.4f}, end={last_reward:.4f})"
        )
    else:
        print(
-            f"::notice title=Training completed::Run has improved (run_name={run_name} start={first:.4f}, end={last:.4f})"
+            f"::notice title=Training completed::Run has improved (run_name={run_name} start={first_reward:.4f}, end={last_reward:.4f})"
        )
@@ -2,7 +2,7 @@

 # pyright: reportPrivateUsage=false

-from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, cast
+from typing import Any, Dict, Iterator, List, Literal, Optional, Sequence, Tuple, cast
 from unittest.mock import AsyncMock, Mock

 import pytest
@@ -32,13 +32,13 @@ class DummyTraceMessagesAdapter(TraceToMessages):
        super().__init__()
        self.seen_spans: Sequence[Span] | None = None

-    def adapt(self, source: List[Span], /) -> List[Dict[str, Any]]:  # type: ignore[override]
+    def adapt(self, source: Sequence[Span], /) -> List[Dict[str, Any]]:  # type: ignore[override]
        self.seen_spans = list(source)
        return [dict(payload="converted")]


 class WrongAdapter(TraceAdapter[List[int]]):
-    def adapt(self, source: List[Span], /) -> List[int]:
+    def adapt(self, source: Sequence[Span], /) -> List[int]:
        return [len(source)]


@@ -79,7 +79,12 @@ class DummyStore:
            return self.wait_results_queue.pop(0)
        return []

-    async def query_spans(self, rollout_id: str) -> List[Span]:
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+        **_: Any,
+    ) -> List[Span]:
        return list(self.query_spans_map.get(rollout_id, []))


@@ -4,7 +4,7 @@ import asyncio
 import logging
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Sequence

 import pytest

@@ -24,7 +24,7 @@ LOGGER_NAME = "agentlightning.algorithm.fast"


 class _AdapterStub(TraceAdapter[Dict[str, Any]]):
-    def adapt(self, source: List[Span], /) -> Dict[str, Any]:
+    def adapt(self, source: Sequence[Span], /) -> Dict[str, Any]:
        return {
            "count": len(source),
            "attempt_ids": sorted({span.attempt_id for span in source}),
@@ -22,3 +22,9 @@
 {"request": {"messages": [{"content": "Return 1.0 if the answer is 8, else 0.0.", "role": "system"}, {"role": "user", "content": "8"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"response": {"title": "Response", "type": "number"}}, "required": ["response"], "title": "OutputType", "type": "object", "additionalProperties": false}}}, "stream": false}, "response": {"id": "chatcmpl-BluaMcYjJNKVhfvlLVLuexmE1srjB", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "created": 1750758630, "choices": [{"index": 0, "message": {"role": "assistant", "content": "{\"response\":1.0}"}, "finish_reason": "stop"}], "system_fingerprint": "fp_178c8d546f", "usage": {"prompt_tokens": 65, "completion_tokens": 8, "total_tokens": 73, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}}}}
 {"request": {"messages": [{"content": "If the question is about math, handoff to MathAgent. Otherwise, handoff to HistoryAgent.", "role": "system"}, {"role": "user", "content": "Who was the first president of the US?"}], "model": "gpt-4.1-mini", "stream": false, "tools": [{"type": "function", "function": {"name": "transfer_to_mathagent", "description": "Handoff to the MathAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}, {"type": "function", "function": {"name": "transfer_to_historyagent", "description": "Handoff to the HistoryAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}]}, "response": {"id": "chatcmpl-BluaNLVzmIYmNkUGGuk7iol8HHBkd", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "created": 1750758631, "choices": [{"index": 0, "message": {"role": "assistant", "tool_calls": [{"id": "call_j9TL7tbHC4v6OpqD66g3k6dL", "type": "function", "function": {"name": "transfer_to_historyagent", "arguments": "{}"}}]}, "finish_reason": "tool_calls"}], "system_fingerprint": "fp_178c8d546f", "usage": {"prompt_tokens": 97, "completion_tokens": 13, "total_tokens": 110, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}}}}
 {"request": {"messages": [{"content": "Answer history questions.", "role": "system"}, {"role": "user", "content": "Who was the first president of the US?"}, {"role": "assistant", "tool_calls": [{"id": "call_j9TL7tbHC4v6OpqD66g3k6dL", "type": "function", "function": {"name": "transfer_to_historyagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_j9TL7tbHC4v6OpqD66g3k6dL", "content": "{\"assistant\": \"HistoryAgent\"}"}], "model": "gpt-4.1-mini", "stream": false}, "response": {"id": "chatcmpl-BluaNHD93ybfyjMqAiF4HqKS93GPf", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "created": 1750758631, "choices": [{"index": 0, "message": {"role": "assistant", "content": "The first president of the United States was George Washington. He served as president from 1789 to 1797."}, "finish_reason": "stop"}], "system_fingerprint": "fp_178c8d546f", "usage": {"prompt_tokens": 53, "completion_tokens": 25, "total_tokens": 78, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}}}}
+{"request": {"messages": [{"content": "If the question is about math, handoff to MathAgent. Otherwise, handoff to HistoryAgent.", "role": "system"}, {"role": "user", "content": "What is 3+5?"}], "model": "gpt-4.1-mini", "tools": [{"type": "function", "function": {"name": "transfer_to_mathagent", "description": "Handoff to the MathAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}, {"type": "function", "function": {"name": "transfer_to_historyagent", "description": "Handoff to the HistoryAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}]}, "response": {"choices": [{"content_filter_results": {}, "finish_reason": "tool_calls", "index": 0, "logprobs": null, "message": {"annotations": [], "content": null, "refusal": null, "role": "assistant", "tool_calls": [{"function": {"arguments": "{}", "name": "transfer_to_mathagent"}, "id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "type": "function"}]}}], "created": 1763565690, "id": "chatcmpl-CdeHqLyOC89BDkmhvbxnhuYa0MV9J", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "jailbreak": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 13, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 95, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 108}}}
+{"request": {"messages": [{"content": "Add two numbers.", "role": "system"}, {"role": "assistant", "content": "For context, here is the conversation so far between the user and the previous agent:\n<CONVERSATION HISTORY>\n1. user: What is 3+5?\n2. function_call: {\"arguments\": \"{}\", \"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"name\": \"transfer_to_mathagent\", \"id\": \"__fake_id__\"}\n3. function_call_output: {\"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"output\": \"{\\\"assistant\\\": \\\"MathAgent\\\"}\"}\n</CONVERSATION HISTORY>"}, {"role": "assistant", "tool_calls": [{"id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "type": "function", "function": {"name": "transfer_to_mathagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "content": "{\"assistant\": \"MathAgent\"}"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"answer": {"title": "Answer", "type": "integer"}}, "required": ["answer"], "title": "MathOutput", "type": "object", "additionalProperties": false}}}, "tools": [{"type": "function", "function": {"name": "add", "description": "", "parameters": {"properties": {"a": {"title": "A", "type": "integer"}, "b": {"title": "B", "type": "integer"}}, "required": ["a", "b"], "title": "add_args", "type": "object", "additionalProperties": false}}}]}, "response": {"choices": [{"content_filter_results": {}, "finish_reason": "tool_calls", "index": 0, "logprobs": null, "message": {"annotations": [], "content": null, "refusal": null, "role": "assistant", "tool_calls": [{"function": {"arguments": "{\"a\":3,\"b\":5}", "name": "add"}, "id": "call_l5lgStmtMUiYuqU0piPW7tY8", "type": "function"}]}}], "created": 1763565694, "id": "chatcmpl-CdeHuHa6HcAsIlClzRimaHMtoq77q", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 18, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 253, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 271}}}
+{"request": {"messages": [{"content": "Add two numbers.", "role": "system"}, {"role": "assistant", "content": "For context, here is the conversation so far between the user and the previous agent:\n<CONVERSATION HISTORY>\n1. user: What is 3+5?\n2. function_call: {\"arguments\": \"{}\", \"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"name\": \"transfer_to_mathagent\", \"id\": \"__fake_id__\"}\n3. function_call_output: {\"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"output\": \"{\\\"assistant\\\": \\\"MathAgent\\\"}\"}\n</CONVERSATION HISTORY>"}, {"role": "assistant", "tool_calls": [{"id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "type": "function", "function": {"name": "transfer_to_mathagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "content": "{\"assistant\": \"MathAgent\"}"}, {"role": "assistant", "tool_calls": [{"id": "call_l5lgStmtMUiYuqU0piPW7tY8", "type": "function", "function": {"name": "add", "arguments": "{\"a\":3,\"b\":5}"}}]}, {"role": "tool", "tool_call_id": "call_l5lgStmtMUiYuqU0piPW7tY8", "content": "8"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"answer": {"title": "Answer", "type": "integer"}}, "required": ["answer"], "title": "MathOutput", "type": "object", "additionalProperties": false}}}, "tools": [{"type": "function", "function": {"name": "add", "description": "", "parameters": {"properties": {"a": {"title": "A", "type": "integer"}, "b": {"title": "B", "type": "integer"}}, "required": ["a", "b"], "title": "add_args", "type": "object", "additionalProperties": false}}}]}, "response": {"choices": [{"content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "protected_material_code": {"filtered": false, "detected": false}, "protected_material_text": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}, "finish_reason": "stop", "index": 0, "logprobs": null, "message": {"annotations": [], "content": "{\"answer\":8}", "refusal": null, "role": "assistant"}}], "created": 1763565699, "id": "chatcmpl-CdeHz4XyD8FskC2OsHRrawVXbUk6A", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 11, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 278, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 289}}}
+{"request": {"messages": [{"content": "Return 1.0 if the answer is 8, else 0.0.", "role": "system"}, {"role": "user", "content": "8"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"response": {"title": "Response", "type": "number"}}, "required": ["response"], "title": "OutputType", "type": "object", "additionalProperties": false}}}}, "response": {"choices": [{"content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "protected_material_code": {"filtered": false, "detected": false}, "protected_material_text": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}, "finish_reason": "stop", "index": 0, "logprobs": null, "message": {"annotations": [], "content": "{\"response\":1.0}", "refusal": null, "role": "assistant"}}], "created": 1763565701, "id": "chatcmpl-CdeI1ifu2oOt06oTodWPpSpNFAqIN", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "jailbreak": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 8, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 65, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 73}}}
+{"request": {"messages": [{"content": "If the question is about math, handoff to MathAgent. Otherwise, handoff to HistoryAgent.", "role": "system"}, {"role": "user", "content": "Who was the first president of the US?"}], "model": "gpt-4.1-mini", "tools": [{"type": "function", "function": {"name": "transfer_to_mathagent", "description": "Handoff to the MathAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}, {"type": "function", "function": {"name": "transfer_to_historyagent", "description": "Handoff to the HistoryAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}]}, "response": {"choices": [{"content_filter_results": {}, "finish_reason": "tool_calls", "index": 0, "logprobs": null, "message": {"annotations": [], "content": null, "refusal": null, "role": "assistant", "tool_calls": [{"function": {"arguments": "{}", "name": "transfer_to_historyagent"}, "id": "call_tXXmrsNYcFlKrG9hMo2Yk0E3", "type": "function"}]}}], "created": 1763565706, "id": "chatcmpl-CdeI6hxlaAwzbn7Sx4ZTXZs0Hx3E5", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "jailbreak": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 13, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 97, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 110}}}
+{"request": {"messages": [{"content": "Answer history questions.", "role": "system"}, {"role": "assistant", "content": "For context, here is the conversation so far between the user and the previous agent:\n<CONVERSATION HISTORY>\n1. user: Who was the first president of the US?\n2. function_call: {\"arguments\": \"{}\", \"call_id\": \"call_tXXmrsNYcFlKrG9hMo2Yk0E3\", \"name\": \"transfer_to_historyagent\", \"id\": \"__fake_id__\"}\n3. function_call_output: {\"call_id\": \"call_tXXmrsNYcFlKrG9hMo2Yk0E3\", \"output\": \"{\\\"assistant\\\": \\\"HistoryAgent\\\"}\"}\n</CONVERSATION HISTORY>"}, {"role": "assistant", "tool_calls": [{"id": "call_tXXmrsNYcFlKrG9hMo2Yk0E3", "type": "function", "function": {"name": "transfer_to_historyagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_tXXmrsNYcFlKrG9hMo2Yk0E3", "content": "{\"assistant\": \"HistoryAgent\"}"}], "model": "gpt-4.1-mini"}, "response": {"choices": [{"content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "protected_material_text": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}, "finish_reason": "stop", "index": 0, "logprobs": null, "message": {"annotations": [], "content": "The first president of the United States was George Washington. He served as president from 1789 to 1797. If you have more questions about U.S. history or any other historical events, feel free to ask!", "refusal": null, "role": "assistant"}}], "created": 1763565707, "id": "chatcmpl-CdeI78NuMQmjmeuwLhx3I0Qqcem99", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 46, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 179, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 225}}}
@@ -1152,7 +1152,7 @@ def test_execute_both_main_algo_runner_ignores_stop(store: LightningStore) -> No
        await asyncio.sleep(0.1)

    # Should complete without hanging despite runners ignoring signals
-    with pytest.raises(RuntimeError, match="Subprocesses failed:"):
+    with pytest.raises(RuntimeError, match="Subprocesses failed"):
        strat.execute(algorithm=algo, runner=_runner_ignores_stop_forever, store=store)


@@ -17,7 +17,7 @@ def test_switchable_authenticated_exporter():
    switchable_authenticated_exporter = BypassableAuthenticatedOTLPExporter(endpoint="http://dummy", jwt="dummy")

    with patch.object(
-        switchable_authenticated_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
+        switchable_authenticated_exporter.__class__.__bases__[-1], "export", return_value=SpanExportResult.SUCCESS
    ) as mock_export:
        enable_agentops_service()
        result = switchable_authenticated_exporter.export([])
@@ -34,7 +34,7 @@ def test_switchable_otlp_metric_exporter():

    switchable_otlp_metric_exporter = BypassableOTLPMetricExporter()
    with patch.object(
-        switchable_otlp_metric_exporter.__class__.__bases__[0], "export", return_value=MetricExportResult.SUCCESS
+        switchable_otlp_metric_exporter.__class__.__bases__[-1], "export", return_value=MetricExportResult.SUCCESS
    ) as mock_export:
        enable_agentops_service()
        result = switchable_otlp_metric_exporter.export(metrics_data=MagicMock())
@@ -51,7 +51,10 @@ def test_switchable_otlp_span_exporter():

    switchable_otlp_span_exporter = BypassableOTLPSpanExporter()
    with patch.object(
-        switchable_otlp_span_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
+        # BypassableOTLPSpanExporter is a subclass of LightningStoreOTLPExporter, which is a subclass of OTLPSpanExporter
+        switchable_otlp_span_exporter.__class__.__bases__[-1].__bases__[0],
+        "export",
+        return_value=SpanExportResult.SUCCESS,
    ) as mock_export:
        enable_agentops_service()
        result = switchable_otlp_span_exporter.export([])
@@ -2,18 +2,23 @@

 import asyncio
 import logging
+import multiprocessing
 import random
 from typing import Any, List, cast

 import litellm
 import openai
+import opentelemetry.trace as trace_api
 import pytest
+from agentops.sdk.core import BatchSpanProcessor
 from litellm.llms.custom_llm import CustomLLM
 from litellm.types.utils import ModelResponse
 from litellm.utils import custom_llm_setup
 from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor

 from agentlightning.llm_proxy import LightningSpanExporter, LLMProxy
+from agentlightning.store import LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.store.threading import LightningStoreThreaded
 from agentlightning.types import Span
@@ -282,7 +287,6 @@ async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixt
            assert response.choices[0].message.content == f"Hi! {restart_idx}"

            error_logs = [record.message for record in caplog.records if record.levelno >= logging.ERROR]
-            error_logs = [message for message in error_logs if "Task was destroyed but it is pending!" not in message]
            assert not error_logs, f"Found error logs: {error_logs}"
            assert not any("Cannot add callback" in record.message for record in caplog.records)

@@ -290,3 +294,87 @@ async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixt
    finally:
        litellm.custom_provider_map = []
        custom_llm_setup()
+
+
+async def llm_proxy_span_exporter_loop(otlp_enabled: bool = False):
+    store = LightningStoreThreaded(InMemoryLightningStore())
+
+    if otlp_enabled:
+        store = LightningStoreServer(store, "127.0.0.1", get_free_port())
+        await store.start()
+
+    llm_instance = TestLLM(f"Hi! I'm a test LLM")
+    litellm.custom_provider_map = [{"provider": "test-llm", "custom_handler": llm_instance}]
+    custom_llm_setup()
+    proxy = LLMProxy(
+        launcher_args=PythonServerLauncherArgs(
+            launch_mode="thread",
+            healthcheck_url="/health",
+            port=get_free_port(),
+        ),
+        store=store,
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "test-llm/any-llm",
+                },
+            }
+        ],
+    )
+    await proxy.start()
+
+    rollout = await store.start_rollout(None)
+    resource = proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id)
+
+    client = openai.AsyncOpenAI(
+        base_url=resource.endpoint,
+        api_key="token-abc123",
+        timeout=5,
+        max_retries=0,
+    )
+    response = await client.chat.completions.create(
+        model="gpt-4o-arbitrary",
+        messages=[{"role": "user", "content": "Hello world"}],
+        stream=False,
+    )
+    assert response.choices[0].message.content == "Hi! I'm a test LLM"
+
+    spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+    assert len(spans) > 0, "Should have captured spans"
+    for span in spans:
+        assert span.rollout_id == rollout.rollout_id, f"Span {span.name} has incorrect rollout_id"
+        assert span.attempt_id == rollout.attempt.attempt_id, f"Span {span.name} has incorrect attempt_id"
+        assert span.sequence_id == 1, f"Span {span.name} has incorrect sequence_id"
+
+    tracer_provider = trace_api.get_tracer_provider()
+
+    have_asserted_loop = False
+    for span_processor in tracer_provider._active_span_processor._span_processors:  # type: ignore
+        if isinstance(span_processor, (SimpleSpanProcessor, BatchSpanProcessor)):
+            if isinstance(span_processor.span_exporter, LightningSpanExporter):
+                if otlp_enabled:
+                    assert span_processor.span_exporter._loop is None  # type: ignore
+                else:
+                    assert span_processor.span_exporter._loop is not None  # type: ignore
+                have_asserted_loop = True
+                break
+    assert have_asserted_loop, f"LightningSpanExporter should be used with otlp_enabled={otlp_enabled}"
+
+    await proxy.stop()
+
+    if isinstance(store, LightningStoreServer):
+        await store.stop()
+
+
+def llm_proxy_span_exporter_loop_sync(otlp_enabled: bool = False):
+    asyncio.run(llm_proxy_span_exporter_loop(otlp_enabled))
+
+
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+def test_llm_proxy_span_exporter_loop(otlp_enabled: bool):
+    context = multiprocessing.get_context("spawn")
+    process = context.Process(target=llm_proxy_span_exporter_loop_sync, args=(otlp_enabled,))
+    process.start()
+    process.join(timeout=30.0)
+    assert process.exitcode == 0
@@ -15,20 +15,20 @@ There are some specific TODOs for each test function.
 import ast
 import asyncio
 import json
-from typing import Any, Dict, List, Type, Union, cast
+from typing import Any, Dict, List, Sequence, Type, Union, cast

 import anthropic
 import openai
 import pytest
 from litellm.integrations.custom_logger import CustomLogger
+from portpicker import pick_unused_port

 from agentlightning import LlmProxyTraceToTriplet
 from agentlightning.llm_proxy import LLMProxy, _reset_litellm_logging_worker  # pyright: ignore[reportPrivateUsage]
-from agentlightning.store import LightningStore, LightningStoreServer
+from agentlightning.store import LightningStore, LightningStoreServer, LightningStoreThreaded
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.types import LLM, Span

-from ..common.network import get_free_port
 from ..common.tracer import clear_tracer_provider
 from ..common.vllm import VLLM_VERSION, RemoteOpenAIServer

@@ -52,7 +52,7 @@ def qwen25_model():
            "--tool-call-parser",
            "hermes",
            "--port",
-            str(get_free_port()),
+            str(pick_unused_port()),
        ],
    ) as server:
        yield server
@@ -69,13 +69,17 @@ def test_qwen25_model_sanity(qwen25_model: RemoteOpenAIServer):


@pytest.mark.asyncio
-async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_basic_integration(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
    clear_tracer_provider()
    inmemory_store = InMemoryLightningStore()
-    store = LightningStoreServer(store=inmemory_store, host="127.0.0.1", port=get_free_port())
-    await store.start()
+    if otlp_enabled:
+        store = LightningStoreServer(store=inmemory_store, host="127.0.0.1", port=pick_unused_port())
+        await store.start()
+    else:
+        store = LightningStoreThreaded(inmemory_store)
    proxy = LLMProxy(
-        port=get_free_port(),
+        port=pick_unused_port(),
        model_list=[
            {
                "model_name": "gpt-4o-arbitrary",
@@ -86,6 +90,7 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
            }
        ],
        store=store,
+        launch_mode="thread" if not otlp_enabled else "mp",
    )

    rollout = await store.start_rollout(None)
@@ -107,7 +112,8 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):

    spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)

-    await store.stop()
+    if isinstance(store, LightningStoreServer):
+        await store.stop()

    # Verify all spans have correct rollout_id, attempt_id, and sequence_id
    assert len(spans) > 0, "Should have captured spans"
@@ -116,6 +122,14 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
        assert span.attempt_id == rollout.attempt.attempt_id, f"Span {span.name} has incorrect attempt_id"
        assert span.sequence_id == 1, f"Span {span.name} has incorrect sequence_id"

+        # Verify start time and end time
+        # TODO: Remove this when this PR is merged: https://github.com/BerriAI/litellm/pull/16558
+        print(f">>> Span: {span.name}")
+        print(f">>> Start time: {span.start_time}")
+        print(f">>> End time: {span.end_time}")
+        assert span.start_time is not None, f"Span {span.name} has no start time"
+        assert span.end_time is not None, f"Span {span.name} has no end time"
+
    # Find the raw_gen_ai_request span and verify token IDs
    raw_gen_ai_spans = [s for s in spans if s.name == "raw_gen_ai_request"]
    assert len(raw_gen_ai_spans) == 1, f"Expected 1 raw_gen_ai_request span, found {len(raw_gen_ai_spans)}"
@@ -179,13 +193,18 @@ async def _make_proxy_and_store(
    retries: int = 0,
    gunicorn: bool = False,
    callbacks: List[Union[Type[CustomLogger], str]] | None = None,
+    otlp_enabled: bool = False,
 ):
    clear_tracer_provider()
    _reset_litellm_logging_worker()  # type: ignore
    store = InMemoryLightningStore()
-    store_server = LightningStoreServer(store=store, host="127.0.0.1", port=get_free_port())
-    # When the server is forked into subprocess, it automatically becomes a client of the store
-    await store_server.start()
+    if otlp_enabled:
+        store = LightningStoreServer(store=store, host="127.0.0.1", port=pick_unused_port())
+        # When the server is forked into subprocess, it automatically becomes a client of the store
+        await store.start()
+    else:
+        # Backward compatibility with legacy thread + non-otlp mode
+        store = LightningStoreThreaded(store)
    proxy = LLMProxy(
        model_list=[
            {
@@ -196,14 +215,15 @@ async def _make_proxy_and_store(
                },
            }
        ],
-        port=get_free_port(),
+        launch_mode="thread" if not otlp_enabled else "mp",
+        port=pick_unused_port(),
        num_workers=4 if gunicorn else 1,
-        store=store_server,
+        store=store,
        num_retries=retries,
        callbacks=callbacks,
    )
    await proxy.start()
-    return proxy, store_server
+    return proxy, store


 async def _new_resource(proxy: LLMProxy, store: LightningStore):
@@ -219,7 +239,7 @@ def _get_async_client_for_resource(resource: LLM):
    return openai.AsyncOpenAI(base_url=resource.endpoint, api_key="token-abc123", timeout=120, max_retries=0)


-def _find_span(spans: list[Span], name: str):
+def _find_span(spans: Sequence[Span], name: str):
    return [s for s in spans if s.name == name]


@@ -228,8 +248,9 @@ def _attr(s: Span, key: str, default: Any = None):  # type: ignore


@pytest.mark.asyncio
-async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer):
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
    try:
        resource, rollout = await _new_resource(proxy, store)
        client = _get_client_for_resource(resource)
@@ -251,13 +272,14 @@ async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer):
        # TODO: Check response contents and token ids for the 3 requests respectively
    finally:
        await proxy.stop()
-        await store.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()


@pytest.mark.asyncio
-@pytest.mark.parametrize("gunicorn", [False, True])
-async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, gunicorn: bool):
-    proxy, store = await _make_proxy_and_store(qwen25_model, gunicorn=gunicorn)
+@pytest.mark.parametrize("mode", ["gunicorn", "thread", "uvicorn"])
+async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, mode: str):
+    proxy, store = await _make_proxy_and_store(qwen25_model, gunicorn=mode == "gunicorn", otlp_enabled=mode != "thread")
    try:
        resource, rollout = await _new_resource(proxy, store)
        aclient = _get_async_client_for_resource(resource)
@@ -280,13 +302,15 @@ async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, gunicor
        # TODO: Check whether the sequence ids get mixed up or not
    finally:
        await proxy.stop()
-        await store.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()


@pytest.mark.asyncio
-async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer):
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
    # litellm proxy accepts Anthropic schema and forwards to OpenAI backend
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
    try:
        resource, rollout = await _new_resource(proxy, store)

@@ -304,12 +328,14 @@ async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer):
        assert len(spans) > 0
    finally:
        await proxy.stop()
-        await store.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()


@pytest.mark.asyncio
-async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
    try:
        resource, rollout = await _new_resource(proxy, store)
        client = _get_client_for_resource(resource)
@@ -373,12 +399,14 @@ async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
        # TODO: Check response contents and token ids for the 2 requests respectively
    finally:
        await proxy.stop()
-        await store.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()


@pytest.mark.asyncio
-async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
    try:
        resource, rollout = await _new_resource(proxy, store)
        client = _get_client_for_resource(resource)
@@ -395,7 +423,8 @@ async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
                if c.delta and getattr(c.delta, "content", None):
                    assert isinstance(c.delta.content, str)
                    collected.append(c.delta.content)
-        assert "apple" in "".join(collected).lower()
+        # Sometimes the model responds with "hello" instead of "apple"
+        assert "apple" in "".join(collected).lower() or "hello" in "".join(collected).lower()

        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
        assert len(spans) > 0
@@ -408,12 +437,14 @@ async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
                assert "gen_ai.completion.0.content" in span.attributes
    finally:
        await proxy.stop()
-        await store.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()


@pytest.mark.asyncio
-async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer):
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
+    proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
    try:
        resource, rollout = await _new_resource(proxy, store)
        adapter = LlmProxyTraceToTriplet()
@@ -470,7 +501,8 @@ async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer):
        assert len(triplets) == 2
    finally:
        await proxy.stop()
-        await store.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()


 class LogprobsCallback(CustomLogger):
@@ -480,9 +512,10 @@ class LogprobsCallback(CustomLogger):


@pytest.mark.asyncio
-async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer):
+@pytest.mark.parametrize("otlp_enabled", [True, False])
+async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
    proxy, store = await _make_proxy_and_store(
-        qwen25_model, callbacks=[LogprobsCallback, "return_token_ids", "opentelemetry"]
+        qwen25_model, callbacks=[LogprobsCallback, "return_token_ids", "opentelemetry"], otlp_enabled=otlp_enabled
    )
    try:
        resource, rollout = await _new_resource(proxy, store)
@@ -534,4 +567,5 @@ async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer):
            # TODO: Check logprobs
    finally:
        await proxy.stop()
-        await store.stop()
+        if isinstance(store, LightningStoreServer):
+            await store.stop()
@@ -1,40 +1,72 @@
 # Copyright (c) Microsoft. All rights reserved.

+from __future__ import annotations
+
+import os
 import time
+from itertools import count
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Sequence
 from unittest.mock import Mock
+from uuid import uuid4

 import pytest
+import pytest_asyncio
 from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Field
 from pytest import FixtureRequest

 from agentlightning.store.base import LightningStore
+from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, KeyValue, ListBasedCollection, Queue
+from agentlightning.store.collection.base import Collection
 from agentlightning.store.memory import InMemoryLightningStore

+if TYPE_CHECKING:
+    from pymongo import AsyncMongoClient
+    from pymongo.asynchronous.database import AsyncDatabase
+
 __all__ = [
    "inmemory_store",
    "mock_readable_span",
+    "sample_items",
+    "sample_collection",
+    "SampleItem",
+    "QueueItem",
+    "deque_queue",
+    "dict_key_value",
+    "dict_key_value_data",
+    "temporary_mongo_database",
 ]


+mongo_uri = os.getenv("AGL_TEST_MONGO_URI", "mongodb://localhost:27017/?replicaSet=rs0")
+
+
@pytest.fixture
 def inmemory_store() -> InMemoryLightningStore:
    """Create a fresh InMemoryLightningStore instance."""
    return InMemoryLightningStore()


-@pytest.fixture
-def sql_store():
-    """Placeholder fixture for SQL store implementation. Returns None until SQL store is ready."""
-    return None
+@pytest_asyncio.fixture
+async def mongo_store(temporary_mongo_database: AsyncDatabase[Any]):
+    """Fixture for MongoDB store implementation."""
+    from agentlightning.store.mongo import MongoLightningStore
+
+    db = MongoLightningStore(client=temporary_mongo_database.client, database_name=temporary_mongo_database.name)
+    try:
+        yield db
+    finally:
+        await db.close()


-# Uncomment this when sql store is ready
-# @pytest.fixture(params=["inmemory_store", "sql_store"])
-@pytest.fixture(params=["inmemory_store"])
-def store_fixture(request: FixtureRequest) -> LightningStore:
-    """Parameterized fixture that provides different store implementations for testing.
-    Currently supports InMemoryLightningStore, with SQL store support planned.
-    """
+@pytest.fixture(
+    params=[
+        "inmemory_store",
+        pytest.param("mongo_store", marks=pytest.mark.mongo),
+    ]
+)
+def store_fixture(request: FixtureRequest) -> AsyncGenerator[LightningStore, None]:
+    """Parameterized fixture that provides different store implementations for testing."""
    return request.getfixturevalue(request.param)


@@ -43,14 +75,20 @@ def mock_readable_span() -> ReadableSpan:
    """Create a mock ReadableSpan for testing."""
    span = Mock()
    span.name = "test_span"
+    context_counter = count(1)
+
+    def _make_context() -> Mock:
+        """Generate a distinct span context each time it is requested."""
+        index = next(context_counter)
+        context = Mock()
+        context.trace_id = 111111
+        context.span_id = 222222 + index
+        context.is_remote = False
+        context.trace_state = {}
+        return context

    # Mock context
-    context = Mock()
-    context.trace_id = 111111
-    context.span_id = 222222
-    context.is_remote = False
-    context.trace_state = {}  # Make it an empty dict instead of Mock
-    span.get_span_context = Mock(return_value=context)
+    span.get_span_context = Mock(side_effect=_make_context)

    # Mock other attributes
    span.parent = None
@@ -66,3 +104,269 @@ def mock_readable_span() -> ReadableSpan:
    span.resource = Mock(attributes={}, schema_url="")

    return span
+
+
+class SampleItem(BaseModel):
+    partition: str
+    index: int
+    name: str
+    status: str
+    tags: List[str] = Field(default_factory=list)
+    score: float | None = None
+    rank: int | None = None
+    updated_time: float | None = None
+    payload: Dict[str, int] = Field(default_factory=dict)
+    metadata: str | None = None
+
+
+class QueueItem(BaseModel):
+    idx: int
+
+
+@pytest_asyncio.fixture
+async def mongo_client():
+    from pymongo import AsyncMongoClient
+
+    client = AsyncMongoClient[Any](mongo_uri, serverSelectionTimeoutMS=5000)
+    try:
+        await client.admin.command("ping")
+    except Exception as exc:  # depends on external service
+        await client.close()
+        raise RuntimeError(f"MongoDB not available: {exc}")
+
+    try:
+        yield client
+    finally:
+        await client.close()
+
+
+@pytest_asyncio.fixture
+async def temporary_mongo_database(mongo_client: AsyncMongoClient[Any]):
+    """Yield a temporary MongoDB database for integration tests."""
+    db_name = f"agentlightning-test-{uuid4().hex}"
+    db = mongo_client[db_name]  # type: ignore
+    try:
+        yield db
+    finally:
+        await mongo_client.drop_database(db_name)
+
+
+### Collection fixtures ###
+
+
+@pytest.fixture()
+def sample_items() -> List[SampleItem]:
+    return [
+        SampleItem(
+            partition="alpha",
+            index=1,
+            name="urgent-phase-one",
+            status="new",
+            tags=["core", "urgent"],
+            score=10.5,
+            rank=3,
+            updated_time=12.0,
+            payload={"priority": 10},
+            metadata="alpha-start",
+        ),
+        SampleItem(
+            partition="alpha",
+            index=2,
+            name="phase-two",
+            status="running",
+            tags=["core"],
+            score=5.0,
+            rank=2,
+            updated_time=None,
+            payload={"priority": 5},
+            metadata=None,
+        ),
+        SampleItem(
+            partition="alpha",
+            index=3,
+            name="delayed-phase",
+            status="blocked",
+            tags=["delayed"],
+            score=None,
+            rank=5,
+            updated_time=15.1,
+            payload={"priority": 8},
+            metadata="delayed-phase",
+        ),
+        SampleItem(
+            partition="beta",
+            index=1,
+            name="beta-critical",
+            status="new",
+            tags=["beta", "urgent"],
+            score=8.0,
+            rank=1,
+            updated_time=7.0,
+            payload={"priority": 7},
+            metadata="beta critical",
+        ),
+        SampleItem(
+            partition="beta",
+            index=2,
+            name="beta optional",
+            status="done",
+            tags=["beta"],
+            score=3.0,
+            rank=None,
+            updated_time=2.0,
+            payload={"priority": 1},
+            metadata="optional path",
+        ),
+        SampleItem(
+            partition="gamma",
+            index=1,
+            name="gamma-phase",
+            status="running",
+            tags=[],
+            score=9.5,
+            rank=4,
+            updated_time=None,
+            payload={"priority": 9},
+            metadata="gamma-phase data",
+        ),
+        SampleItem(
+            partition="gamma",
+            index=2,
+            name="gamma-late",
+            status="done",
+            tags=["late", "core"],
+            score=1.0,
+            rank=6,
+            updated_time=20.0,
+            payload={"priority": 2},
+            metadata="gamma late entry",
+        ),
+        SampleItem(
+            partition="delta",
+            index=1,
+            name="delta misc",
+            status="archived",
+            tags=["misc"],
+            score=4.2,
+            rank=7,
+            updated_time=11.0,
+            payload={"priority": 3},
+            metadata="delta misc block",
+        ),
+    ]
+
+
+### Generic collection fixtures ###
+
+
+@pytest.fixture()
+def sample_collection_memory(sample_items: Sequence[SampleItem]) -> ListBasedCollection[SampleItem]:
+    collection: Collection[SampleItem] = ListBasedCollection(list(sample_items), SampleItem, ("partition", "index"))
+    setattr(collection, "_test_backend", "memory")
+    return collection
+
+
+@pytest_asyncio.fixture
+async def sample_collection_mongo(temporary_mongo_database: AsyncDatabase[Any], sample_items: Sequence[SampleItem]):
+    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        collection = MongoBasedCollection(
+            client_pool,
+            temporary_mongo_database.name,
+            "sample-items",
+            "partition-123",
+            ["partition", "index"],
+            SampleItem,
+        )
+        await collection.insert(sample_items)
+        setattr(collection, "_test_backend", "mongo")
+        yield collection
+
+
+@pytest.fixture(
+    params=[
+        "memory",
+        pytest.param("mongo", marks=pytest.mark.mongo),
+    ]
+)
+def sample_collection(request: pytest.FixtureRequest):
+    backend = request.param
+    return request.getfixturevalue("sample_collection_" + backend)
+
+
+### Generic queue fixtures ###
+
+
+@pytest.fixture
+def deque_queue_memory() -> DequeBasedQueue[QueueItem]:
+    return DequeBasedQueue(QueueItem, [QueueItem(idx=i) for i in range(3)])
+
+
+@pytest_asyncio.fixture
+async def deque_queue_mongo(temporary_mongo_database: AsyncDatabase[Any]):
+    from agentlightning.store.collection.mongo import MongoBasedQueue, MongoClientPool
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        queue = MongoBasedQueue[QueueItem](
+            client_pool,
+            temporary_mongo_database.name,
+            "queue-items",
+            "partition-1",
+            QueueItem,
+        )
+        await queue.enqueue([QueueItem(idx=i) for i in range(3)])
+        yield queue
+
+
+@pytest.fixture(
+    params=[
+        "memory",
+        pytest.param("mongo", marks=pytest.mark.mongo),
+    ]
+)
+def deque_queue(request: pytest.FixtureRequest) -> AsyncGenerator[Queue[QueueItem], None]:
+    backend = request.param
+    return request.getfixturevalue("deque_queue_" + backend)
+
+
+### Generic key-value fixtures ###
+
+
+@pytest.fixture()
+def dict_key_value_data() -> Dict[str, int]:
+    return {"alpha": 1, "beta": 2}
+
+
+@pytest.fixture()
+def dict_key_value_memory(dict_key_value_data: Dict[str, int]) -> DictBasedKeyValue[str, int]:
+    return DictBasedKeyValue(dict_key_value_data)
+
+
+@pytest_asyncio.fixture
+async def dict_key_value_mongo(temporary_mongo_database: AsyncDatabase[Any], dict_key_value_data: Dict[str, int]):
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            "key-value-items",
+            "partition-1",
+            str,
+            int,
+        )
+        for key, value in dict_key_value_data.items():
+            await key_value.set(key, value)
+        yield key_value
+
+
+@pytest.fixture(
+    params=[
+        "memory",
+        pytest.param("mongo", marks=pytest.mark.mongo),
+    ]
+)
+def dict_key_value(request: pytest.FixtureRequest) -> AsyncGenerator[KeyValue[str, int], None]:
+    backend = request.param
+    return request.getfixturevalue("dict_key_value_" + backend)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Yuge Zhang	872c52986e	remove debug info	2025-11-21 10:30:47 +08:00
Yuge Zhang	ccde956743	update tests	2025-11-21 09:58:13 +08:00
Yuge Zhang	8d3f62e8fb	.	2025-11-21 08:47:11 +08:00
Yuge Zhang	d0b4d29088	.	2025-11-21 01:46:01 +08:00
Yuge Zhang	c98027439c	.	2025-11-21 01:45:06 +08:00
Yuge Zhang	6a026bcb68	.	2025-11-21 00:50:05 +08:00
Yuge Zhang	c78d881449	fix client pool close	2025-11-21 00:26:21 +08:00
Yuge Zhang	972e5da20a	debug mongo failure	2025-11-20 22:05:08 +08:00
Yuge Zhang	56f5554caa	.	2025-11-20 21:43:35 +08:00
Yuge Zhang	595b171b86	.	2025-11-20 21:19:53 +08:00
Yuge Zhang	fc63eb58e0	fix import	2025-11-20 20:55:20 +08:00
Yuge Zhang	8bc38e8122	.	2025-11-20 20:29:15 +08:00
Yuge Zhang	12e8a63ae9	Update mongodb tests	2025-11-20 20:06:12 +08:00
Yuge Zhang	c5c93492d7	Merge branch 'main' of github.com:microsoft/agent-lightning into feature/mongo	2025-11-20 20:02:50 +08:00
Yuge Zhang	711dee2c30	exclude in tests	2025-11-20 19:52:02 +08:00
Yuge Zhang	2d65d1717d	setup mongo	2025-11-20 19:45:59 +08:00
Yuge Zhang	1d43ceff60	exclude mongo tests by default	2025-11-20 19:45:07 +08:00
Yuge Zhang	d581cbcd63	Upgrade VM image (#325 )	2025-11-20 17:49:16 +08:00
Yuge Zhang	3459caa1de	Fix OpenAI Agents 0.6 compatibility and pin vLLM < 0.11.1 (#322 )	2025-11-20 07:13:15 +08:00
Yuge Zhang	71d25351b4	setup mongo	2025-11-19 23:47:15 +08:00
Yuge Zhang	6c7405a30e	update tests	2025-11-19 22:17:45 +08:00
Yuge Zhang	f862f75518	fix client server	2025-11-19 22:09:17 +08:00
Yuge Zhang	8c69348b34	minor fix	2025-11-19 20:57:28 +08:00
Yuge Zhang	f3fd58e72a	Put store init in the right place of tracer (#321 )	2025-11-19 20:35:27 +08:00
Yuge Zhang	973b2859c1	fix transient errors	2025-11-19 15:21:05 +08:00
Yuge Zhang	6b90f67f5e	atomic decorator	2025-11-19 15:09:41 +08:00
Yuge Zhang	0fcdd1c940	fix all except transient error	2025-11-19 14:04:02 +08:00
Yuge Zhang	388ea9bfa2	fix collection-based store implementation	2025-11-19 13:37:14 +08:00
Yuge Zhang	f24d48b8eb	fix client pool	2025-11-19 13:25:30 +08:00
Yuge Zhang	eda7187f01	.	2025-11-19 12:47:18 +08:00
Yuge Zhang	51e6bb9982	refactor memory for prep	2025-11-19 09:38:00 +08:00
Yuge Zhang	127ee3253e	add ensure collection	2025-11-19 01:05:53 +08:00
Yuge Zhang	33471ed243	fix mongo tests	2025-11-18 23:48:59 +08:00
Yuge Zhang	b3cb5e1337	Minor improvements to make RL workflow more robust (#319 )	2025-11-18 15:40:51 +08:00
Yuge Zhang	c1a10b43e4	add mongo tests	2025-11-18 15:23:28 +08:00
Yuge Zhang	11dcddcd0f	initiate mongo	2025-11-18 14:05:41 +08:00
Yuge Zhang	cce793cdfc	make size async	2025-11-18 12:50:58 +08:00
Yuge Zhang	3761c0f54c	Support native advanced queries in LightningStore (#318 )	2025-11-18 10:54:55 +08:00
Yuge Zhang	d4334182be	Adding check traces with reward for VERL (#317 )	2025-11-17 21:18:15 +08:00
Yuge Zhang	57c3c0525e	Collection-based Lightning Store (#315 )	2025-11-17 18:51:51 +08:00
Yuge Zhang	e356593f73	Bump to 0.3.0 (#316 )	2025-11-17 17:32:42 +08:00
Yuge Zhang	0e033831d5	Support OTLP in LightningStore (#313 )	2025-11-15 16:34:09 +08:00
xiaochulaoban	0d721228d5	Added the README and script files for training sql_agent on NPU (#272 ) Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>	2025-11-15 01:27:07 +08:00
Yuge Zhang	e49b75b7d8	Check all matching jobs per variant (#310 )	2025-11-13 17:10:50 +00:00
Yuge Zhang	eab691b1a1	Refactor logging (#306 )	2025-11-13 22:48:52 +08:00
Yuge Zhang	fd6494873d	Make health timeout configurable (#305 )	2025-11-13 19:46:02 +08:00
Yuge Zhang	6cbfc1fee0	Fix CI Badge and make Calc-X pipeline faster (#304 )	2025-11-13 18:06:18 +08:00
Yuge Zhang	b986ae132a	Use PythonServerLauncher in LightningStoreServer (#303 )	2025-11-13 14:22:54 +08:00
Yuge Zhang	f24a47969e	Increase graceful timeout on CI (#302 )	2025-11-13 10:15:14 +08:00
Yuge Zhang	a0bc1827d9	[Release] v0.2.2 (#298 )	2025-11-12 23:54:35 +08:00
Yuge Zhang	f2869cea30	Fix local model support in VERL (#299 )	2025-11-12 22:56:10 +08:00