Merge branch 'main' into dev/database-store

Use unified server launcher for LiteLLM Proxy (#292 )
fix lint warning and update uv.lock
2025-11-12 12:02:35 +08:00 · 2025-11-12 02:45:29 +08:00 · 2025-11-11 16:53:01 +08:00 · 2025-11-11 16:29:02 +08:00 · 2025-11-11 00:27:02 +08:00 · 2025-11-10 22:30:00 +08:00
422 changed files with 87618 additions and 4096 deletions
@@ -0,0 +1,32 @@
+name: Backport Merged Pull Request
+on:
+  pull_request_target:
+    types: [closed]
+permissions:
+  contents: write
+  issues: write
+  pull-requests: write
+
+# NOTE:
+# Microsoft requires rotating BOT_PAT every 3 months.
+# Log onto agent-lightning-bot account and rotate the PAT if needed.
+
+jobs:
+  backport:
+    name: Backport pull request
+    runs-on: ubuntu-latest
+    # Don't run on closed unmerged pull requests
+    if: github.event.pull_request.merged
+    steps:
+      - uses: actions/checkout@v4
+      - name: Create backport pull requests
+        uses: korthout/backport-action@v3
+        with:
+          branch_name: 'backport/${pull_number}/${target_branch}'
+          label_pattern: ^(stable/[^ ]+)$
+          github_token: ${{ secrets.BOT_PAT }}
+          add_labels: backport
+          add_author_as_assignee: true
+          git_committer_name: agent-lightning-bot
+          # This email address is not monitored.
+          git_committer_email: agl.msft@outlook.com
@@ -0,0 +1,29 @@
+name: Badge - APO
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - APO
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-apo.yml', label: 'apo', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,29 @@
+name: Badge - Calc-X
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Calc-X
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-calc-x.yml', label: 'calc-x', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,35 @@
+name: Badge - Examples
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Calc-X
+      - Examples - Spider
+      - Examples - APO
+      - Examples - Unsloth
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-calc-x.yml', label: 'examples-calc-x.stable', variants: ['stable'] },
+              { workflow: 'examples-spider.yml', label: 'examples-spider.stable', variants: ['stable'] },
+              { workflow: 'examples-apo.yml', label: 'examples-apo.stable', variants: ['stable'] },
+              { workflow: 'examples-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,37 @@
+name: Badge - Latest
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Calc-X
+      - Examples - Spider
+      - Examples - APO
+      - Examples - Unsloth
+      - GPU Test
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-calc-x.yml', label: 'calc-x.latest', variants: ['latest'] },
+              { workflow: 'examples-spider.yml', label: 'spider.latest', variants: ['latest'] },
+              { workflow: 'examples-apo.yml', label: 'apo.latest', variants: ['latest'] },
+              { workflow: 'examples-unsloth.yml', label: 'unsloth.latest', variants: ['latest'] },
+              { workflow: 'tests-full.yml', label: 'tests-full.latest', variants: ['latest'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,29 @@
+name: Badge - Spider
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Spider
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-spider.yml', label: 'spider', variants: ['stable', 'legacy'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,29 @@
+name: Badge - Unsloth
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Unsloth
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,33 @@
+name: Dashboard
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 5 AM UTC+8
+    - cron: '0 21 * * *'
+
+  workflow_dispatch:
+
+  push:
+    branches: [ main, stable/**/* ]
+
+jobs:
+  dashboard:
+    name: Chromatic
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Run Chromatic
+        uses: chromaui/action@v13
+        with:
+          projectToken: ${{ secrets.CHROMATIC_PROJECT_TOKEN }}
+          workingDir: dashboard
+          exitZeroOnChanges: false
@@ -8,6 +8,10 @@ on:
      - 'v*'
  workflow_dispatch:

+concurrency:
+  group: docs-deploy
+  cancel-in-progress: false
+
 permissions:
  contents: write
  pages: write
@@ -20,15 +24,14 @@ jobs:
      - uses: actions/checkout@v4
        with:
          fetch-depth: 0
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - uses: actions/setup-python@v6
        with:
          python-version: '3.12'
-
-      - name: Install dependencies
-        run: |
-          ./scripts/setup_stable.sh
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev

      - name: Configure Git
        run: |
@@ -51,10 +54,11 @@ jobs:
      - name: Deploy versioned docs
        if: startsWith(github.ref, 'refs/tags/')
        run: |
-          mike deploy --push --update-aliases ${{ steps.version.outputs.version }} stable
+          uv run --locked --no-sync mike deploy --push --update-aliases ${{ steps.version.outputs.version }} stable

      - name: Deploy dev docs
        if: github.ref == 'refs/heads/main'
        run: |
-          mike deploy --push latest
-          mike set-default --push latest
+          uv run --locked --no-sync mike deploy --push latest
+          # Always set stable to default
+          uv run --locked --no-sync mike set-default --push stable
@@ -0,0 +1,116 @@
+name: Examples - APO
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 3 AM UTC+8
+    - cron: '0 19 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-apo, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'PR #{0} - Label {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('APO - {0}', github.event_name) }}
+
+jobs:
+  apo:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-apo' ||
+      github.event.action == 'ci-all'
+    name: APO (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    # This job is run on GitHub hosted runners rather than self-hosted runners because it needs no GPU.
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: |
+          uv sync --frozen --no-default-groups --extra apo \
+            --group dev --group experiment --group agents --group core-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: |
+          uv sync --frozen --no-default-groups --extra apo \
+            --group dev --group experiment --group agents --group core-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-apo-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: APO custom algorithm
+        run: |
+          set -ex
+          cd examples/apo
+          uv run apo_custom_algorithm_trainer.py | tee _ci_apo.log
+          # Check whether the log contains "Best prompt found:"
+          grep "Best prompt found:" _ci_apo.log
+        env:
+          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: APO custom algorithm debugger
+        run: |
+          set -ex
+          cd examples/apo
+          uv run apo_debug.py --mode runner
+          uv run apo_debug.py --mode hook
+          uv run apo_debug.py --mode trainer
+        env:
+          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: APO built-in algorithm
+        run: |
+          set -ex
+          cd examples/apo
+          uv run room_selector_apo.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        if: matrix.setup-script != 'legacy'
@@ -0,0 +1,205 @@
+name: Examples - Calc-X
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 3 AM UTC+8
+    - cron: '0 19 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-calc-x, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'PR #{0} - Label {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Calc-X - {0}', github.event_name) }}
+
+jobs:
+  calc-x:
+    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 }})
+    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-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: Prepare Calc-X dataset
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
+          unzip calc-x-data.zip -d data
+          rm calc-x-data.zip
+
+      - name: Calc-X MCP sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run tests/test_mcp_calculator.py
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: Calc-X sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run legacy_calc_agent_debug.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      # Calc-X training suddenly works after running the sanity check.
+      # And it has to be run before Spider training.
+      # The client side used to hang in many of my attempts.
+      # Don't ask why. Don't touch this.
+      - name: Calc-X training
+        run: |
+          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
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train
+
+      - name: Validate Calc-X training
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Calc-X training 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 --llm-proxy
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_llm_proxy
+
+      - name: Calc-X training with external store
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+
+          agl store --port 4747 &
+          sleep 5
+          AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=runner python train_calc_agent.py --external-store-address http://localhost:4747 --val-file data/test_mini.parquet --ci-fast &
+          sleep 5
+          AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=algorithm python train_calc_agent.py --external-store-address http://localhost:4747 --val-file data/test_mini.parquet --ci-fast
+
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+          pkill -f train_calc_agent.py && echo "SIGTERM sent to train_calc_agent.py" || echo "No train_calc_agent.py process found"
+          while pgrep -f train_calc_agent.py; do
+            echo "Waiting for train_calc_agent.py to finish..."
+            sleep 5
+          done
+          echo "train_calc_agent.py has finished."
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_external_store
+
+      - name: Calc-X training with role-based environment variables
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+
+          PYTHONUNBUFFERED=1 AGL_SERVER_HOST=127.0.0.1 AGL_SERVER_PORT=5858 AGL_CURRENT_ROLE=runner python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast &
+          sleep 5
+          PYTHONUNBUFFERED=1 AGL_SERVER_HOST=0.0.0.0 AGL_SERVER_PORT=5858 AGL_CURRENT_ROLE=algorithm python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast
+
+          pkill -f train_calc_agent.py && echo "SIGTERM sent to train_calc_agent.py" || echo "No train_calc_agent.py process found"
+          while pgrep -f train_calc_agent.py; do
+            echo "Waiting for train_calc_agent.py to finish..."
+            sleep 5
+          done
+          echo "train_calc_agent.py has finished."
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -0,0 +1,151 @@
+name: Examples - Backward Compatibility
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 6 AM UTC+8
+    - cron: '0 22 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-compat, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'PR #{0} - Label {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Backward Compatibility - {0}', github.event_name) }}
+
+jobs:
+  backward-compatibility:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-compat' ||
+      github.event.action == 'ci-all'
+    name: Backward Compatibility (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups --extra apo --extra verl \
+            --group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
+      - name: Override VERL (stable)
+        run: |
+          uv pip install verl==0.5.0
+        if: matrix.setup-script == 'stable'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-backward-compatibility-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+      - name: Prepare Calc-X dataset
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
+          unzip calc-x-data.zip -d data
+          rm calc-x-data.zip
+
+      - name: APO example (legacy client-server style)
+        run: |
+          set -ex
+          cd examples/apo
+          uv run legacy_apo_client.py &
+          sleep 3  # Wait for the client to be up
+          uv run legacy_apo_server.py
+          pkill -f legacy_apo_client.py && echo "SIGTERM sent to legacy_apo_client.py" || echo "No legacy_apo_client.py process found"
+          while pgrep -f legacy_apo_client.py; do
+            echo "Waiting for legacy_apo_client.py to finish..."
+            sleep 5
+          done
+          echo "legacy_apo_client.py has finished."
+          sleep 10
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: Calc-X MCP sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run tests/test_mcp_calculator.py
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: Calc-X sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run legacy_calc_agent_debug.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: Calc-X training (legacy client-server style)
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python legacy_calc_agent.py &
+          bash legacy_train.sh
+          pkill -f legacy_calc_agent.py && echo "SIGTERM sent to legacy_calc_agent.py" || echo "No legacy_calc_agent.py process found"
+          while pgrep -f legacy_calc_agent.py; do
+            echo "Waiting for legacy_calc_agent.py to finish..."
+            sleep 5
+          done
+          echo "legacy_calc_agent.py has finished."
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train
+
+      - name: Validate Calc-X training
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -0,0 +1,127 @@
+name: Examples - Spider
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: '0 20 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-spider, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'PR #{0} - Label {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Spider - {0}', github.event_name) }}
+
+jobs:
+  spider:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-spider' ||
+      github.event.action == 'ci-all'
+    name: Spider (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+        - 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-spider-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: Prepare Spider dataset
+        run: |
+          set -ex
+          cd examples/spider
+          uv run gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
+          unzip -q spider-data.zip -d data
+          rm spider-data.zip
+
+      - name: Spider sanity check
+        run: |
+          set -ex
+          cd examples/spider
+          uv run sql_agent.py
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        if: success() || failure()
+
+      - name: Spider training
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/spider
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_sql_agent.py fast
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: spider_train
+
+      - name: Validate Spider training
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -0,0 +1,129 @@
+name: Examples - Unsloth
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 5 AM UTC+8
+    - cron: '0 21 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-unsloth, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'PR #{0} - Label {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Unsloth - {0}', github.event_name) }}
+
+jobs:
+  unsloth:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-unsloth' ||
+      github.event.action == 'ci-all'
+    name: Unsloth (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        # Legacy versions are not supported for Unsloth examples.
+        include:
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group trl --group agents --group torch-gpu-stable
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-unsloth-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Prepare Unsloth model
+        run: |
+          set -ex
+          cd examples/unsloth
+          rm -rf models
+          uv run hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
+
+      - name: Unsloth SFT example
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/unsloth
+
+          agl store --port 4747 &
+          sleep 5
+          python sft_rollout_runners.py &
+          sleep 5
+          python sft_algorithm.py
+
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+          pkill -f sft_rollout_runners.py && echo "SIGTERM sent to sft_rollout_runners.py" || echo "No sft_rollout_runners.py process found"
+          while pgrep -f sft_rollout_runners.py; do
+            echo "Waiting for sft_rollout_runners.py to finish..."
+            sleep 5
+          done
+          echo "sft_rollout_runners.py has finished."
+          sleep 10
+
+          # Check models/version_2 must exist
+          if [ ! -d "models/version_2" ]; then
+            echo "models/version_2 does not exist"
+            exit 1
+          fi
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Unsloth SFT example all-in-one
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/unsloth
+          rm -rf models/version_1 models/version_2
+
+          python sft_allinone.py
+          if [ ! -d "models/version_2" ]; then
+            echo "models/version_2 does not exist"
+            exit 1
+          fi
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
@@ -1,156 +0,0 @@
-name: GPU Test
-permissions:
-  contents: read
-on:
-  schedule:
-    # Every day at 3 AM UTC+8
-    - cron: '0 19 * * *'
-
-  workflow_dispatch:
-
-jobs:
-  examples:
-    runs-on: [self-hosted, linux, gpu]
-    timeout-minutes: 60
-    strategy:
-      matrix:
-        setup: [stable, latest]
-      fail-fast: false
-    container:
-      image: ghcr.io/microsoft/agent-lightning/base:latest
-      options: --gpus all --ipc=host --interactive --tty
-    steps:
-      - name: Check GPU status
-        run: nvidia-smi
-      - uses: actions/checkout@v4
-      - name: Create a virtual environment
-        run: python3 -m venv .venv
-      - name: Install deps inside the container (${{ matrix.setup }})
-        run: |
-          . .venv/bin/activate
-          ./scripts/setup_${{ matrix.setup }}_gpu.sh
-      - name: Freeze dependencies
-        run: |
-          . .venv/bin/activate
-          which python
-          which pip
-          which uvx
-          pip list | tee requirements-freeze.txt
-      - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: dependencies-${{ matrix.setup }}
-          path: requirements-freeze.txt
-          compression-level: 0
-      - name: Prepare Spider dataset
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/spider
-          gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
-          unzip -q spider-data.zip -d data
-          rm spider-data.zip
-      - name: Prepare Calc-X dataset
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/calc_x
-          gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
-          unzip calc-x-data.zip -d data
-          rm calc-x-data.zip
-      - name: Spider sanity check
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/spider
-          python sql_agent.py --trainer.n-workers 1 --trainer.dev true --trainer.max-tasks 2
-        env:
-          VERL_API_BASE: http://localhost:9999/
-          OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }}
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-      - name: Calc-X MCP sanity check
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/calc_x
-          python tests/test_mcp_calculator.py
-        env:
-          OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }}
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-      - name: Calc-X sanity check
-        run: |
-          set -ex
-          . .venv/bin/activate
-          cd examples/calc_x
-          python calc_agent_dev.py
-        env:
-          OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }}
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-
-      # Calc-X training suddenly works after running the sanity check.
-      # And it has to be run before Spider training.
-      # The client side used to hang in many of my attempts.
-      # Don't ask why. Don't touch this.
-      - name: Calc-X training
-        run: |
-          set -ex
-          source .venv/bin/activate
-          cd examples/calc_x
-          ../../scripts/restart_ray.sh
-          sleep 5
-          PYTHONUNBUFFERED=1 python calc_agent.py &
-          bash train_ci.sh
-          pkill -f calc_agent.py && echo "SIGTERM sent to calc_agent.py" || echo "No calc_agent.py process found"
-          while pgrep -f calc_agent.py; do
-            echo "Waiting for calc_agent.py to finish..."
-            sleep 5
-          done
-          echo "calc_agent.py has finished."
-          sleep 10
-        shell: bash
-        env:
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-        id: calc_x_train
-
-      - name: Validate Calc-X training
-        run: |
-          set -ex
-          . .venv/bin/activate
-          python scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
-        env:
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-
-      - name: Spider training
-        run: |
-          set -ex
-          source .venv/bin/activate
-          cd examples/spider
-          ../../scripts/restart_ray.sh
-          sleep 5
-          PYTHONUNBUFFERED=1 python sql_agent.py --trainer.n-workers 10 &
-          bash train_ci.sh
-          pkill -f sql_agent.py && echo "SIGTERM sent to sql_agent.py" || echo "No sql_agent.py process found"
-          while pgrep -f sql_agent.py; do
-            echo "Waiting for sql_agent.py to finish..."
-            sleep 5
-          done
-          echo "sql_agent.py has finished."
-          sleep 10
-        shell: bash
-        env:
-          VERL_API_BASE: http://localhost:9991/
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-        id: spider_train
-        if: success() || failure()
-
-      - name: Validate Spider training
-        run: |
-          set -ex
-          . .venv/bin/activate
-          python scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }}
-        env:
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-
-      - name: Cleanup
-        run: ./scripts/cleanup.sh
-        if: success() || failure()
@@ -0,0 +1,309 @@
+name: Issue Comment
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  pull-requests: write
+  issues: write
+  contents: write
+  actions: read
+
+jobs:
+  dispatch:
+    # Only run for comments on pull requests AND when the comment starts with "/ci"
+    if: >
+      github.event.issue.pull_request != null &&
+      startsWith(github.event.comment.body, '/ci')
+    runs-on: ubuntu-latest
+    outputs:
+      dispatched: ${{ steps.dispatch.outputs.dispatched }}
+      event_types: ${{ steps.dispatch.outputs.event_types }}
+      correlation_id: ${{ steps.dispatch.outputs.correlation_id }}
+      trigger_comment_id: ${{ steps.dispatch.outputs.trigger_comment_id }}
+      ack_comment_id: ${{ steps.ack.outputs.comment_id }}
+    steps:
+      - name: Guardrail — allow only members/collaborators
+        id: guard
+        uses: actions/github-script@v8
+        with:
+          script: |
+            const allowed = ['MEMBER','OWNER','COLLABORATOR'];
+            const assoc = context.payload.comment.author_association;
+            if (!allowed.includes(assoc)) {
+              core.notice(`Ignoring /ci from ${context.payload.comment.user.login} (author_association=${assoc}).`);
+              core.setOutput('skip', 'true');
+            }
+
+      - name: Trigger repository dispatch
+        id: dispatch
+        if: steps.guard.outputs.skip != 'true'
+        uses: actions/github-script@v8
+        with:
+          script: |
+            const owner = context.repo.owner;
+            const repo = context.repo.repo;
+            const pull_number = context.payload.issue.number;
+            const comment = context.payload.comment;
+
+            // Fetch current PR state
+            const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number });
+
+            // Add reaction so folks know we saw it
+            try {
+              await github.rest.reactions.createForIssueComment({
+                owner,
+                repo,
+                comment_id: comment.id,
+                content: 'rocket'
+              });
+            } catch (e) {
+              core.info('Could not add reaction (likely due to permissions). Continuing.');
+            }
+
+            const labels = (pr.labels ?? []).map(label => label.name);
+            const directCiLabels = labels.filter(label => label.startsWith('ci-'));
+            const hasCiAll = directCiLabels.includes('ci-all');
+            const dedupe = new Set(
+              directCiLabels.filter(label => label !== 'ci-all')
+            );
+
+            if (!hasCiAll && dedupe.size === 0) {
+              core.notice('No ci-* labels found on the pull request; nothing to dispatch.');
+              core.setOutput('dispatched', 'false');
+              core.setOutput('event_types', '');
+              return;
+            }
+
+            const correlation_id = `id-${comment.id}-${Date.now().toString(36)}`;
+
+            const clientPayload = {
+              correlation_id,
+              pull_number,
+              pr_ref: `refs/pull/${pull_number}/merge`,
+              pr_head_ref: pr.head.ref,
+              pr_head_sha: pr.head.sha,
+              pr_base_ref: pr.base.ref,
+              pr_base_sha: pr.base.sha,
+              trigger_comment_id: comment.id,
+              trigger_comment_user: comment.user.login,
+            };
+
+            const eventTypes = hasCiAll
+              ? ['ci-all']
+              : Array.from(dedupe);
+            for (const eventType of eventTypes) {
+              await github.rest.repos.createDispatchEvent({
+                owner,
+                repo,
+                event_type: eventType,
+                client_payload: { ...clientPayload, ci_label: eventType }
+              });
+              core.notice(`Dispatched '${eventType}' event for PR #${pull_number}.`);
+            }
+
+            core.setOutput('dispatched', 'true');
+            core.setOutput('event_types', eventTypes.join(','));
+            core.setOutput('correlation_id', correlation_id);
+            core.setOutput('trigger_comment_id', String(comment.id));
+
+      - name: Acknowledge in thread (optional)
+        if: steps.guard.outputs.skip != 'true' && steps.dispatch.outputs.dispatched == 'true'
+        id: ack
+        uses: actions/github-script@v8
+        env:
+          EVENT_TYPES: ${{ steps.dispatch.outputs.event_types }}
+          CORRELATION_ID: ${{ steps.dispatch.outputs.correlation_id }}
+        with:
+          script: |
+            const eventTypes = (process.env.EVENT_TYPES || '')
+              .split(',')
+              .map(label => label.trim())
+              .filter(Boolean);
+            const formatted = eventTypes.map(label => `\`repository_dispatch:${label}\``).join(', ');
+            const { owner, repo } = context.repo;
+            const issue_number = context.payload.issue.number;
+            const body = [
+              `✅ CI trigger requested by @${context.payload.comment.user.login}.`,
+              `Fired ${formatted}.`,
+              '',
+              `_Collecting run links for correlation \`${process.env.CORRELATION_ID}\`…_`
+            ].join('\n');
+            const { data: comment } = await github.rest.issues.createComment({
+              owner, repo, issue_number,
+              body
+            });
+            core.setOutput('comment_id', String(comment.id));
+
+      - name: Notify missing ci label
+        if: steps.guard.outputs.skip != 'true' && steps.dispatch.outputs.dispatched != 'true'
+        uses: actions/github-script@v8
+        with:
+          script: |
+            const { owner, repo } = context.repo;
+            const issue_number = context.payload.issue.number;
+            await github.rest.issues.createComment({
+              owner,
+              repo,
+              issue_number,
+              body: `⚠️ CI trigger ignored because the pull request has no \`ci-*\` labels (e.g. \`ci-apo\`, \`ci-calc-x\`). Add the desired labels and try \`/ci\` again.`
+            });
+
+  watch:
+    needs: dispatch
+    if: needs.dispatch.outputs.dispatched == 'true'
+    runs-on: ubuntu-latest
+    timeout-minutes: 180
+    steps:
+      - name: Track dispatched runs and update comment
+        uses: actions/github-script@v8
+        env:
+          CORRELATION_ID: ${{ needs.dispatch.outputs.correlation_id }}
+          ACK_COMMENT_ID: ${{ needs.dispatch.outputs.ack_comment_id }}
+          TRIGGER_COMMENT_ID: ${{ needs.dispatch.outputs.trigger_comment_id }}
+        with:
+          script: |
+            const owner = context.repo.owner;
+            const repo = context.repo.repo;
+            const correlationId = process.env.CORRELATION_ID;
+            if (!correlationId) {
+              core.warning('No correlation id supplied; nothing to watch.');
+              return;
+            }
+
+            const ackCommentId = Number(process.env.ACK_COMMENT_ID || 0);
+            if (!ackCommentId) {
+              core.warning('No comment id available for updates; skipping watch.');
+              return;
+            }
+            const triggerCommentId = Number(process.env.TRIGGER_COMMENT_ID || 0);
+            if (!triggerCommentId) {
+              core.warning('No trigger comment id available; skipping watch.');
+              return;
+            }
+
+            const prefix = `🚀 CI Watcher for correlation ${correlationId} triggered by comment ${triggerCommentId}`;
+            core.notice(`Watching workflow runs for correlation '${correlationId}' using comment ${ackCommentId}.`);
+
+            function fmt(run) {
+              const status = run.status;
+              const conclusion = run.conclusion;
+              const badge = status === 'completed'
+                ? (conclusion === 'success' ? '🟢' : conclusion === 'failure' ? '🔴' : '🟡')
+                : (status === 'in_progress' ? '🟣' : '⚪️');
+              const title = run.display_title || run.name || `run ${run.id}`;
+              const statusText = status === 'completed' ? `${status}/${conclusion}` : status;
+              return `- ${badge} [${title}](${run.html_url}) — \`${statusText}\``;
+            }
+
+            const signatureOf = runs =>
+              runs
+                .map(run => `${run.id}:${run.status}/${run.conclusion || ''}`)
+                .sort()
+                .join('|');
+
+            const deadlineMs = Date.now() + 175 * 60 * 1000; // 175 minutes
+            let found = [];
+
+            async function searchOnce() {
+              const runs = await github.paginate(
+                github.rest.actions.listWorkflowRunsForRepo,
+                { owner, repo, event: 'repository_dispatch', per_page: 100 }
+              );
+              const cutoff = new Date(Date.now() - 60 * 60 * 1000); // last hour
+              return runs.filter(run => {
+                const createdAt = new Date(run.created_at);
+                const title = String(run.display_title || run.name || '');
+                return createdAt >= cutoff && title.includes(correlationId);
+              });
+            }
+
+            while (Date.now() < deadlineMs) {
+              found = await searchOnce();
+              if (found.length > 0) {
+                core.notice(`Discovered ${found.length} workflow run(s) for correlation '${correlationId}'.`);
+                break;
+              }
+              core.notice(`No runs found yet for correlation '${correlationId}'; retrying shortly.`);
+              await new Promise(res => setTimeout(res, 10000));
+            }
+
+            if (found.length === 0) {
+              core.notice(`Watcher timed out with no runs for correlation '${correlationId}'; notifying thread.`);
+              await github.rest.issues.updateComment({
+                owner,
+                repo,
+                comment_id: ackCommentId,
+                body: [
+                  prefix,
+                  `⚠️ I couldn't find any workflow runs for correlation \`${correlationId}\`.`,
+                  `They may be delayed or misconfigured.`
+                ].join('\n')
+              });
+              return;
+            }
+
+            const runIds = new Set(found.map(run => run.id));
+            let lastSignature = '';
+
+            async function refreshRuns() {
+              const ids = Array.from(runIds);
+              const refreshed = [];
+              for (const id of ids) {
+                const { data } = await github.rest.actions.getWorkflowRun({
+                  owner,
+                  repo,
+                  run_id: id
+                });
+                refreshed.push(data);
+              }
+              return refreshed;
+            }
+
+            async function updateCommentIfChanged(runs, allDone) {
+              const signature = signatureOf(runs);
+              if (signature === lastSignature) {
+                // Run statuses unchanged; skipping comment update.
+                return;
+              }
+              lastSignature = signature;
+              core.notice(`Updating comment ${ackCommentId} with ${runs.length} run status entries (allDone=${allDone}).`);
+              await github.rest.issues.updateComment({
+                owner,
+                repo,
+                comment_id: ackCommentId,
+                body: [
+                  prefix,
+                  `🏃‍♀️ Tracking ${runs.length} workflow run(s):`,
+                  '',
+                  ...runs.map(fmt),
+                  '',
+                  allDone ? '✅ All runs completed.' : '_Still running…_'
+                ].join('\n')
+              });
+            }
+
+            await updateCommentIfChanged(found, found.every(run => run.status === 'completed'));
+
+            while (Date.now() < deadlineMs) {
+              const latest = await searchOnce();
+              for (const run of latest) {
+                if (!runIds.has(run.id)) {
+                  runIds.add(run.id);
+                  core.notice(`Detected additional run ${run.id} (${run.name || run.display_title || 'unnamed'}) for correlation '${correlationId}'.`);
+                }
+              }
+              const current = await refreshRuns();
+              const allDone = current.every(run => run.status === 'completed');
+              await updateCommentIfChanged(current, allDone);
+              if (allDone) {
+                core.notice(`All runs for correlation '${correlationId}' completed; stopping watcher.`);
+                break;
+              }
+              await new Promise(res => setTimeout(res, 60000));
+            }
+
+            if (Date.now() >= deadlineMs) {
+              core.warning(`Watcher hit the deadline while monitoring correlation '${correlationId}'.`);
+            }
@@ -2,8 +2,8 @@ name: PyPI Nightly Build

 on:
  schedule:
-    # Run daily at 6:00 AM UTC
-    - cron: '0 6 * * *'
+    # Run daily at 6:00 AM UTC+8
+    - cron: '0 22 * * *'
  workflow_dispatch:  # Allow manual trigger

 jobs:
@@ -14,18 +14,25 @@ jobs:
      contents: read

    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v6
        with:
          python-version: '3.12'
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev

-      - name: Install build dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e .[dev]
+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Build dashboard
+        run: cd dashboard && npm run build

      - name: Get current version
        id: get_version
@@ -44,16 +51,9 @@ jobs:

      - name: Build package
        run: |
-          hatch build
+          uv build

      - name: Publish to Test PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://test.pypi.org/legacy/
-
-      - name: Test installation from Test PyPI
-        run: |
-          # Wait a bit for the package to be available
-          sleep 30
-          pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
-          python -c "import agentlightning; print('Package installed successfully')"
@@ -48,34 +48,34 @@ jobs:
      contents: read

    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v6
        with:
          python-version: '3.12'
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev

-      - name: Install build dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e .[dev]
+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Build dashboard
+        run: cd dashboard && npm run build

      - name: Build package
        run: |
-          hatch build
+          uv build

      - name: Verify package contents
        run: |
-          python -m tarfile -l dist/*.tar.gz
-          python -m zipfile -l dist/*.whl
+          uv run --locked --no-sync python -m tarfile -l dist/*.tar.gz
+          uv run --locked --no-sync python -m zipfile -l dist/*.whl

      - name: Publish to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
-
-      - name: Test installation from PyPI
-        run: |
-          # Wait a bit for the package to be available
-          sleep 30
-          pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
-          python -c "import agentlightning; print('Package installed successfully')"
@@ -0,0 +1,97 @@
+name: GPU Test
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 5 AM UTC+8
+    - cron: '0 21 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-gpu, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'PR #{0} - Label {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('GPU Test - {0}', github.event_name) }}
+
+jobs:
+  tests-full:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-gpu' ||
+      github.event.action == 'ci-all'
+    name: 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 --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 --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-${{ 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
+        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
@@ -5,9 +5,9 @@ permissions:

 on:
  push:
-    branches: [ main ]
+    branches: [ main, stable/**/* ]
  pull_request:
-    branches: [ main ]
+    branches: [ main, stable/**/* ]
  workflow_dispatch:

  schedule:
@@ -17,42 +17,93 @@ on:
 jobs:

  lint:
-    name: Lint with Black
+    strategy:
+      matrix:
+        setup: [fast, slow]
+    name: Lint - ${{ matrix.setup }}
    runs-on: ubuntu-latest
    timeout-minutes: 10
    steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v7
        with:
+          enable-cache: true
          python-version: '3.12'
-      - name: Install dependencies
+      - name: Sync dependencies (fast)
+        run: uv sync --frozen --group dev --no-default-groups
+        if: matrix.setup == 'fast'
+      - name: Sync dependencies (slow)
        run: |
-          python -m pip install --upgrade pip
-          pip install -e .[dev]
+          uv sync --frozen \
+            --extra apo \
+            --extra verl \
+            --group dev \
+            --group torch-cpu \
+            --group torch-stable \
+            --group trl \
+            --group tinker \
+            --group agents \
+            --no-default-groups
+        if: matrix.setup == 'slow'
+      # This pre-commit skips JavaScript on purpose.
+      - name: Run pre-commit
+        uses: pre-commit/action@v3.0.1
+      - name: Check Python headers
+        run: uv run --locked --no-sync scripts/check_headers.py
      - name: Run Black
-        run: |
-          black --check --diff --line-length=120 .
+        run: uv run --locked --no-sync black --check .
+      - name: Run isort
+        run: uv run --locked --no-sync isort --check-only .
+      - name: Run pyright (fast)
+        run: uv run --locked --no-sync pyright -p pyrightconfig.fast.json
+        if: matrix.setup == 'fast'
+      - name: Run pyright (slow)
+        run: uv run --locked --no-sync pyright -p pyrightconfig.json
+        if: matrix.setup == 'slow'
+
+  lint-js:
+    name: Lint - JavaScript
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - name: Install dependencies
+        run: cd dashboard && npm ci
+      - name: Run ESLint
+        run: cd dashboard && npm run eslint
+      - name: Run Prettier
+        run: cd dashboard && npm run prettier
+      - name: Run Stylelint
+        run: cd dashboard && npm run stylelint
+      - name: Run Typecheck
+        run: cd dashboard && npm run typecheck
+      - name: Verify build
+        run: cd dashboard && npm run build

  docs:
    name: Build documentation
    runs-on: ubuntu-latest
    timeout-minutes: 10
    steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
        with:
          fetch-depth: 0
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v6
        with:
          python-version: '3.12'
-      - name: Install documentation dependencies
-        run: |
-          ./scripts/setup_stable.sh
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev
      - name: Set source commit for docs
        run: |
          echo "SOURCE_COMMIT=${{ github.sha }}" >> $GITHUB_ENV
      - name: Build documentation
-        run: |
-          mkdocs build --strict
+        run: uv run --locked --no-sync mkdocs build --strict
      - name: Upload docs artifact
        uses: actions/upload-artifact@v4
        with:
@@ -65,35 +116,78 @@ jobs:
      matrix:
        include:
          - python-version: '3.10'
+            setup-script: 'legacy'
+          - python-version: '3.11'
            setup-script: 'stable'
          - python-version: '3.12'
+            setup-script: 'stable'
+          - python-version: '3.13'
            setup-script: 'latest'
-          - python-version: '3.12'
-            setup-script: 'stable'
      fail-fast: false

    name: Test with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v7
        with:
+          enable-cache: true
          python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
-        run: |
-          ./scripts/setup_${{ matrix.setup-script }}.sh
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
      - name: Freeze dependencies
        run: |
-          pip list | tee requirements-freeze-${{ matrix.python-version }}-${{ matrix.setup-script }}.txt
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-python-${{ matrix.python-version }}-${{ matrix.setup-script }}
-          path: requirements-freeze-${{ matrix.python-version }}-${{ matrix.setup-script }}.txt
+          name: dependencies-${{ matrix.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: Run tests
        run: |
-          pytest -v tests
+          uv run pytest -v --durations=0 tests
        env:
          PYTEST_ADDOPTS: "--color=yes"
+
+  test-js:
+    name: Test - JavaScript
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-node@v6
+        with:
+          node-version: '22'
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: '3.12'
+      - name: Sync Python dependencies
+        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-stable
+      - name: Install JavaScript dependencies
+        run: cd dashboard && npm ci
+      - name: Run vitest
+        run: cd dashboard && npm run vitest
@@ -183,12 +183,15 @@ cython_debug/
 .abstra/

 # Visual Studio Code
-#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
 #  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  and can be added to the global gitignore or merged into this file. However, if you prefer,
 #  you could uncomment the following to ignore the enitre vscode folder
 .vscode/

+# Emacs backup files
+*~
+
 # Ruff stuff:
 .ruff_cache/

@@ -201,3 +204,17 @@ cython_debug/
 #  refer to https://docs.cursor.com/context/ignore-files
 .cursorignore
 .cursorindexingignore
+
+# Claude
+.claude/*.local.json
+
+# Temporary and backup files
+*.tmp
+*.bak
+*.backup
+
+# Dashboard generated files
+agentlightning/dashboard/**/*.css
+agentlightning/dashboard/**/*.js
+agentlightning/dashboard/**/*.html
+agentlightning/dashboard/**/*.svg
@@ -1,8 +1,76 @@
 repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: check-yaml
+        exclude: ^mkdocs\.yml$
+      - id: check-toml
+      - id: check-added-large-files
+        args: ["--maxkb=1024"]
+        exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)
+      - id: check-shebang-scripts-are-executable
+      - id: detect-private-key
+  - repo: https://github.com/pycqa/isort
+    rev: 6.0.1
+    hooks:
+      - id: isort
+        args: ["."]
  - repo: https://github.com/psf/black
    rev: 25.1.0
    hooks:
-    - id: black
+      - id: black
+        pass_filenames: false
+        always_run: true
+        args: ["."]
+
+  - repo: local
+    hooks:
+    - id: prettier
+      name: prettier (dashboard)
+      language: system
      pass_filenames: false
      always_run: true
-      args: ["--line-length=120", "."]
+      entry: >
+        bash -c '
+          cd dashboard || exit 1
+          if [ -d node_modules ]; then
+            echo "✅ node_modules already exists"
+            npx prettier --cache --write "**/*.{ts,tsx,mjs,cjs}"
+          else
+            echo "⚠️ node_modules not found — npx is not reliable. Skipping."
+          fi
+        '
+
+    - id: eslint
+      name: eslint (dashboard)
+      language: system
+      pass_filenames: false
+      always_run: true
+      entry: >
+        bash -c '
+          cd dashboard || exit 1
+          if [ -d node_modules ]; then
+            echo "✅ node_modules already exists"
+            npx eslint --cache --fix .
+          else
+            echo "⚠️ node_modules not found — npx is not reliable. Skipping."
+          fi
+        '
+
+    - id: stylelint
+      name: stylelint (dashboard)
+      language: system
+      pass_filenames: false
+      always_run: true
+      entry: >
+        bash -c '
+          cd dashboard || exit 1
+          if [ -d node_modules ]; then
+            echo "✅ node_modules already exists"
+            npx stylelint --cache --fix "**/*.css"
+          else
+            echo "⚠️ node_modules not found — npx is not reliable. Skipping."
+          fi
+        '
@@ -0,0 +1 @@
+3.12
@@ -16,4 +16,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-THE SOFTWARE.
+THE SOFTWARE.
@@ -0,0 +1,77 @@
+# Responsible AI Transparency Documentation - Agent Lightning
+
+## OVERVIEW
+
+Agent Lightning is a flexible and extensible framework that enables seamless agent optimization for any existing agent framework. Agent optimization includes various data-driven techniques to customize the agent for better performance, including but not limited to model fine-tuning, prompt tuning, and model selection. And the agent frameworks refer to popular and easy-to-use agent developing frameworks such as OpenAI Agents SDK, Microsoft AutoGen, and LangChain.
+
+### WHAT CAN AGENT LIGHTNING DO
+Agent lightning was developed to bridge the gap between agent workflow development and agent optimization, empowering developers to go beyond static, pre-trained models and unlock the full potential of adaptive, learning-based agents. Agent Lightning is a training framework which can be used for any LLMs.
+
+### INTENDED USES
+Agent Lightning is best suited for agent researchers and developers. They can easily fine-tune models in existing agent frameworks with Agent Lightning. This can improve model performance on the targeted scenarios.
+
+### OUT-OF-SCOPE USES
+Agent Lightning is not well-suited for users who are not familiar with agent development and machine learning concepts.
+
+We do not recommend using Agent Lightning in commercial or real-world applications without further testing and development. It is being released for research purposes.
+
+Agent Lightning was not designed or evaluated for all possible downstream purposes. Developers should consider its inherent limitations as they select use cases, and evaluate and mitigate for accuracy, safety, and fairness concerns specific to each intended downstream use.
+
+Agent Lightning should not be used in highly regulated domains where inaccurate outputs could suggest actions that lead to injury or negatively impact an individual's legal, financial, or life opportunities.
+
+We do not recommend using Agent Lightning in the context of high-risk decision making (e.g. in law enforcement, legal, finance, or healthcare).
+
+## HOW TO GET STARTED
+To begin using Agent Lightning, here are some instructions.
+1.	Install dependencies, including Python, uv, PyTorch, FlashAttention, vLLM, verl.
+2.	Clone and install Agent Lightning.
+3.	Convert the dataset (provided by the user) into parquet file, which contains multiple columns. Each column contains a data id, an input and an expected output.
+4.	Run agent, which is developed by the user.
+5.	Run the training process via “bash train.sh”
+
+## EVALUATION
+Agent Lightning was evaluated on its ability to correctly complete 3 example tasks: (1) Math. The model needs to answer some math questions, and when answering one question, the model can use the calculator as its tool to help answer. (2) Text2SQL. The model is given a question related to the database, and it is required to generate a SQL which can query the database, find the information to answer the question. (3) Retrieval-Augmented Generation (RAG). The model is given a question which needs some information from Wikipedia to answer. The model is required to generate some queries to find the related information in Wikipedia, and answer the question according to retrieved documents.
+
+### EVALUATION METHODS AND RESULTS
+For detailed evaluation methods and results, please refer to the latest version of our  [technical report](https://arxiv.org/abs/2508.03680).
+
+
+## LIMITATIONS
+Agent Lightning was developed for research and experimental purposes. Further testing and validation are needed before considering its application in commercial or real-world scenarios.
+
+Agent Lightning was designed and tested using the English language. Performance in other languages may vary and should be assessed by someone who is both an expert in the expected outputs and a native speaker of that language.
+
+Outputs generated by AI may include factual errors, fabrication, or speculation. Users are responsible for assessing the accuracy of generated content. All decisions leveraging outputs of the system should be made with human oversight and not be based solely on system outputs.
+Agent Lightning inherits any biases, errors, or omissions produced by its base model. Developers are advised to choose an appropriate base LLM/MLLM carefully, depending on the intended use case.
+We use some demo cases to show the effectiveness of our training framework. See their links to understand the capabilities and limitations of this model.
+
+## BEST PRACTICES
+Better performance can be achieved by following the instructions in how to get started section.
+
+We strongly encourage users to use LLMs/MLLMs that support robust Responsible AI mitigations, such as Azure Open AI (AOAI) services. Such services continually update their safety and RAI mitigations with the latest industry standards for responsible use. For more on AOAI’s best practices when employing foundations models for scripts and applications:
+- [Blog post on responsible AI features in AOAI that were presented at Ignite 2023](https://techcommunity.microsoft.com/t5/ai-azure-ai-services-blog/announcing-new-ai-safety-amp-responsible-ai-features-in-azure/ba-p/3983686)
+- [Overview of Responsible AI practices for Azure OpenAI models](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/overview)
+- [Azure OpenAI Transparency Note](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/transparency-note)
+- [OpenAI’s Usage policies](https://openai.com/policies/usage-policies)
+- [Azure OpenAI’s Code of Conduct](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/code-of-conduct)
+
+Users are responsible for sourcing their datasets legally and ethically. This could include securing appropriate rights, ensuring consent for use of audio/images, and/or the anonymization of data prior to use in research.
+
+Users are reminded to be mindful of data privacy concerns and are encouraged to review the privacy policies associated with any models and data storage solutions interfacing with Agent Lightning.
+
+It is the user’s responsibility to ensure that the use of Agent Lightning complies with relevant data protection regulations and organizational guidelines.
+
+## LICENSE
+We use the MIT license.
+
+## CONTACT
+We welcome feedback and collaboration from our audience. If you have suggestions, questions, or observe unexpected/offensive behavior in our technology, please contact us at agent-lightning@microsoft.com.
+
+If the team receives reports of undesired behavior or identifies issues independently, we will update this repository with appropriate mitigations.
+
+
+
+---
+
+*Last updated: September 6, 2025*
+*Document version: 1.0*
@@ -1,11 +1,14 @@
-![Agent-lightning-banner](docs/assets/readme-banner.png)
+<p align="center">
+  <img src="docs/assets/readme-banner.svg" alt="Agent-lightning-banner" style="width:600px"/>
+</p>

 # Agent Lightning⚡

-[![CPU Test](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml)
-[![GPU Test](https://github.com/microsoft/agent-lightning/actions/workflows/examples.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples.yml)
+[![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)
+[![Documentation](https://img.shields.io/badge/GitHub%20Pages-Documentation-blue)](https://microsoft.github.io/agent-lightning/)
 [![PyPI version](https://badge.fury.io/py/agentlightning.svg)](https://badge.fury.io/py/agentlightning)
 [![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
+[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/microsoft/agent-lightning)
 [![Discord](https://img.shields.io/badge/Discord-Join-5865F2?logo=discord&logoColor=white)](https://discord.gg/RYk7CdvDR7)

 **The absolute trainer to light up AI agents.**
@@ -15,127 +18,61 @@ Join our [Discord community](https://discord.gg/RYk7CdvDR7) to connect with othe
 ## ⚡ Core Features

 - Turn your agent into an optimizable beast with **ZERO CODE CHANGE** (almost)! 💤
- Build with **ANY** agent framework (LangChain, OpenAI Agent SDK, AutoGen, CrewAI, ...); or even WITHOUT agent framework (Python OpenAI). You name it! 🤖
+- Build with **ANY** agent framework (LangChain, OpenAI Agent SDK, AutoGen, CrewAI, Microsoft Agent Framework...); or even WITHOUT agent framework (Python OpenAI). You name it! 🤖
 - **Selectively** optimize one or more agents in a multi-agent system. 🎯
- Embraces Reinforcement Learning, Automatic Prompt Optimization and more **algorithms**. 🤗
+- Embraces **Algorithms** like Reinforcement Learning, Automatic Prompt Optimization, Supervised Fine-tuning and more. 🤗

-![Agent-Lightning-code-diff](docs/assets/readme-diff.png)
+Read more on our [documentation website](https://microsoft.github.io/agent-lightning/).

-## ⚡ Resources
-
- 8/11/2025 [Training AI Agents to Write and Self-correct SQL with Reinforcement Learning](https://medium.com/@yugez/training-ai-agents-to-write-and-self-correct-sql-with-reinforcement-learning-571ed31281ad) Medium.
- 8/5/2025 [Agent Lightning: Train ANY AI Agents with Reinforcement Learning](https://arxiv.org/abs/2508.03680) arXiv paper.
- 7/26/2025 [We discovered an approach to train any AI agent with RL, with (almost) zero code changes.](https://www.reddit.com/r/LocalLLaMA/comments/1m9m670/we_discovered_an_approach_to_train_any_ai_agent/) Reddit.
- 6/6/2025 [Agent Lightning - Microsoft Research](https://www.microsoft.com/en-us/research/project/agent-lightning/) Project page.
+<p align="center">
+  <img src="docs/assets/readme-diff.svg" alt="Agent-Lightning Core Quickstart" style="width:100%"/>
+</p>

 ## ⚡ Installation

-First, let's get your environment set up. We'll be using `/path/to/agentlightning` to refer to the directory containing this README file.
-
-### 1. Set Up Your Environment
-
-We strongly recommend creating a new virtual environment to avoid conflicts with other packages. You can use either `conda` or `venv`. **Python 3.10 or later** is recommended.
-
-### 2. Install Core Training Dependencies (Optional)
-
-If you are running RL with Agent-Lightning, the next step is to install the essential packages: `PyTorch`, `FlashAttention`, `vLLM` and `VERL`. The following versions and installation order have been tested and are confirmed to work.
-
-```bash
-pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
-pip install flash-attn --no-build-isolation
-pip install vllm==0.9.2
-pip install verl==0.5.0
-```
-
-See `scripts/setup_stable_gpu.sh` for a full installation script.
-
-### 3. Install Agent Lightning
-
-Now, you're ready to install Agent Lightning itself.
-
 ```bash
 pip install agentlightning
 ```

-### 4. Install Agent Frameworks (Optional)
+Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.

-If you plan to use other agent frameworks, you can install them with the following commands. If you don't need these, feel free to skip this step.
-We recommend doing this as the final step to avoid dependency versions being overwritten by mistake.
+To start using Agent-lightning, check out our [documentation](https://microsoft.github.io/agent-lightning/) and [examples](./examples).

-```bash
-# AutoGen (Recommended to install first)
-pip install "autogen-agentchat" "autogen-ext[openai]"
+## ⚡ Articles

-# LiteLLM
-pip install "litellm[proxy]"
+- 11/4/2025 [Tuning ANY AI agent with Tinker ✕ Agent-lightning](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-1-1d8c9a397f0e) Medium. See also [Part 2](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-2-332c5437f0dc).
+- 10/22/2025 [No More Retokenization Drift: Returning Token IDs via the OpenAI Compatible API Matters in Agent RL](https://blog.vllm.ai/2025/10/22/agent-lightning.html) vLLM blog. See also [Zhihu writeup](https://zhuanlan.zhihu.com/p/1965067274642785725).
+- 8/11/2025 [Training AI Agents to Write and Self-correct SQL with Reinforcement Learning](https://medium.com/@yugez/training-ai-agents-to-write-and-self-correct-sql-with-reinforcement-learning-571ed31281ad) Medium.
+- 8/5/2025 [Agent Lightning: Train ANY AI Agents with Reinforcement Learning](https://arxiv.org/abs/2508.03680) arXiv paper.
+- 7/26/2025 [We discovered an approach to train any AI agent with RL, with (almost) zero code changes.](https://www.reddit.com/r/LocalLLaMA/comments/1m9m670/we_discovered_an_approach_to_train_any_ai_agent/) Reddit.
+- 6/6/2025 [Agent Lightning - Microsoft Research](https://www.microsoft.com/en-us/research/project/agent-lightning/) Project page.

-# MCP
-pip install mcp
+## ⚡ Community Projects

-# UV
-pip install uv
-
-# OpenAI Agents
-pip install openai-agents
-
-# LangChain
-pip install langgraph "langchain[openai]" langchain-community langchain-text-splitters
-
-# SQL-related dependencies
-pip install sqlparse nltk
-```
-
-Don't worry if dependency conflicts arise during this step. Follow the installation order above and the conflicts generally do not matter.
-
-## ⚡ Examples
-
-For more detailed examples, please see the `examples` folder:
-
-1. [calc_x](examples/calc_x): An agent built with AutoGen with calculator tool use, trained on Calc-X dataset with Reinforcement Learning.
-2. [spider](examples/spider): A write-check-rewrite looped agent with LangGraph with SQL execution; selectively optimize write and rewrite on Spider dataset with Reinforcement Learning.
-3. [apo](examples/apo): An example to customize an optimization algorithm: Automatic Prompt Optimization.
-
-## ⚡ Important Caveats
-
-1.  **AgentOps Integration**: Agent Lightning uses [AgentOps](https://github.com/AgentOps-AI/agentops) for agent tracking by default. If you're already using AgentOps in your own code, you'll need to disable our managed AgentOps client by modifying the `tracer` parameter of trainer.
-2.  **Debugging Traces**: If you encounter issues with tracing, you can visualize the trace tree using `tracer.last_trace().visualize("tree_graph")`. Please note that this API is experimental and may change in future releases.
-3.  **Launching the Server and Agents**: Currently, the training server and agent clients must be launched in separate processes. You can open two terminal windows or run one of them in the background. The launching order generally doesn't matter.
-4.  **Environment Variables**: The environment variables and working directory at the time of `ray init` are important. If you run into "file not found" errors, try restarting Ray from your current working directory.
-5.  **Handling Timeouts**: The training server may hang if samples fail or time out on the agent side. To prevent this, we recommend setting limits on the prompt and response lengths, as this is the most common cause of failures.
-6.  **VERL Failures**: Save checkpoints frequently, as VERL with vLLM may sometimes experience out-of-memory issues. If you encounter a VERL failure, you can resume training from the last checkpoint.
+- [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
+- [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.

 ## ⚡ Architecture

-Currently, Agent Lightning is built around a **training server** and one or multiple **agents**.
+Agent Lightning keeps the moving parts to a minimum so you can focus on your idea, not the plumbing. Your agent continues to run as usual; you can still use any agent framework you like; you drop in the lightweight `agl.emit_xxx()` helper, or let the tracer collect every prompt, tool call, and reward. Those events become structured spans that flow into the LightningStore, a central hub that keeps tasks, resources, and traces in sync.

-* The **server** manages the training data, prepares samples for the agents, and provides the LLM endpoint.
-* **Agents** retrieve samples from the server, process them (which may involve interacting with the LLM), and send the results back. These results, or "trajectories," are lists of prompts and responses from the LLM.
-* The **server** then collects these trajectories and computes the losses to optimize the language models.
+On the other side of the store sits the algorithm you choose, or write yourself. The algorithm reads spans, learns from them, and posts updated resources such as refined prompt templates or new policy weights. The Trainer ties it all together: it streams datasets to runners, ferries resources between the store and the algorithm, and updates the inference engine when improvements land. You can either stop there, or simply let the same loop keep turning.

-![Agent-Lightning-architecture](docs/assets/readme-architecture.png)
+No rewrites, no lock-in, just a clear path from first rollout to steady improvement.

-## ⚡ Development Instructions
+<p align="center">
+  <img src="docs/assets/readme-architecture.svg" alt="Agent-lightning Architecture" style="width:100%"/>
+</p>

-Install with development dependencies:
+## ⚡ CI Status

-```
-git clone https://github.com/microsoft/agent-lightning
-cd agent-lightning
-pip install -e .[dev]
-```
-
-Please run pre-commit hooks before checking in code:
-
-```
-pre-commit install
-pre-commit run --all-files --show-diff-on-failure --color=always
-```
-
-Serve documentation locally:
-
-```bash
-mkdocs serve
-```
+| Workflow | Status |
+|----------|--------|
+| CPU Tests | [![tests workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml) |
+| 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) |
+| 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) |

 ## ⚡ Citation

@@ -143,19 +80,19 @@ If you find Agent Lightning useful in your research or projects, please cite our

 ```bibtex
@misc{luo2025agentlightningtrainai,
-      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning}, 
+      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning},
      author={Xufang Luo and Yuge Zhang and Zhiyuan He and Zilong Wang and Siyun Zhao and Dongsheng Li and Luna K. Qiu and Yuqing Yang},
      year={2025},
      eprint={2508.03680},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
-      url={https://arxiv.org/abs/2508.03680}, 
+      url={https://arxiv.org/abs/2508.03680},
 }
 ```

 ## ⚡ Contributing

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

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

@@ -11,4 +11,4 @@ For security reporting information, locations, contact information, and policies
 please review the latest guidance for Microsoft repositories at
 [https://aka.ms/SECURITY.md](https://aka.ms/SECURITY.md).

-<!-- END MICROSOFT SECURITY.MD BLOCK -->
+<!-- END MICROSOFT SECURITY.MD BLOCK -->
@@ -1,10 +1,19 @@
-__version__ = "0.1.2"
+# Copyright (c) Microsoft. All rights reserved.

-from .client import AgentLightningClient, DevTaskLoader
-from .config import lightning_cli
-from .litagent import LitAgent
-from .logging import configure_logger
-from .reward import reward
-from .server import AgentLightningServer
-from .trainer import Trainer
+__version__ = "0.2.2"
+
+from .adapter import *
+from .algorithm import *
+from .client import AgentLightningClient, DevTaskLoader  # deprecated  # type: ignore
+from .config import *
+from .emitter import *
+from .execution import *
+from .litagent import *
+from .llm_proxy import *
+from .logging import *
+from .runner import *
+from .server import AgentLightningServer  # deprecated  # type: ignore
+from .store import *
+from .tracer import *
+from .trainer import *
 from .types import *
@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import Adapter, OtelTraceAdapter, TraceAdapter
+from .messages import TraceToMessages
+from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase
+
+__all__ = [
+    "TraceAdapter",
+    "OtelTraceAdapter",
+    "Adapter",
+    "TraceToTripletBase",
+    "TracerTraceToTriplet",
+    "LlmProxyTraceToTriplet",
+    "TraceToMessages",
+]
@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Generic, List, TypeVar
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import Span
+
+T_from = TypeVar("T_from")
+T_to = TypeVar("T_to")
+
+
+class Adapter(Generic[T_from, T_to]):
+    """Base class for synchronous adapters that convert data from one format to another.
+
+    The class defines a minimal protocol so that adapters can be treated like callables while
+    still allowing subclasses to supply the concrete transformation logic.
+
+    !!! note
+        Subclasses must override [`adapt()`][agentlightning.Adapter.adapt] to provide
+        the actual conversion.
+
+    Type Variables:
+
+        T_from: Source data type supplied to the adapter.
+
+        T_to: Target data type produced by the adapter.
+
+    Examples:
+        >>> class IntToStrAdapter(Adapter[int, str]):
+        ...     def adapt(self, source: int) -> str:
+        ...         return str(source)
+        ...
+        >>> adapter = IntToStrAdapter()
+        >>> adapter(42)
+        '42'
+    """
+
+    def __call__(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        This method delegates to [`adapt()`][agentlightning.Adapter.adapt] so that an
+        instance of [`Adapter`][agentlightning.Adapter] can be used like a standard
+        function.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        return self.adapt(source)
+
+    def adapt(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        Subclasses must override this method with the concrete transformation logic. The base
+        implementation raises `NotImplementedError` to make the requirement explicit.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        raise NotImplementedError("Adapter.adapt() is not implemented")
+
+
+class OtelTraceAdapter(Adapter[List[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
+    `opentelemetry.sdk.trace.ReadableSpan` instances and produces any target format, such as
+    reinforcement learning trajectories, structured logs, or analytics-ready payloads.
+
+    Examples:
+        >>> class TraceToDictAdapter(OtelTraceAdapter[dict]):
+        ...     def adapt(self, spans: List[ReadableSpan]) -> dict:
+        ...         return {"count": len(spans)}
+        ...
+        >>> adapter = TraceToDictAdapter()
+        >>> adapter([span1, span2])
+        {'count': 2}
+    """
+
+
+class TraceAdapter(Adapter[List[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
+    [`Span`][agentlightning.Span] instances emitted by Agent Lightning instrumentation.
+    Subclasses receive entire trace slices and return a format suited for the downstream consumer,
+    for example reinforcement learning training data or observability metrics.
+    """
@@ -0,0 +1,270 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+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 pydantic import TypeAdapter
+
+from agentlightning.types import Span
+
+from .base import TraceAdapter
+
+if TYPE_CHECKING:
+    from openai.types.chat import (
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+
+
+class OpenAIMessages(TypedDict):
+    """OpenAI-style chat messages with optional tool definitions.
+
+    Attributes:
+        messages: Ordered chat messages that describe the conversation.
+        tools: Tool specifications available to the assistant, if any.
+    """
+
+    messages: List[ChatCompletionMessageParam]
+    tools: Optional[List[ChatCompletionFunctionToolParam]]
+
+
+class _RawSpanInfo(TypedDict):
+    """Intermediate representation parsed from a span.
+
+    Attributes:
+        prompt: Prompt messages reconstructed from span attributes.
+        completion: Assistant completions following tool invocations.
+        request: Request payload recorded in the trace.
+        response: Response payload recorded in the trace.
+        tools: Tool call metadata extracted from child spans.
+    """
+
+    prompt: List[Dict[str, Any]]
+    completion: List[Dict[str, Any]]
+    request: Dict[str, Any]
+    response: Dict[str, Any]
+    tools: List[Dict[str, Any]]
+
+
+def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
+    """Convert flattened trace attributes into nested structures.
+
+    Attributes emitted by the tracing pipeline often arrive as dotted paths (for example
+    `gen_ai.prompt.0.role`). This helper groups those keys into nested dictionaries or lists so that
+    downstream processing can operate on structured data.
+
+    Args:
+        data: Flat dictionary whose keys are dotted paths.
+        prefix: Top-level key (for example `gen_ai.prompt`) that determines which attributes are
+            grouped.
+
+    Returns:
+        A nested dictionary (no numeric index detected) or list (numeric indices detected) containing
+        the grouped values.
+    """
+    result: Union[Dict[str, Any], List[Any]] = {}
+
+    # Collect keys that match the prefix
+    relevant = {k[len(prefix) + 1 :]: v for k, v in data.items() if k.startswith(prefix + ".")}
+
+    # Detect if we have numeric indices (-> list) or not (-> dict)
+    indexed = any(part.split(".")[0].isdigit() for part in relevant.keys())
+
+    if indexed:
+        # Group by index
+        grouped: Dict[int, Dict[str, Any]] = defaultdict(dict)
+        for k, v in relevant.items():
+            parts = k.split(".")
+            if not parts[0].isdigit():
+                continue
+            idx, rest = int(parts[0]), ".".join(parts[1:])
+            grouped[idx][rest] = v
+        # Recursively build
+        result = []
+        for i in sorted(grouped.keys()):
+            result.append(group_genai_dict({f"{prefix}.{rest}": val for rest, val in grouped[i].items()}, prefix))
+    else:
+        # No indices: build dict
+        nested: Dict[str, Any] = defaultdict(dict)
+        for k, v in relevant.items():
+            if "." in k:
+                head, _tail = k.split(".", 1)
+                nested[head][f"{prefix}.{k}"] = v
+            else:
+                result[k] = v
+        # Recurse into nested dicts
+        for head, subdict in nested.items():
+            result[head] = group_genai_dict(subdict, prefix + "." + head)
+
+    return result
+
+
+def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]:
+    """Convert raw trace payloads into OpenAI-style chat messages.
+
+    The function consumes an iterable produced by
+    [`TraceToMessages.adapt()`][agentlightning.TraceToMessages.adapt] and yields
+    structures that match the OpenAI fine-tuning JSONL schema, including tool definitions.
+
+    Args:
+        prompt_completion_list: Raw prompt/completion/tool payloads extracted from a trace.
+
+    Returns:
+        A generator that yields [`OpenAIMessages`][agentlightning.adapter.messages.OpenAIMessages]
+        entries compatible with the OpenAI Functions fine-tuning format.
+    """
+
+    # Import locally to avoid legacy OpenAI version type import errors
+    from openai.types.chat import (
+        ChatCompletionAssistantMessageParam,
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+
+    for pc_entry in prompt_completion_list:
+        messages: List[ChatCompletionMessageParam] = []
+
+        # Extract messages
+        for msg in pc_entry["prompt"]:
+            role = msg["role"]
+
+            if role == "assistant" and "tool_calls" in msg:
+                # Use the tool_calls directly
+                # This branch is usually not used in the wild.
+                tool_calls: List[ChatCompletionMessageFunctionToolCallParam] = [
+                    ChatCompletionMessageFunctionToolCallParam(
+                        id=call["id"],
+                        type="function",
+                        function={"name": call["name"], "arguments": call["arguments"]},
+                    )
+                    for call in msg["tool_calls"]
+                ]
+                messages.append(
+                    ChatCompletionAssistantMessageParam(role="assistant", content=None, tool_calls=tool_calls)
+                )
+            else:
+                # Normal user/system/tool content
+                message = cast(
+                    ChatCompletionMessageParam,
+                    TypeAdapter(ChatCompletionMessageParam).validate_python(
+                        dict(role=role, content=msg.get("content", ""), tool_call_id=msg.get("tool_call_id", None))
+                    ),
+                )
+                messages.append(message)
+
+        # Extract completions (assistant outputs after tool responses)
+        for comp in pc_entry["completion"]:
+            if comp.get("role") == "assistant":
+                content = comp.get("content")
+                if pc_entry["tools"]:
+                    tool_calls = [
+                        ChatCompletionMessageFunctionToolCallParam(
+                            id=tool["call"]["id"],
+                            type=tool["call"]["type"],
+                            function={"name": tool["name"], "arguments": tool["parameters"]},
+                        )
+                        for tool in pc_entry["tools"]
+                    ]
+                    messages.append(
+                        ChatCompletionAssistantMessageParam(role="assistant", content=content, tool_calls=tool_calls)
+                    )
+                else:
+                    messages.append(ChatCompletionAssistantMessageParam(role="assistant", content=content))
+
+        # Build tools definitions (if available)
+        if "functions" in pc_entry["request"]:
+            tools = [
+                ChatCompletionFunctionToolParam(
+                    type="function",
+                    function={
+                        "name": fn["name"],
+                        "description": fn.get("description", ""),
+                        "parameters": (
+                            json.loads(fn["parameters"]) if isinstance(fn["parameters"], str) else fn["parameters"]
+                        ),
+                    },
+                )
+                for fn in pc_entry["request"]["functions"]
+            ]
+            yield OpenAIMessages(messages=messages, tools=tools)
+        else:
+            yield OpenAIMessages(messages=messages, tools=None)
+
+
+class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
+    """Convert trace spans into OpenAI-compatible conversation messages.
+
+    The adapter reconstructs prompts, completions, tool calls, and function definitions from
+    `gen_ai.*` span attributes. The resulting objects match the JSONL structure expected by the
+    OpenAI fine-tuning pipeline.
+
+    !!! warning
+        The adapter assumes all spans share a common trace and that tool call spans are direct
+        children of the associated completion span.
+    """
+
+    def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable[Dict[str, Any]]:
+        """Yield tool call payloads for a completion span.
+
+        Args:
+            completion: The completion span whose descendants should be inspected.
+            all_spans: The complete span list belonging to the trace.
+
+        Yields:
+            Dictionaries describing tool calls with identifiers, names, and arguments.
+
+        Raises:
+            ValueError: If a candidate tool span cannot be converted into a dictionary.
+        """
+        # Get all the spans that are children of the completion span
+        children = [span for span in all_spans if span.parent_id == completion.span_id]
+        # Get the tool calls from the children
+        for maybe_tool_call in children:
+            tool_call = group_genai_dict(maybe_tool_call.attributes, "tool")
+            if not isinstance(tool_call, dict):
+                raise ValueError(f"Extracted tool call from trace is not a dict: {tool_call}")
+            if tool_call:
+                yield tool_call
+
+    def adapt(self, source: List[Span], /) -> List[OpenAIMessages]:
+        """Transform trace spans into OpenAI chat payloads.
+
+        Args:
+            source: Spans containing `gen_ai.*` attributes emitted by the tracing pipeline.
+
+        Returns:
+            A list of [`OpenAIMessages`][agentlightning.adapter.messages.OpenAIMessages] entries that
+            capture prompts, completions, tools, and metadata.
+        """
+        raw_prompt_completions: List[_RawSpanInfo] = []
+
+        for span in source:
+            attributes = {k: v for k, v in span.attributes.items()}
+
+            # Get all related information from the trace span
+            prompt = group_genai_dict(attributes, "gen_ai.prompt") or []
+            completion = group_genai_dict(attributes, "gen_ai.completion") or []
+            request = group_genai_dict(attributes, "gen_ai.request") or {}
+            response = group_genai_dict(attributes, "gen_ai.response") or {}
+            if not isinstance(prompt, list):
+                raise ValueError(f"Extracted prompt from trace is not a list: {prompt}")
+            if not isinstance(completion, list):
+                raise ValueError(f"Extracted completion from trace is not a list: {completion}")
+            if not isinstance(request, dict):
+                raise ValueError(f"Extracted request from trace is not a dict: {request}")
+            if not isinstance(response, dict):
+                raise ValueError(f"Extracted response from trace is not a dict: {response}")
+            if prompt or completion or request or response:
+                tools = list(self.get_tool_calls(span, source)) or []
+                raw_prompt_completions.append(
+                    _RawSpanInfo(
+                        prompt=prompt or [], completion=completion, request=request, response=response, tools=tools
+                    )
+                )
+
+        return list(convert_to_openai_messages(raw_prompt_completions))
@@ -0,0 +1,887 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
+
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel
+
+from agentlightning.types import Span, SpanNames, Triplet
+
+from .base import TraceAdapter
+
+logger = logging.getLogger(__name__)
+
+
+class Transition(BaseModel):
+    """A single transition within a reinforcement learning trajectory.
+
+    Attributes:
+        state: Token identifiers describing the model input state.
+        action: Token identifiers representing the model output.
+        response_id: Identifier of the LLM response used to deduplicate spans.
+        agent_name: Human-readable agent name captured from the trace.
+        reward: Scalar reward associated with the transition, if available.
+    """
+
+    state: List[int]
+    action: List[int]
+    response_id: Optional[str]
+    # action_logprobs: List[float]
+    agent_name: str
+    reward: Optional[float]
+
+
+class RewardMatchPolicy(str, Enum):
+    """Strategies for matching rewards to LLM call spans.
+
+    !!! note
+        Each reward span must expose a payload shaped like `{"type": "reward", "value": <float>|None}`
+        as described in `reward.py`.
+    """
+
+    FIRST_SIBLING = "first_sibling"
+    """Use the first sibling in the current trace subtree as the reward unless another LLM call match is found."""
+
+    FIRST_OCCURRENCE = "first_occurrence"
+    """Use the first reward encountered in chronological order after the current LLM call match."""
+
+
+class TraceTree:
+    """Tree representation of a trace span and its descendants.
+
+    Attributes:
+        id: Unique identifier for the span node.
+        span: [`Span`][agentlightning.Span] backing this node.
+        children: Child nodes connected to the current span.
+    """
+
+    def __init__(
+        self,
+        id: str,
+        span: Span,
+        children: Optional[List["TraceTree"]] = None,
+    ):
+        self.id = id
+        self.span = span
+        self.children = children or []
+
+    @property
+    def start_time(self):
+        return self.span.start_time
+
+    @property
+    def end_time(self):
+        return self.span.end_time
+
+    def find_id(self, id: str) -> "TraceTree | None":
+        if self.id == id:
+            return self
+        for child in self.children:
+            found = child.find_id(id)
+            if found:
+                return found
+        return None
+
+    def add_child(self, child: "TraceTree") -> None:
+        self.children.append(child)
+
+    def visualize(self, filename: str, interested_span_match: str | None = None) -> None:
+        """Render the trace tree with Graphviz for debugging purposes.
+
+        Args:
+            filename: Base filename for the generated `.png` diagram.
+            interested_span_match: Optional regular expression used to keep only matching spans
+                (and their ancestors) in the output.
+
+        !!! note
+            The method requires the optional `graphviz` dependency to be available in the runtime
+            environment.
+        """
+        import graphviz
+
+        dot = graphviz.Digraph(comment="Trace Tree")
+
+        should_visit_cache: Dict[str, bool] = {}
+
+        def should_visit(node: "TraceTree") -> bool:
+            if node.id in should_visit_cache:
+                return should_visit_cache[node.id]
+            if interested_span_match is not None:
+                if re.search(interested_span_match, node.span.name):
+                    should_visit_cache[node.id] = True
+                    return True
+                else:
+                    should_visit_cache[node.id] = False
+                    for child in node.children:
+                        if should_visit(child):
+                            should_visit_cache[node.id] = True
+
+                    return should_visit_cache[node.id]
+            else:
+                return True
+
+        def visit(node: "TraceTree") -> bool:
+            if not should_visit(node):
+                return False
+            agent_name = node.agent_name()
+            vis_name = node.id[:8] + " (" + node.span.name + ")"
+            if agent_name is not None:
+                vis_name += " [" + agent_name + "]"
+            dot.node(node.id, vis_name)  # type: ignore
+            for child in node.children:
+                if visit(child):
+                    dot.edge(node.id, child.id)  # type: ignore
+            return True
+
+        visit(self)
+        dot.render(filename, format="png", cleanup=True)  # type: ignore
+
+    def names_tuple(self) -> Tuple[str, List[Any]]:
+        """Return the span name alongside nested child names.
+
+        Returns:
+            A tuple of the current span name and a list of tuples for each child containing the
+            child name and its descendants.
+        """
+        name = self.span.name
+        agent_name = self.agent_name()
+        if agent_name is not None:
+            name += " [" + agent_name + "]"
+        children_names: List[Tuple[str, List[Any]]] = []
+        for child in self.children:
+            child_name, child_children = child.names_tuple()
+            children_names.append((child_name, child_children))
+        return name, children_names
+
+    def traverse(self) -> List["TraceTree"]:
+        """Traverse the tree depth first and return every node."""
+        spans: List["TraceTree"] = [self]
+        for child in self.children:
+            spans.extend(child.traverse())
+        return spans
+
+    def to_json(self) -> dict[str, Any]:
+        """Convert the tree node into a JSON-serialisable structure."""
+        if isinstance(self.span, ReadableSpan):
+            span_data = json.loads(self.span.to_json())
+        else:
+            span_data = self.span.model_dump()
+        return {
+            "id": self.id,
+            "span": span_data,
+            "children": [child.to_json() for child in self.children],
+        }
+
+    @classmethod
+    def from_spans(cls, spans: List[Span]) -> "TraceTree":
+        """Construct a tree from a flat list of spans.
+
+        Args:
+            spans: Spans that collectively form a single trace segment.
+
+        Returns:
+            A [`TraceTree`][agentlightning.adapter.triplet.TraceTree] rooted at either the
+            discovered root span or a synthetic root when multiple roots are present.
+
+        Raises:
+            ValueError: If the span list is empty or no root span can be inferred.
+        """
+
+        if not spans:
+            raise ValueError("No spans provided to create TraceTree.")
+
+        # Process trace items in topological order
+        id_to_span = {span.span_id: span for span in spans}
+
+        forward_graph: dict[str, list[str]] = {}
+        root_ids: list[str] = []
+        for span in spans:
+            span_id = span.span_id
+            if span.parent_id is None:
+                root_ids.append(span.span_id)
+            else:
+                if span.parent_id not in forward_graph:
+                    forward_graph[span.parent_id] = []
+                forward_graph[span.parent_id].append(span_id)
+
+        # Diff between span with data and forward_graph keys
+        # Sometimes the top-level session span is lost.
+        unfound_roots = set(forward_graph.keys()) - set(id_to_span.keys())
+        for unfound_root in unfound_roots:
+            root_ids.append(unfound_root)
+
+        def visit(node_id: str) -> "TraceTree":
+            children: list[TraceTree] = []
+            if node_id in forward_graph:
+                for child_id in forward_graph[node_id]:
+                    children.append(visit(child_id))
+
+            if node_id not in id_to_span:
+                assert len(children) > 0
+                virtual_span = Span.from_attributes(
+                    rollout_id=children[0].span.rollout_id,
+                    attempt_id=children[0].span.attempt_id,
+                    sequence_id=children[0].span.sequence_id,
+                    trace_id=children[0].span.trace_id,
+                    span_id=node_id,
+                    parent_id=None,
+                    attributes={},
+                    start_time=min(child.start_time for child in children if child.start_time is not None),
+                    end_time=max(child.end_time for child in children if child.end_time is not None),
+                )
+                return cls(node_id, virtual_span, children=children)
+            else:
+                return cls(
+                    node_id,
+                    id_to_span[node_id],
+                    children=children,
+                )
+
+        # Create a virtual root span if multiple root spans are found
+        if len(root_ids) > 1:
+            root_spans = [visit(root_id) for root_id in root_ids]
+            virtual_root = TraceTree(
+                id="virtual-root",
+                span=Span.from_attributes(
+                    rollout_id=root_spans[0].span.rollout_id,
+                    attempt_id=root_spans[0].span.attempt_id,
+                    sequence_id=root_spans[0].span.sequence_id,
+                    trace_id=root_spans[0].span.trace_id,
+                    span_id=None,  # Generate one
+                    parent_id=None,
+                    name="virtual-root",
+                    attributes={},
+                    start_time=root_spans[0].start_time,
+                    end_time=root_spans[-1].end_time,
+                ),
+                children=root_spans,
+            )
+            return virtual_root
+        elif len(root_ids) == 0:
+            # No root spans found
+            raise ValueError("No root spans found in the trace.")
+        else:
+            root_span = visit(root_ids[0])
+            return root_span
+
+    def agent_name(self) -> Optional[str]:
+        """Return the agent name associated with the span, if any.
+
+        Returns:
+            Agent name extracted from known attributes, otherwise `None`.
+        """
+        attributes = self.span.attributes
+        if attributes is None:  # type: ignore
+            return None
+
+        # Case 1: OpenAI Agent SDK
+        agent_name = cast(Optional[str], attributes.get("agent.name"))
+        if agent_name is not None:
+            return agent_name
+
+        # Case 2: Agentops decorator @agent
+        is_agent = attributes.get("agentops.span.kind") == "agent"
+        if is_agent:
+            agent_name = cast(Optional[str], attributes.get("operation.name"))
+            if agent_name is not None:
+                return agent_name
+
+        # Case 3: Autogen team
+        agent_name = cast(Optional[str], attributes.get("recipient_agent_type"))
+        if agent_name is not None:
+            return agent_name
+
+        # Case 4: LangGraph
+        agent_name = cast(Optional[str], attributes.get("langchain.chain.type"))
+        if agent_name is not None:
+            return agent_name
+
+        # Case 5: agent-framework
+        agent_name = cast(Optional[str], attributes.get("executor.id"))
+        if agent_name is not None:
+            return agent_name
+
+    def maybe_reward_dict(self) -> dict[str, Any]:
+        """Return a reward payload if the span encodes one.
+
+        Returns:
+            Dictionary containing reward metadata, or an empty dictionary when no reward is found.
+        """
+        for key in [
+            "agentops.task.output",  # newer versions of agentops
+            "agentops.entity.output",
+        ]:
+            output = self.span.attributes.get(key)  # type: ignore
+            if output:
+                if isinstance(output, dict):
+                    return output
+                elif isinstance(output, str):
+                    try:
+                        return json.loads(output)
+                    except json.JSONDecodeError:
+                        return {}
+
+        # Latest emit reward format
+        if self.span.name == SpanNames.REWARD.value and self.span.attributes:
+            return {"type": "reward", "value": self.span.attributes.get("reward", None)}
+        return {}
+
+    def is_reward_span(self) -> bool:
+        """Return whether the span explicitly encodes a reward.
+
+        Returns:
+            `True` when the span payload describes a reward, otherwise `False`.
+        """
+        maybe_reward = self.maybe_reward_dict()
+        return maybe_reward and maybe_reward.get("type") == "reward"  # type: ignore
+
+    def find_llm_calls(
+        self,
+        *,
+        llm_call_match: str,
+        agent_match: Optional[str],
+        within_matching_subtree: str | None = None,
+        within_reward: Optional[bool] = None,
+        within_llm_call: Optional[bool] = None,
+        existing_llm_call_response_ids: Optional[set[str]] = None,
+    ) -> List[Tuple["TraceTree", str]]:
+        """Find LLM call spans matching the supplied filters.
+
+        Args:
+            llm_call_match: Regular expression used to match span names that qualify as LLM calls.
+            agent_match: Optional regular expression that must match the enclosing agent span name.
+            within_matching_subtree: Marker propagated through recursive calls to record matching agents.
+            within_reward: When `True`, suppresses LLM matches under reward spans.
+            within_llm_call: When `True`, prevents duplicate matches for nested LLM calls.
+            existing_llm_call_response_ids: Known response identifiers used to deduplicate spans.
+
+        Returns:
+            A list of tuples pairing the matching node with the agent subtree label that triggered the
+            match.
+        """
+        llm_calls: List[Tuple[TraceTree, str]] = []
+
+        is_llm_call = True
+        if within_matching_subtree is None or within_reward is True:
+            # We must be in an interesting agent subtree, and not in a reward span.
+            is_llm_call = False
+        if re.search(llm_call_match, self.span.name) is None:
+            # The span name does not match the LLM call match.
+            is_llm_call = False
+        if is_llm_call:
+            # Check the response id
+            response_id: Optional[str] = self.span.attributes.get("gen_ai.response.id")  # type: ignore
+            if response_id is None and within_llm_call is True:
+                is_llm_call = False
+            if (
+                response_id is not None
+                and existing_llm_call_response_ids is not None
+                and response_id in existing_llm_call_response_ids
+            ):
+                is_llm_call = False
+
+            if is_llm_call:
+                llm_calls.append((self, within_matching_subtree))  # type: ignore
+                existing_llm_call_response_ids = existing_llm_call_response_ids or set()
+                if response_id is not None:
+                    existing_llm_call_response_ids.add(response_id)
+                if within_llm_call is not None:
+                    within_llm_call = True
+
+        agent_name = self.agent_name()
+        if agent_name is not None:
+            if agent_match is None or re.search(agent_match, agent_name):
+                within_matching_subtree = agent_name
+            else:
+                within_matching_subtree = None
+
+        if within_reward is not None and self.is_reward_span():
+            within_reward = True
+
+        for child in self.children:
+            llm_calls.extend(
+                child.find_llm_calls(
+                    llm_call_match=llm_call_match,
+                    agent_match=agent_match,
+                    within_matching_subtree=within_matching_subtree,
+                    within_reward=within_reward,
+                    within_llm_call=within_llm_call,
+                    existing_llm_call_response_ids=existing_llm_call_response_ids,
+                )
+            )
+
+        return llm_calls
+
+    def repair_hierarchy(self) -> None:
+        """Repair missing parent-child relationships introduced by mixed tracing systems.
+
+        Some agent frameworks emit spans via multiple subsystems, which can cause LLM completion
+        spans to float directly under the root span instead of being nested under the correct agent.
+        The method re-parents those spans to the closest ancestor that fully envelopes the child in
+        time.
+
+        If we don't, when we want to select the LLM completion span with agent as filter.
+        We will never get the correct span underneath.
+        """
+        # If the current node has only one child, recursively repair its hierarchy directly.
+        # This special-case handling is needed because when a trace is manually ended
+        # (via agentops.end_trace), the AgentOps provider automatically wraps all spans
+        # under an extra synthetic root node (e.g., "run_one.session").
+        if len(self.children) == 1:
+            self.children[0].repair_hierarchy()
+            return
+
+        nodes_to_repair = list(self.children)
+
+        for repair_node in nodes_to_repair:
+            if len(self.children) == 1:
+                # If there is only one child, we don't need to repair the hierarchy.
+                break
+            # Find the closest parent span (but not the root itself)
+            closest_parent = None
+            closest_duration = float("inf")
+            for node in self.traverse():
+                if node.id == repair_node.id:
+                    continue
+                if node is self:
+                    continue
+                if node.start_time <= repair_node.start_time and node.end_time >= repair_node.end_time:  # type: ignore
+                    duration_delta = node.end_time - repair_node.end_time + repair_node.start_time - node.start_time  # type: ignore
+                    if duration_delta > 0 and duration_delta < closest_duration:
+                        closest_duration = duration_delta  # type: ignore
+                        closest_parent = node
+
+            # Repair the hierarchy
+            if closest_parent is not None:
+                self.children.remove(repair_node)
+                closest_parent.children.append(repair_node)
+
+    def match_rewards(self, reward_match: str, llm_calls: List["TraceTree"]) -> dict[str, Optional[float]]:
+        """Assign rewards to previously matched LLM calls.
+
+        Args:
+            reward_match: Strategy identifier from
+                [`RewardMatchPolicy`][agentlightning.adapter.triplet.RewardMatchPolicy].
+            llm_calls: Trace nodes representing LLM call spans.
+
+        Returns:
+            Mapping from span identifier to reward value or `None` when no reward is available.
+        """
+        llm_call_ids = set([llm_call.id for llm_call in llm_calls])
+        rewards: dict[str, Optional[float]] = {}
+
+        if reward_match == RewardMatchPolicy.FIRST_OCCURRENCE:
+            time_sorted: List[TraceTree] = cast(List[TraceTree], sorted(self.traverse(), key=lambda x: x.start_time))  # type: ignore
+            assign_to: List[Tuple[str, int]] = []  # type: ignore
+            for item in time_sorted:
+                if item.id in llm_call_ids:
+                    assign_to.append((item.id, item.end_time))  # type: ignore
+
+                # get reward
+                agentops_output = item.maybe_reward_dict()
+                if agentops_output and agentops_output.get("type") == "reward":
+                    for assign_to_id, assign_to_end_time in reversed(assign_to):
+                        # This reward happens before the end of the LLM call.
+                        if assign_to_end_time > item.start_time:  # type: ignore
+                            continue
+                        # Ok, we found someone to assign to
+                        if assign_to_id in rewards:
+                            # If the reward is already set, skip
+                            continue
+                        rewards[assign_to_id] = agentops_output.get("value", None)
+                        break
+
+        elif reward_match == RewardMatchPolicy.FIRST_SIBLING:
+            for item in self.traverse():
+                assign_to: List[Tuple[str, int]] = []
+                for child in item.children:
+                    if child.id in llm_call_ids:
+                        assign_to.append(child.id)  # type: ignore
+
+                    agentops_output = item.maybe_reward_dict()
+                    if agentops_output and agentops_output.get("type") == "reward":
+                        for assign_to_id, assign_to_end_time in reversed(assign_to):
+                            if assign_to_end_time > item.start_time:  # type: ignore
+                                # This reward happens before the end of the LLM call.
+                                continue
+                            if assign_to_id in rewards:
+                                continue
+                            rewards[assign_to_id] = agentops_output.get("value", None)
+                            break
+
+        return rewards
+
+    def span_to_triplet(self, span: Span, agent_name: str) -> Triplet:
+        """Convert a span to a triplet.
+
+        Subclass can override this method to add more fields to the triplet,
+        such as chat messages and tool calls.
+        """
+        prompt_token_ids = span.attributes.get("prompt_token_ids", [])  # type: ignore
+        response_token_ids = span.attributes.get("response_token_ids", [])  # type: ignore
+        response_id = span.attributes.get("gen_ai.response.id", None)  # type: ignore
+
+        logprobs_content = span.attributes.get("logprobs.content", None)  # type: ignore
+        if isinstance(logprobs_content, str):
+            logprobs_content = json.loads(logprobs_content)
+            response: Dict[str, Any] = {"token_ids": response_token_ids, "logprobs": logprobs_content}
+        else:
+            response = {"token_ids": response_token_ids}
+
+        return Triplet(
+            prompt={"token_ids": prompt_token_ids},
+            response=response,
+            reward=None,
+            metadata=dict(response_id=response_id, agent_name=agent_name),
+        )
+
+    def to_trajectory(
+        self,
+        llm_call_match: str = r"openai\.chat\.completion",
+        agent_match: Optional[str] = None,
+        exclude_llm_call_in_reward: bool = True,
+        dedup_llm_call: bool = True,
+        reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
+        final_reward: Optional[float] = None,
+        _skip_empty_token_spans: bool = False,
+    ) -> List[Triplet]:
+        """Convert the trace tree into a trajectory of [`Triplet`][agentlightning.Triplet] items.
+
+        Args:
+            llm_call_match: Regular expression for LLM call span names.
+            agent_match: Optional regular expression for agent span names.
+            exclude_llm_call_in_reward: When `True`, prevents searching for rewards under the LLM
+                call subtree.
+            dedup_llm_call: When `True`, deduplicates spans using the LLM response identifier.
+            reward_match: Reward matching policy used to associate reward spans with LLM calls.
+            final_reward: Optional reward appended to the final transition when provided.
+
+        Returns:
+            A list of [`Triplet`][agentlightning.Triplet] objects ordered by call sequence.
+        """
+        # Find all LLM calls
+        llm_calls = self.find_llm_calls(
+            llm_call_match=llm_call_match,
+            agent_match=agent_match,
+            within_matching_subtree="*" if agent_match is None else None,
+            within_reward=False if exclude_llm_call_in_reward else None,
+            within_llm_call=False if dedup_llm_call else None,
+            existing_llm_call_response_ids=set(),
+        )
+
+        id_transitions: List[Tuple[str, Triplet]] = []
+        # We need to filter out the LLM calls with unrecorded token IDs
+        filtered_llm_calls: List[Tuple[TraceTree, str]] = []
+        for llm_call, agent_name in llm_calls:
+            triplet = self.span_to_triplet(llm_call.span, agent_name)
+            # This is a hot-fix for Tinker+CrewAI, which has some anonymous requests outside the trained agent.
+            # TODO: We might need to reconsider this.
+            if _skip_empty_token_spans and (
+                not triplet.prompt.get("token_ids") or not triplet.response.get("token_ids")
+            ):
+                logger.warning(f"Skipping LLM call with unrecorded token IDs: {triplet}")
+                continue
+            filtered_llm_calls.append((llm_call, agent_name))
+            id_transitions.append((llm_call.id, triplet))
+
+        rewards = self.match_rewards(reward_match, [call for call, _ in filtered_llm_calls])
+        transitions = [
+            transition.model_copy(update={"reward": rewards.get(id, None)}) for id, transition in id_transitions
+        ]
+        if final_reward is not None and len(transitions) > 0:
+            # Add the final reward to the last transition
+            transitions[-1] = transitions[-1].model_copy(update={"reward": final_reward})
+        return transitions
+
+    def __repr__(self):
+        return (
+            f"TraceTree(id={self.id}, span={self.span}, start_time={self.start_time}, "
+            + f"end_time={self.end_time}, children={self.children})"
+        )
+
+
+class TraceToTripletBase(TraceAdapter[List[Triplet]]):
+    """Base class for adapters that emit [`Triplet`][agentlightning.Triplet] trajectories."""
+
+
+class TracerTraceToTriplet(TraceToTripletBase):
+    """Convert tracer-emitted spans into triplet trajectories.
+
+    Attributes:
+        repair_hierarchy: When `True`, repair the span tree using
+            [`TraceTree.repair_hierarchy()`][agentlightning.adapter.triplet.TraceTree.repair_hierarchy]
+            before matching calls and rewards.
+        llm_call_match: Regular expression pattern that selects LLM call span names.
+        agent_match: Optional regular expression pattern for agent span names. When omitted, spans
+            from any agent are considered.
+        exclude_llm_call_in_reward: When `True`, ignore matches under reward spans while searching
+            for rewards.
+        reward_match: Strategy used to associate rewards with LLM calls.
+    """
+
+    def __init__(
+        self,
+        repair_hierarchy: bool = True,
+        llm_call_match: str = r"openai\.chat\.completion",
+        agent_match: Optional[str] = None,
+        exclude_llm_call_in_reward: bool = True,
+        reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
+        _skip_empty_token_spans: bool = False,
+    ):
+        self.repair_hierarchy = repair_hierarchy
+        self.llm_call_match = llm_call_match
+        self.agent_match = agent_match
+        self.exclude_llm_call_in_reward = exclude_llm_call_in_reward
+        self.reward_match = reward_match
+        self._skip_empty_token_spans = _skip_empty_token_spans
+
+    def visualize(
+        self,
+        source: Union[List[Span], List[ReadableSpan]],
+        /,
+        filename: str = "trace_tree",
+        interested_span_match: str | None = None,
+    ) -> TraceTree:
+        """Visualize the trace tree built from the supplied spans.
+
+        Args:
+            source: Collection of Agent Lightning [`Span`][agentlightning.Span] objects
+                or raw `opentelemetry.sdk.trace.ReadableSpan` instances.
+            filename: Base filename for the generated image; `.png` is appended automatically.
+            interested_span_match: Optional regular expression used to highlight a subset of spans.
+
+        Returns:
+            The [`TraceTree`][agentlightning.adapter.triplet.TraceTree] built from the provided
+            spans.
+        """
+        source_normalized = [
+            Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
+            for span in source
+        ]
+        trace_tree = TraceTree.from_spans(source_normalized)
+        if self.repair_hierarchy:
+            trace_tree.repair_hierarchy()
+        trace_tree.visualize(filename, interested_span_match=interested_span_match)
+        return trace_tree
+
+    def adapt(self, source: Union[List[Span], List[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
+        """Convert tracer spans into [`Triplet`][agentlightning.Triplet] trajectories.
+
+        Args:
+            source: Agent Lightning spans or raw OpenTelemetry spans that form a trace.
+
+        Returns:
+            Ordered list of trajectory transitions with prompt, response, and reward information.
+        """
+        source_normalized = [
+            Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
+            for span in source
+        ]
+        trace_tree = TraceTree.from_spans(source_normalized)
+        if self.repair_hierarchy:
+            trace_tree.repair_hierarchy()
+        trajectory = trace_tree.to_trajectory(
+            llm_call_match=self.llm_call_match,
+            agent_match=self.agent_match,
+            exclude_llm_call_in_reward=self.exclude_llm_call_in_reward,
+            reward_match=self.reward_match,
+            _skip_empty_token_spans=self._skip_empty_token_spans,
+        )
+        return trajectory
+
+
+class LlmProxyTraceToTriplet(TraceToTripletBase):
+    """Convert telemetry emitted by the LLM Proxy into triplet trajectories.
+
+    !!! warning
+        This adapter is experimental and might be merged with
+        [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] in the future.
+
+    !!! danger
+        Do not rely on timestamps when using this adapter. Proxy spans can originate on different
+        machines with unsynchronised clocks, so `sequence_id` is treated as the sole source of
+        ordering.
+
+    Strategy:
+
+    1. Sort spans by `(sequence_id, start_time)` for deterministic processing.
+    2. Extract token identifiers from `litellm_request` or `raw_gen_ai_request` spans.
+    3. Extract rewards from spans exposing AgentOps-style payloads or explicit reward spans.
+    4. Match each reward to the most recent unmatched LLM call whose sequence is smaller.
+    """
+
+    def _literal_eval_maybe(self, v: Any) -> Any:
+        import ast
+
+        if isinstance(v, str):
+            try:
+                return ast.literal_eval(v)
+            except Exception:
+                return v
+        return v
+
+    def _extract_tokens_from_raw(self, attrs: Dict[str, Any]) -> Tuple[List[int], List[int]]:
+        """Extract token ids from raw_gen_ai_request attributes.
+
+        - llm.hosted_vllm.prompt_token_ids: string -> List[int]
+        - llm.hosted_vllm.response_token_ids: string -> List[List[int]] -> take first
+        - llm.hosted_vllm.choices: string -> [{'token_ids': [...]}] -> take first
+        """
+        prompt_ids: List[int] = []
+        resp_ids: List[int] = []
+
+        # prompt
+        p = attrs.get("llm.hosted_vllm.prompt_token_ids")
+        p = self._literal_eval_maybe(p)
+        if isinstance(p, list) and all(isinstance(x, int) for x in p):  # type: ignore
+            prompt_ids = cast(List[int], p)
+
+        # response preferred path
+        r = attrs.get("llm.hosted_vllm.response_token_ids")
+        r = self._literal_eval_maybe(r)
+        if isinstance(r, list) and len(r) > 0 and isinstance(r[0], list):  # type: ignore
+            first = cast(List[Any], r[0])
+            if all(isinstance(x, int) for x in first):
+                resp_ids = cast(List[int], first)
+
+        # fallback via choices
+        if not resp_ids:
+            choices = attrs.get("llm.hosted_vllm.choices")
+            choices = self._literal_eval_maybe(choices)
+            if isinstance(choices, list) and choices:
+                cand = cast(Any, choices[0])
+                if isinstance(cand, dict):
+                    tids = cast(Dict[str, Any], cand).get("token_ids")
+                    if isinstance(tids, list) and all(isinstance(x, int) for x in tids):  # type: ignore
+                        resp_ids = cast(List[int], tids)
+
+        return prompt_ids, resp_ids
+
+    def _extract_tokens_from_openai(self, attrs: Dict[str, Any]) -> Tuple[List[int], List[int]]:
+        prompt_ids = cast(Any, attrs.get("prompt_token_ids") or [])
+        resp_ids = cast(Any, attrs.get("response_token_ids") or [])
+        prompt_ids = self._literal_eval_maybe(prompt_ids)
+        resp_ids = self._literal_eval_maybe(resp_ids)
+        if not (isinstance(prompt_ids, list) and all(isinstance(x, int) for x in prompt_ids)):  # type: ignore
+            prompt_ids = []
+        if not (isinstance(resp_ids, list) and all(isinstance(x, int) for x in resp_ids)):  # type: ignore
+            resp_ids = []
+        return cast(List[int], prompt_ids), cast(List[int], resp_ids)
+
+    def _maybe_reward_value(self, span: Span) -> Optional[float]:
+        """Parse reward from typical AgentOps payloads or explicit reward spans."""
+        attrs = span.attributes or {}
+
+        # AgentOps new/old keys
+        for k in ("agentops.task.output", "agentops.entity.output"):
+            v = attrs.get(k)
+            v = self._literal_eval_maybe(v)
+            if isinstance(v, dict) and cast(Dict[str, Any], v).get("type") == "reward":
+                rv = cast(Dict[str, Any], v).get("value", None)
+                if rv is None or isinstance(rv, (int, float)):
+                    return None if rv is None else float(rv)
+
+        # Explicit reward span
+        if span.name == SpanNames.REWARD.value:
+            rv = attrs.get("reward", None)
+            if rv is None or isinstance(rv, (int, float)):
+                return None if rv is None else float(rv)
+
+        return None
+
+    def _request_id_from_attrs(self, attrs: Dict[str, Any]) -> Optional[str]:
+        # Prefer OpenAI-like id if present, else proxy raw id.
+        rid = attrs.get("gen_ai.response.id") or attrs.get("llm.hosted_vllm.id")
+        return str(rid) if isinstance(rid, str) and rid else None
+
+    def adapt(self, source: List[Span], /) -> List[Triplet]:  # type: ignore
+        """Convert LLM Proxy spans into [`Triplet`][agentlightning.Triplet] trajectories.
+
+        Args:
+            source: Spans emitted by the LLM Proxy containing prompt, response, and reward data.
+
+        Returns:
+            Ordered trajectory transitions matched purely by `sequence_id`.
+        """
+        # 1) Sort deterministically by (sequence_id, start_time).
+        spans = sorted(
+            source,
+            key=lambda s: (s.sequence_id, s.start_time),
+        )
+
+        # 2) Collect LLM calls with token IDs.
+        llm_items: List[Dict[str, Any]] = []
+        seen_request_ids: set[str] = set()
+        for s in spans:
+            attrs = s.attributes or {}
+            prompt_ids: List[int] = []
+            resp_ids: List[int] = []
+
+            if s.name == "raw_gen_ai_request":
+                prompt_ids, resp_ids = self._extract_tokens_from_raw(attrs)
+            elif s.name == "litellm_request":
+                # Some proxies never include token ids here. Ignore unless present.
+                prompt_ids, resp_ids = self._extract_tokens_from_openai(attrs)
+
+            if prompt_ids and resp_ids:
+                rid = self._request_id_from_attrs(attrs)
+                if rid:
+                    # Duplicated request ID. This request is already handled.
+                    if rid in seen_request_ids:
+                        continue
+                    seen_request_ids.add(rid)
+                llm_items.append(
+                    dict(
+                        span=s,
+                        seq=s.sequence_id,
+                        response_ids=resp_ids,
+                        prompt_ids=prompt_ids,
+                        request_id=rid,
+                    )
+                )
+
+        # Order LLM items by sequence only.
+        llm_items.sort(key=lambda x: x["seq"])
+
+        # Collect rewards by sequence only.
+        rewards: List[Tuple[int, Optional[float]]] = []
+        for s in spans:
+            val = self._maybe_reward_value(s)
+            if val is not None:
+                rewards.append((s.sequence_id, val))
+
+        # First-occurrence matching by sequence_id only:
+        # For reward at sequence R, assign to the most recent unmatched LLM with seq < R.
+        assigned: Dict[str, Optional[float]] = {}
+        for r_seq, r_val in sorted(rewards, key=lambda x: x[0]):
+            for item in reversed(llm_items):
+                sid = item["span"].span_id
+                if sid in assigned:
+                    continue
+                if item["seq"] < r_seq:
+                    assigned[sid] = r_val
+                    break
+
+        # Build triplets in LLM sequence order.
+        triplets: List[Triplet] = []
+        for item in llm_items:
+            s = item["span"]
+            triplets.append(
+                Triplet(
+                    prompt={"token_ids": item["prompt_ids"]},
+                    response={"token_ids": item["response_ids"]},
+                    reward=assigned.get(s.span_id, None),
+                    metadata=dict(
+                        # This is called response_id to align with the other adapters.
+                        response_id=item["request_id"],
+                    ),
+                )
+            )
+
+        return triplets
@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .base import Algorithm
+from .decorator import algo
+from .fast import Baseline, FastAlgorithm
+
+if TYPE_CHECKING:
+    from .apo import APO as APOType
+    from .verl import VERL as VERLType
+
+__all__ = ["Algorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]
+
+# Shortcuts for usages like algo.APO(...)
+
+
+def APO(*args: Any, **kwargs: Any) -> APOType[Any]:
+    from .apo import APO as APOImplementation
+
+    return APOImplementation(*args, **kwargs)
+
+
+def VERL(*args: Any, **kwargs: Any) -> VERLType:
+    from .verl import VERL as VERLImplementation
+
+    return VERLImplementation(*args, **kwargs)
@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .apo import APO
+
+__all__ = ["APO"]
@@ -0,0 +1,863 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+APO with textual gradients that read rollout spans and outputs to modify the prompt.
+
+- algo: beam search with span-aware textual gradients -> apply_edit via LLM
+- rollout: same pattern as your example, but task is a dict (T_task)
+"""
+
+import asyncio
+import logging
+import random
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Counter, Dict, Generic, Iterator, List, Optional, Sequence, Set, Tuple, TypedDict, TypeVar, cast
+
+import poml
+from openai import AsyncOpenAI
+
+from agentlightning.adapter.messages import TraceToMessages
+from agentlightning.algorithm.base import Algorithm
+from agentlightning.algorithm.utils import batch_iter_over_dataset
+from agentlightning.reward import find_final_reward
+from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
+
+logger = logging.getLogger(__name__)
+
+T_task = TypeVar("T_task")
+
+
+class RolloutResultForAPO(TypedDict):
+    """This must be all JSON serializable to be processable by POML."""
+
+    status: RolloutStatus
+    final_reward: Optional[float]
+    spans: List[Dict[str, Any]]
+    messages: List[Any]
+
+
+@dataclass
+class VersionedPromptTemplate:
+    version: str
+    prompt_template: PromptTemplate
+    score: Optional[float] = None
+
+
+GRADIENT_PROMPT_FILES = [
+    Path(__file__).parent / "prompts" / "text_gradient_variant01.poml",
+    Path(__file__).parent / "prompts" / "text_gradient_variant02.poml",
+    Path(__file__).parent / "prompts" / "text_gradient_variant03.poml",
+]
+
+APPLY_EDIT_PROMPT_FILES = [
+    Path(__file__).parent / "prompts" / "apply_edit_variant01.poml",
+    Path(__file__).parent / "prompts" / "apply_edit_variant02.poml",
+]
+
+
+class APO(Algorithm, Generic[T_task]):
+    """Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search.
+
+    APO is an iterative prompt optimization algorithm that uses LLM-generated textual gradients
+    to improve prompts through a beam search process. It evaluates prompts on rollouts,
+    computes critiques based on the results, and applies edits to generate improved prompts.
+
+    The algorithm operates in rounds, where each round:
+
+    1. Samples parent prompts from the current beam
+    2. Generates new prompts by computing textual gradients and applying edits
+    3. Evaluates all candidates on a validation set
+    4. Selects the top-k prompts for the next round
+
+    Based on the ideas from:
+
+    - [ProTeGi](https://aclanthology.org/2023.emnlp-main.494.pdf)
+    - [TextGrad](https://github.com/zou-group/textgrad)
+    """
+
+    def __init__(
+        self,
+        async_openai_client: AsyncOpenAI,
+        *,
+        gradient_model: str = "gpt-5-mini",
+        apply_edit_model: str = "gpt-4.1-mini",
+        diversity_temperature: float = 1.0,
+        gradient_batch_size: int = 4,
+        val_batch_size: int = 16,
+        beam_width: int = 4,
+        branch_factor: int = 4,
+        beam_rounds: int = 3,
+        rollout_batch_timeout: float = 3600.0,
+        run_initial_validation: bool = True,
+        # Internal flags for debugging
+        _poml_trace: bool = False,
+    ):
+        """
+        Initialize the APO algorithm with configuration parameters.
+
+        Args:
+            async_openai_client: AsyncOpenAI client for making LLM API calls.
+            gradient_model: Model name for computing textual gradients (critiques).
+            apply_edit_model: Model name for applying edits based on critiques.
+            diversity_temperature: Temperature parameter for LLM calls to control diversity.
+            gradient_batch_size: Number of rollout results to sample for gradient computation.
+            val_batch_size: Number of validation examples to use for evaluation.
+            beam_width: Number of top-scoring prompts to keep in the beam at each round.
+            branch_factor: Number of new prompt candidates to generate from each parent prompt
+                by applying textual gradient edits. This controls the expansion of the search tree.
+            beam_rounds: Number of beam search rounds to perform.
+            rollout_batch_timeout: Maximum time in seconds to wait for rollout batch completion.
+            run_initial_validation: If True, runs validation on the seed prompt before starting
+                optimization to establish a baseline score. Defaults to True.
+        """
+        self.async_openai_client = async_openai_client
+        self.gradient_model = gradient_model
+        self.apply_edit_model = apply_edit_model
+        self.diversity_temperature = diversity_temperature
+        self.gradient_batch_size = gradient_batch_size
+        self.val_batch_size = val_batch_size
+        self.beam_width = beam_width
+        self.branch_factor = branch_factor
+        self.beam_rounds = beam_rounds
+        self.rollout_batch_timeout = rollout_batch_timeout
+        self.run_initial_validation = run_initial_validation
+
+        self._history_best_prompt: Optional[PromptTemplate] = None
+        self._history_best_score: float = float("-inf")
+        self._history_best_version: Optional[str] = None
+
+        self._version_counter: int = 0
+
+        self._poml_trace = _poml_trace
+
+    def _create_versioned_prompt(
+        self,
+        prompt_template: PromptTemplate,
+        *,
+        score: Optional[float] = None,
+    ) -> VersionedPromptTemplate:
+        """
+        Wrap a prompt template with a new monotonically increasing version identifier.
+        """
+        version = f"v{self._version_counter}"
+        self._version_counter += 1
+        return VersionedPromptTemplate(version=version, prompt_template=prompt_template, score=score)
+
+    def _format_log_prefix(
+        self,
+        *,
+        round_num: Optional[int] = None,
+        beam_idx: Optional[int] = None,
+        branch_idx: Optional[int] = None,
+        prompt_version: Optional[str] = None,
+    ) -> str:
+        """
+        Construct the standardized log prefix.
+        """
+        parts: List[str] = []
+        if round_num is not None:
+            parts.append(f"Round {round_num:02d}")
+        if beam_idx is not None:
+            parts.append(f"Beam {beam_idx:02d}")
+        if branch_idx is not None:
+            parts.append(f"Branch {branch_idx:02d}")
+        if prompt_version is not None:
+            parts.append(f"Prompt {prompt_version}")
+        if not parts:
+            return ""
+        return f"[{' | '.join(parts)}]"
+
+    def _log(self, level: int, message: str, *, prefix: Optional[str] = None) -> None:
+        """
+        Log a message with an optional standardized prefix.
+        """
+        effective_prefix = prefix
+        if effective_prefix:
+            logger.log(level, f"{effective_prefix} {message}")
+        else:
+            logger.log(level, message)
+
+    def get_seed_prompt_template(self) -> Tuple[str, PromptTemplate]:
+        """
+        Extract the initial prompt template from the algorithm's resources.
+
+        Returns:
+            A tuple of (resource_name, prompt_template) representing the seed prompt.
+
+        Raises:
+            ValueError: If initial_resources is not set or no PromptTemplate is found.
+        """
+        initial_resources = self.get_initial_resources()
+        if initial_resources is None:
+            raise ValueError(
+                "initial_resources are not set for APO algorithm. "
+                "Use algorithm.set_initial_resources() to set initial resources or set it in Trainer()"
+            )
+        for name, resource in initial_resources.items():
+            if isinstance(resource, PromptTemplate):
+                return name, resource
+        raise ValueError("No prompt template resource found in initial_resources")
+
+    def get_adapter(self) -> TraceToMessages:
+        """
+        Get the adapter for converting spans to messages.
+
+        Returns:
+            The TraceToMessages instance for this algorithm.
+
+        Raises:
+            ValueError: If the adapter is not a TraceToMessages.
+        """
+        adapter = super().get_adapter()
+        if not isinstance(adapter, TraceToMessages):
+            raise ValueError("Adapter must be a TraceToMessages for APO algorithm")
+        return adapter
+
+    def get_best_prompt(self) -> PromptTemplate:
+        """
+        Retrieve the best prompt discovered during optimization.
+
+        Returns:
+            The prompt template with the highest validation score found so far.
+
+        Raises:
+            ValueError: If no best prompt has been found yet (run() not called).
+        """
+        if self._history_best_prompt is None:
+            raise ValueError("No best prompt found")
+        return self._history_best_prompt
+
+    async def compute_textual_gradient(
+        self,
+        current_prompt: VersionedPromptTemplate,
+        rollout_results: List[RolloutResultForAPO],
+        *,
+        prefix: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Compute a textual gradient (critique) for the current prompt based on rollout results.
+
+        This method samples rollout results, sends them to an LLM along with the current prompt,
+        and generates a critique describing how the prompt could be improved.
+
+        Args:
+            current_prompt: The prompt template to critique.
+            rollout_results: List of rollout results containing spans, messages, and rewards.
+
+        Returns:
+            A textual critique generated by the LLM, or None if generation fails.
+        """
+        tg_template = random.choice(GRADIENT_PROMPT_FILES)
+
+        if len(rollout_results) < self.gradient_batch_size:
+            self._log(
+                logging.WARNING,
+                f"Only {len(rollout_results)} rollouts available, but {self.gradient_batch_size} are needed. Using all rollouts.",
+                prefix=prefix,
+            )
+            sampled_rollout_results = rollout_results
+        else:
+            sampled_rollout_results = random.sample(rollout_results, self.gradient_batch_size)
+
+        self._log(
+            logging.INFO,
+            f"Gradient will be computed with {self.gradient_model} for {len(sampled_rollout_results)} rollouts with template: {tg_template.name}",
+            prefix=prefix,
+        )
+
+        tg_msg = poml.poml(  # type: ignore
+            tg_template,
+            context={
+                "experiments": sampled_rollout_results,
+                "prompt_template": current_prompt.prompt_template.template,
+            },
+            format="openai_chat",
+        )
+        self._log(
+            logging.DEBUG,
+            f"Gradient computed with {self.gradient_model} prompt: {tg_msg}",
+            prefix=prefix,
+        )
+        critique_response = await self.async_openai_client.chat.completions.create(
+            model=self.gradient_model,
+            messages=tg_msg["messages"],  # type: ignore
+            temperature=self.diversity_temperature,
+        )
+        critique_text = critique_response.choices[0].message.content
+        self._log(
+            logging.INFO,
+            f"Gradient computed with {self.gradient_model} has result: {critique_text}",
+            prefix=prefix,
+        )
+
+        return critique_text
+
+    async def textual_gradient_and_apply_edit(
+        self,
+        current_prompt: VersionedPromptTemplate,
+        rollout: List[RolloutResultForAPO],
+        *,
+        prefix: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Generate an improved prompt by computing a textual gradient and applying an edit.
+
+        This is the main optimization step that:
+
+        1. Computes a critique (textual gradient) based on rollout performance
+        2. Uses another LLM to apply the critique and generate an improved prompt
+
+        Args:
+            current_prompt: The current prompt template to improve.
+            rollout: List of rollout results to base the critique on.
+
+        Returns:
+            The improved prompt text, or the original prompt if gradient computation fails.
+        """
+        # 1) Critique
+        critique_text = await self.compute_textual_gradient(
+            current_prompt,
+            rollout,
+            prefix=prefix,
+        )
+        if not critique_text:
+            self._log(
+                logging.ERROR,
+                "Failed to compute critique for prompt.",
+                prefix=prefix,
+            )
+            return current_prompt.prompt_template.template
+
+        # 2) Apply edit
+        ae_template = random.choice(APPLY_EDIT_PROMPT_FILES)
+        self._log(
+            logging.INFO,
+            f"Edit will be generated by {self.apply_edit_model} with template: {ae_template.name}",
+            prefix=prefix,
+        )
+        ae_msg = poml.poml(  # type: ignore
+            ae_template,
+            context={
+                "prompt_template": current_prompt.prompt_template.template,
+                "critique": critique_text,
+            },
+            format="openai_chat",
+        )
+
+        ae_response = await self.async_openai_client.chat.completions.create(
+            model=self.apply_edit_model,
+            messages=ae_msg["messages"],  # type: ignore
+            temperature=self.diversity_temperature,
+        )
+        new_prompt = ae_response.choices[0].message.content
+        if new_prompt:
+            self._log(
+                logging.INFO,
+                f"Edit generated by {self.apply_edit_model}: {new_prompt[:50]}...",
+                prefix=prefix,
+            )
+        return new_prompt
+
+    async def get_rollout_results(
+        self,
+        rollout: List[Rollout],
+        *,
+        prefix: Optional[str] = None,
+    ) -> List[RolloutResultForAPO]:
+        """
+        Convert completed rollouts to APO-compatible result format.
+
+        Fetches spans for each rollout, adapts them to messages, and packages them
+        with rewards and status information for gradient computation.
+
+        Args:
+            rollout: List of completed rollout metadata.
+
+        Returns:
+            List of rollout results formatted for APO processing.
+        """
+        rollout_results: List[RolloutResultForAPO] = []
+        store = self.get_store()
+        adapter = self.get_adapter()
+        for r in rollout:
+            spans = await store.query_spans(r.rollout_id)
+            messages = adapter.adapt(spans)
+            rollout_result = RolloutResultForAPO(
+                status=r.status,
+                final_reward=find_final_reward(spans),
+                spans=[span.model_dump() for span in spans],
+                messages=messages,
+            )
+            self._log(
+                logging.DEBUG,
+                f"Rollout result for {r.rollout_id}: status {rollout_result['status']} with final reward {rollout_result['final_reward']}. "
+                f"{len(rollout_result['spans'])} spans and {len(rollout_result['messages'])} messages.",
+                prefix=prefix,
+            )
+            rollout_results.append(rollout_result)
+        return rollout_results
+
+    async def evaluate_prompt_on_batch(
+        self,
+        prompt: VersionedPromptTemplate,
+        resource_name: str,
+        dataset: Sequence[T_task],
+        mode: RolloutMode,
+        *,
+        prefix: Optional[str] = None,
+    ) -> Tuple[List[RolloutResultForAPO], float]:
+        """
+        Evaluate a prompt on a batch of tasks by running rollouts and computing average reward.
+
+        This method:
+
+        1. Adds the prompt as a named resource to the store
+        2. Enqueues rollouts for each task in the dataset
+        3. Waits for rollouts to complete (with timeout)
+        4. Computes and returns the average reward
+
+        Args:
+            prompt: The prompt template string to evaluate.
+            resource_name: The name to register the prompt under in the store.
+            dataset: Sequence of tasks to evaluate the prompt on.
+            mode: Rollout mode ("train" or "val") for logging/tracking.
+
+        Returns:
+            A tuple of (rollout_results, average_reward) where rollout_results contains
+            detailed information for each rollout and average_reward is the mean final reward.
+        """
+        store = self.get_store()
+        preview = prompt.prompt_template.template[:50]
+        self._log(
+            logging.INFO,
+            f'Evaluating prompt "{preview}..." on {len(dataset)} tasks in {mode} mode',
+            prefix=prefix,
+        )
+
+        # Install prompt as named resource
+        resources: NamedResources = {resource_name: prompt.prompt_template}
+        resource_update = await store.update_resources(prompt.version, resources)
+
+        rollout_ids: List[str] = []
+        for t in dataset:
+            r = await store.enqueue_rollout(input=t, mode=mode, resources_id=resource_update.resources_id)
+            rollout_ids.append(r.rollout_id)
+
+        deadline = time.time() + self.rollout_batch_timeout
+        finished: List[Rollout] = []
+        while time.time() < deadline:
+            finished = await store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=0.0)
+            if len(finished) >= len(rollout_ids):
+                self._log(
+                    logging.INFO,
+                    f"All {len(rollout_ids)} rollouts finished within timeout.",
+                    prefix=prefix,
+                )
+                break
+            else:
+                self._log(
+                    logging.DEBUG,
+                    f"Only {len(finished)} rollouts finished within timeout. Waiting for remaining {len(rollout_ids) - len(finished)} rollouts.",
+                    prefix=prefix,
+                )
+                # Sleep to avoid busy-waiting
+                await asyncio.sleep(2.0)
+
+        rollout_results = await self.get_rollout_results(
+            finished,
+            prefix=prefix,
+        )
+        final_rewards = [rr["final_reward"] for rr in rollout_results]
+
+        avg = float(sum([r or 0.0 for r in final_rewards]) / max(1, len(final_rewards)))
+        status_counter = Counter([rr["status"] for rr in rollout_results])
+
+        self._log(
+            logging.INFO,
+            f"Evaluated {len(rollout_results)} rollouts. Statuses: {status_counter}. Rewards: {final_rewards}, average is {avg}",
+            prefix=prefix,
+        )
+        return rollout_results, avg
+
+    def _initialize_beam(
+        self,
+        train_dataset: Optional[Dataset[T_task]],
+        val_dataset: Optional[Dataset[T_task]],
+    ) -> Tuple[str, PromptTemplate, Iterator[Sequence[T_task]], Iterator[Sequence[T_task]]]:
+        """
+        Initialize the beam search with seed prompt and dataset iterators.
+
+        Args:
+            train_dataset: Dataset for computing gradients.
+            val_dataset: Dataset for evaluating prompts.
+
+        Returns:
+            Tuple of (resource_name, seed_prompt, grad_iterator, val_iterator).
+
+        Raises:
+            ValueError: If either dataset is None.
+        """
+        resource_name, seed_prompt = self.get_seed_prompt_template()
+
+        if train_dataset is None:
+            raise ValueError("train_dataset is required for APO algorithm")
+        if val_dataset is None:
+            raise ValueError("val_dataset is required for APO algorithm")
+
+        grad_dataset_iterator = batch_iter_over_dataset(train_dataset, self.gradient_batch_size)
+        val_dataset_iterator = batch_iter_over_dataset(val_dataset, self.val_batch_size)
+
+        # Initialize history tracking
+        self._history_best_prompt = seed_prompt
+        self._history_best_score = float("-inf")
+
+        return resource_name, seed_prompt, grad_dataset_iterator, val_dataset_iterator
+
+    def _sample_parent_prompts(
+        self,
+        beam: List[VersionedPromptTemplate],
+        round_num: int,
+    ) -> List[Tuple[int, VersionedPromptTemplate]]:
+        """
+        Sample parent prompts from the current beam for generating new candidates.
+
+        If the beam has fewer prompts than beam_width, replicates existing prompts.
+        Otherwise, randomly samples beam_width prompts.
+
+        Args:
+            beam: Current list of prompt templates in the beam.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of parent prompts to generate children from.
+        """
+        display_round = round_num + 1
+        if len(beam) < self.beam_width:
+            prefix = self._format_log_prefix(round_num=display_round)
+            self._log(
+                logging.WARNING,
+                f"Beam width is currently {self.beam_width}, but only {len(beam)} prompts in beam. Replicating all prompts.",
+                prefix=prefix,
+            )
+            return [(i % len(beam), beam[i % len(beam)]) for i in range(self.beam_width)]
+
+        selected_indices = random.sample(range(len(beam)), self.beam_width)
+        return [(idx, beam[idx]) for idx in selected_indices]
+
+    async def _generate_candidate_prompts(
+        self,
+        parent_prompts: List[Tuple[int, VersionedPromptTemplate]],
+        resource_name: str,
+        grad_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[VersionedPromptTemplate]:
+        """
+        Generate new candidate prompts from parents using textual gradients.
+
+        For each parent prompt, generates branch_factor new candidates by:
+
+        1. Evaluating the parent on a training batch
+        2. Computing textual gradient
+        3. Applying edit to generate improved prompt
+
+        Args:
+            parent_prompts: List of parent prompts to generate children from.
+            resource_name: Name to register prompts under in the store.
+            grad_dataset_iterator: Iterator over training data batches.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of newly generated prompt templates.
+        """
+        display_round = round_num + 1
+        round_prefix = self._format_log_prefix(round_num=display_round)
+        self._log(
+            logging.INFO,
+            f"Applying {self.branch_factor} edits to each of the {len(parent_prompts)} parents based on "
+            "gradients computed on training dataset",
+            prefix=round_prefix,
+        )
+
+        parent_prompts_str = [
+            f"{p.version}:{p.score:.3f}" if p.score is not None else p.version for _, p in parent_prompts
+        ]
+        self._log(
+            logging.INFO,
+            f"Parent prompts: {', '.join(parent_prompts_str)}",
+            prefix=round_prefix,
+        )
+
+        candidates: List[VersionedPromptTemplate] = []
+        used_beam_indices: Set[int] = set()
+        for real_beam_idx, (beam_idx, prompt) in enumerate(parent_prompts):
+            if beam_idx in used_beam_indices:
+                beam_prefix = self._format_log_prefix(
+                    round_num=display_round,
+                    beam_idx=beam_idx + 1,
+                    prompt_version=prompt.version,
+                )
+                self._log(
+                    logging.WARNING,
+                    "Duplicated beam index found. Might be caused by beam_width too high. "
+                    + f"The real index of this beam is {real_beam_idx + 1}.",
+                    prefix=beam_prefix,
+                )
+            else:
+                used_beam_indices.add(beam_idx)
+            for branch_idx in range(self.branch_factor):
+                parent_prefix = self._format_log_prefix(
+                    round_num=display_round,
+                    beam_idx=beam_idx + 1,
+                    branch_idx=branch_idx + 1,
+                    prompt_version=prompt.version,
+                )
+                baseline_score = f"{prompt.score:.3f}" if prompt.score is not None else "N/A"
+                self._log(
+                    logging.INFO,
+                    f"Use parent prompt {prompt.version} as a baseline to generate a new prompt. Baseline score: {baseline_score}",
+                    prefix=parent_prefix,
+                )
+                grad_samples = next(grad_dataset_iterator)
+                rollout_results, _ = await self.evaluate_prompt_on_batch(
+                    prompt,
+                    resource_name,
+                    grad_samples,
+                    mode="train",
+                    prefix=parent_prefix,
+                )
+                new_prompt = await self.textual_gradient_and_apply_edit(
+                    prompt,
+                    rollout_results,
+                    prefix=parent_prefix,
+                )
+                if not new_prompt:
+                    self._log(
+                        logging.ERROR,
+                        f"Failed to compute edit for prompt: {prompt.prompt_template.template}",
+                        prefix=parent_prefix,
+                    )
+                    continue
+                new_prompt_template = PromptTemplate(template=new_prompt, engine="f-string")
+                versioned_candidate = self._create_versioned_prompt(new_prompt_template)
+                self._log(
+                    logging.INFO,
+                    f"New prompt template created from parent {prompt.version}: {versioned_candidate.version}",
+                    prefix=parent_prefix,
+                )
+                candidate_prefix = self._format_log_prefix(
+                    round_num=display_round, prompt_version=versioned_candidate.version
+                )
+                self._log(
+                    logging.INFO,
+                    f"New prompt template created from parent {prompt.version}:\n```\n{new_prompt}\n```",
+                    prefix=candidate_prefix,
+                )
+                candidates.append(versioned_candidate)
+
+        return candidates
+
+    async def _evaluate_and_select_beam(
+        self,
+        candidates: List[VersionedPromptTemplate],
+        resource_name: str,
+        val_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[VersionedPromptTemplate]:
+        """
+        Evaluate all candidate prompts on validation data and select top-k for the beam.
+
+        Args:
+            candidates: List of candidate prompts to evaluate.
+            resource_name: Name to register prompts under in the store.
+            val_dataset_iterator: Iterator over validation data batches.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of top beam_width prompts sorted by validation score (best first).
+
+        Raises:
+            ValueError: If no candidates remain after evaluation.
+        """
+        display_round = round_num + 1
+        round_prefix = self._format_log_prefix(round_num=display_round)
+        self._log(
+            logging.INFO,
+            f"Evaluating {len(candidates)} candidates on validation dataset",
+            prefix=round_prefix,
+        )
+
+        val_batch = next(val_dataset_iterator)
+
+        for prompt in candidates:
+            candidate_prefix = self._format_log_prefix(
+                round_num=display_round,
+                prompt_version=prompt.version,
+            )
+            _, score = await self.evaluate_prompt_on_batch(
+                prompt,
+                resource_name,
+                val_batch,
+                mode="val",
+                prefix=candidate_prefix,
+            )
+            prompt.score = score
+            self._log(
+                logging.INFO,
+                f"Candidate score: {score:.3f}",
+                prefix=candidate_prefix,
+            )
+
+        # Sort by score (descending) and select top beam_width
+        sorted_prompts = [p for p in sorted(candidates, key=lambda x: cast(float, x.score), reverse=True)]
+        selected_prompts = sorted_prompts[: self.beam_width]
+        selected_versions = [
+            f"{prompt.version}:{prompt.score:.3f}" if prompt.score is not None else prompt.version
+            for prompt in selected_prompts
+        ]
+        self._log(
+            logging.INFO,
+            f"Top {len(selected_prompts)} candidates on validation dataset: {selected_versions}",
+            prefix=round_prefix,
+        )
+
+        if len(selected_prompts) == 0:
+            raise ValueError("No beam candidates any more")
+
+        return selected_prompts
+
+    async def _update_best_prompt(
+        self,
+        beam: List[VersionedPromptTemplate],
+        resource_name: str,
+        val_dataset: Dataset[T_task],
+        round_num: int,
+    ) -> None:
+        """
+        Evaluate the best prompt in the beam on the full validation set and update history.
+
+        Args:
+            beam: Current beam of prompts (sorted, best first).
+            resource_name: Name to register prompts under in the store.
+            val_dataset: Full validation dataset.
+            round_num: Current round number (for logging, 0-indexed).
+        """
+        display_round = round_num + 1
+        best_prompt = beam[0]
+        prefix = self._format_log_prefix(round_num=display_round, prompt_version=best_prompt.version)
+        _, best_score = await self.evaluate_prompt_on_batch(
+            best_prompt,
+            resource_name,
+            cast(Sequence[T_task], val_dataset),
+            mode="val",
+            prefix=prefix,
+        )
+        self._log(
+            logging.INFO,
+            f"Beam leader score: {best_score:.3f}",
+            prefix=prefix,
+        )
+
+        if best_score > self._history_best_score:
+            prev = self._history_best_score
+            self._log(
+                logging.INFO,
+                f"Best prompt updated. New best score: {best_score:.3f} (prev: {prev:.3f})",
+                prefix=prefix,
+            )
+            self._history_best_prompt = best_prompt.prompt_template
+            self._history_best_score = best_score
+            self._history_best_version = best_prompt.version
+        else:
+            self._log(
+                logging.WARNING,
+                f"Best prompt not updated. Current score: {best_score:.3f} vs. history best: {self._history_best_score:.3f})",
+                prefix=prefix,
+            )
+
+    async def run(
+        self,
+        train_dataset: Optional[Dataset[T_task]] = None,
+        val_dataset: Optional[Dataset[T_task]] = None,
+    ) -> None:
+        """
+        Execute the APO algorithm to optimize prompts through beam search with textual gradients.
+
+        The algorithm performs iterative prompt optimization over multiple rounds:
+
+        - Each round: samples parent prompts, generates new candidates via textual gradients,
+          evaluates all candidates on validation data, and keeps the top performers
+        - Tracks the historically best prompt across all rounds
+        - Uses different training data samples for each gradient computation to ensure diversity
+
+        Args:
+            train_dataset: Dataset of tasks for computing textual gradients. Required.
+            val_dataset: Dataset of tasks for evaluating and selecting prompts. Required.
+
+        Raises:
+            ValueError: If train_dataset or val_dataset is None, or if resources are not set.
+        """
+        # Initialize beam search
+        resource_name, seed_prompt, grad_iterator, val_iterator = self._initialize_beam(train_dataset, val_dataset)
+
+        if self._poml_trace:
+            poml.set_trace(trace_dir="pomltrace")
+
+        # Validation datasets are guaranteed to be non-None after initialization
+        assert val_dataset is not None
+
+        # Start with seed prompt in the beam
+        seed_versioned = self._create_versioned_prompt(seed_prompt)
+        beam: List[VersionedPromptTemplate] = [seed_versioned]
+        self._history_best_prompt = seed_prompt
+        self._history_best_version = seed_versioned.version
+
+        # Optionally evaluate seed prompt on validation set to establish baseline
+        if self.run_initial_validation:
+            seed_prefix = self._format_log_prefix(round_num=0, prompt_version=seed_versioned.version)
+            self._log(
+                logging.INFO,
+                "Evaluating seed prompt on validation dataset before optimization...",
+                prefix=seed_prefix,
+            )
+            _, seed_score = await self.evaluate_prompt_on_batch(
+                seed_versioned,
+                resource_name,
+                cast(Sequence[T_task], val_dataset),
+                mode="val",
+                prefix=seed_prefix,
+            )
+            self._log(
+                logging.INFO,
+                f"Seed prompt baseline score: {seed_score:.3f}",
+                prefix=seed_prefix,
+            )
+            self._history_best_prompt = seed_prompt
+            self._history_best_score = seed_score
+            self._history_best_version = seed_versioned.version
+
+        # Run beam search for specified number of rounds
+        for rnd in range(self.beam_rounds):
+            display_round = rnd + 1
+            round_prefix = self._format_log_prefix(round_num=display_round)
+            self._log(
+                logging.INFO,
+                f"Round {display_round}/{self.beam_rounds}...",
+                prefix=round_prefix,
+            )
+
+            # Sample parent prompts from current beam
+            parent_prompts = self._sample_parent_prompts(beam, rnd)
+
+            # Generate new candidate prompts from parents
+            new_candidates = await self._generate_candidate_prompts(parent_prompts, resource_name, grad_iterator, rnd)
+
+            # Combine existing beam with new candidates
+            all_candidates = [*beam, *new_candidates]
+
+            # Evaluate and select top-k prompts for next beam
+            beam = await self._evaluate_and_select_beam(all_candidates, resource_name, val_iterator, rnd)
+
+            # Update historically best prompt if improved
+            await self._update_best_prompt(beam, resource_name, val_dataset, rnd)
@@ -0,0 +1,22 @@
+<poml>
+  <p>Revise the given prompt template using the critique as constraints and improvement guide.</p>
+  <cp caption="Revision Rules">
+    <list listStyle="decimal">
+      <item>Rewrite or restructure the prompt if critique implies it.</item>
+      <item>Explicitly include any requested output format, structure, or word limit, if requested by the critique.</item>
+      <item>Prioritize mechanism-first phrasing: define what to do, then how to do it.</item>
+      <item>Preserve placeholder variables inside curly brackets.</item>
+    </list>
+  </cp>
+  <output-format>
+    Return only the improved prompt template with placeholders intact. Do not include other explanations on how you did it, or headers and introductory texts.
+  </output-format>
+  <human-msg>
+    <cp caption="Prompt Template">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="Critique">
+      <text whiteSpace="pre">{{ critique }}</text>
+    </cp>
+  </human-msg>
+</poml>
@@ -0,0 +1,18 @@
+<!-- Conservative Edit Prompt -->
+
+<poml>
+  <p>Revise the prompt to address ONE critique point clearly and effectively. Preserve all variable names in curly-brackets.</p>
+  <p>Do not address more than one critique point. Focus on the single most critical issue.</p>
+  <p>Keep the new prompt close in tone, length, and structure to the original.</p>
+  <output-format>
+    Return only the revised full prompt. Do not include explanations, comparisons, or other text.
+  </output-format>
+  <human-msg>
+    <cp caption="PROMPT" level="3">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="CRITIQUE" level="3">
+      <text whiteSpace="pre">{{ critique }}</text>
+    </cp>
+  </human-msg>
+</poml>
@@ -0,0 +1,18 @@
+<poml>
+  <p>You optimize a prompt template.</p>
+  <cp caption="Original Prompt Template">
+    <text whiteSpace="pre">{{ prompt_template }}</text>
+  </cp>
+  <cp caption="Experiments with Original Prompt Template">
+    <cp for="experiment in experiments" caption="Experiment {{ loop.index + 1 }}">
+      <p>This experiment has {{ experiment.status }}. It gets a final reward: {{ experiment.final_reward }}</p>
+      <cp caption="Rollout Traces (Chat Messages, Grader Requests included)">
+        <object data="{{ experiment.messages }}" />
+      </cp>
+    </cp>
+  </cp>
+  <cp caption="Your Task">
+    Produce a brief critique listing specific causes for the error or ways to raise reward next time.
+    Return a bullet list with concrete, testable changes (format, constraints, ordering, definitions).
+  </cp>
+</poml>
@@ -0,0 +1,16 @@
+<poml>
+  <role>You are a prompt engineer.</role>
+  <task>Analyze where the current prompt failed to elicit the right mechanism.</task>
+  <cp caption="Current Prompt Template">
+    <text whiteSpace="pre">{{ prompt_template }}</text>
+  </cp>
+  <cp caption="Sample Runs with Current Prompt Template">
+    <p>The following are the OpenTelemetry spans collected from the sample runs with the current prompt template. They should contain both prompt, responses and rewards.</p>
+    <cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }} Diagnostics">
+      <object for="span in experiment.spans" data="{{ span }}" />
+    </cp>
+  </cp>
+  <output-format>
+    Write 3-5 short bullets titled 'Critique:' focusing on missing constraints, ordering, or formatting.
+  </output-format>
+</poml>
@@ -0,0 +1,107 @@
+<poml>
+
+  <role>You are an expert prompt engineer.</role>
+
+  <task>Your task is to analyze the prompt and provide a critique of the prompt. Follow the steps below to create the critique.
+
+  <cp caption="1. Structural Issues">
+    <p>These flaws block clarity and logic. Always check them first.</p>
+
+    <list>
+      <item><b>Missing goal</b>: The prompt never defines what success looks like. Ask: <i>Can I summarize its output goal in one line?</i></item>
+      <item><b>Contradictions</b>: Two or more instructions conflict. Search for words like *never*, *always*, *except*, *but also*.</item>
+      <item><b>Circular dependencies</b>: The model is told to do A before B and B before A.</item>
+      <item><b>No stop condition</b>: The prompt doesn’t say when the task is done. Flag any open-ended verbs: <i>explore,</i> <i>analyze further,</i> <i>continue indefinitely.</i></item>
+    </list>
+  </cp>
+
+  <cp caption="2. Instruction Quality">
+    <p>Examine how the instructions are stated and ordered to ensure clarity and enforceability.</p>
+    <list>
+      <item><b>Vague verbs</b>: Avoid terms like <i>optimize,</i> <i>improve,</i> and <i>ensure.</i> Use precise, measurable instructions.</item>
+      <item><b>Lack of hierarchy</b>: All rules appear equally important, making conflict resolution impossible. Clarify rule precedence.</item>
+      <item><b>Mixed abstraction</b>: High-level policies are interleaved with implementation details. Keep principles separate from step-by-step actions.</item>
+      <item><b>Overlapping scope</b>: Similar instructions appear in several sections with minor changes. Identify and consolidate duplicates.</item>
+    </list>
+  </cp>
+
+  <cp caption="3. Control and Behavior">
+    <p>Review boundaries on model autonomy, tool use, and communication style.</p>
+    <list>
+      <item><b>No tool limits</b>: Limits on tool calls, retries, or time not specified. Define boundaries for operations.</item>
+      <item><b>Unclear uncertainty handling</b>: Conflicting instructions regarding clarifying uncertainties vs. never asking users. Select one behavior.</item>
+      <item><b>Verbosity confusion</b>: Some parts demand detailed answers, others specify brevity. Highlight and resolve inconsistency.</item>
+      <item><b>Feedback omission</b>: No plan for progress reporting or preamble during multi-step operations.</item>
+    </list>
+  </cp>
+
+  <cp caption="4. Input and Output Specification">
+    <p>Assess if required data and expected output formats are clearly defined.</p>
+    <list>
+      <item><b>No input defaults</b>: What should happen if a needed value is absent or invalid isn’t explained.</item>
+      <item><b>Output schema missing</b>: Expected response format or sections are not spelled out.</item>
+      <item><b>Format inconsistency</b>: Output style (Markdown, JSON, XML, etc.) shifts mid-prompt. Ensure format requirements are stable.</item>
+      <item><b>No validation</b>: Lacks steps like <i>verify results before submitting</i> or <i>summarize at end.</i></item>
+    </list>
+  </cp>
+
+  <cp caption="5. Scope and Safety">
+    <p>Ensure prompt actions remain within safe, authorized boundaries.</p>
+    <list>
+      <item><b>Scope creep</b>: Open-ended statements such as <i>feel free to enhance</i> can justify unrelated changes.</item>
+      <item><b>Unsafe actions</b>: Allows deletions or modifications without explicit user approval.</item>
+      <item><b>No error handling</b>: What happens if a tool call fails or data is missing is not addressed.</item>
+      <item><b>User authority ambiguity</b>: Model may act for multiple users or perform irreversible actions without checks.</item>
+    </list>
+  </cp>
+
+  <cp caption="6. Efficiency and Maintainability">
+    <p>Consider the prompt’s length, redundancy, and future comprehensibility.</p>
+    <list>
+      <item><b>Overexplained</b>: Verbose explanations where concise, numbered steps suffice.</item>
+      <item><b>Redundancy</b>: Similar rules scattered in multiple aliases; centralize and summarize them.</item>
+      <item><b>Hidden assumptions</b>: Implicit defaults (like timezone, language) are not stated.</item>
+      <item><b>Poor auditability</b>: Lacks section markers (e.g., <code>&lt;policy&gt;</code>, <code>&lt;procedure&gt;</code>). Structure prompt for easy review.</item>
+    </list>
+  </cp>
+
+  <cp caption="7. Testing Method">
+    <p>Methodical approach for reviewing a prompt:</p>
+    <list>
+      <item>Read the prompt fully; highlight all unclear or contradictory instructions.</item>
+      <item>For each main area, answer:
+      <list listStyle="decimal">
+        <item>What is the intended outcome?</item>
+        <item>What is the stop or completion condition?</item>
+        <item>How are conflicts between rules resolved?</item>
+        <item>What are the explicit limits (tools, run time, tokens)?</item>
+        <item>What should the output format be?</item>
+      </list>
+      </item>
+      <item>Rate each section: <i>clear</i>, <i>incomplete</i>, <i>contradictory</i>, or <i>redundant</i>.</item>
+      <item>Summarize findings under categories: structure, control, scope, format, safety.</item>
+    </list>
+    <p>This method surfaces issues such as ambiguity, contradiction, missing boundaries, and output uncertainty—core failure modes in prompting identified by the GPT-5 prompting guide.</p>
+  </cp>
+  </task>
+
+  <output-format>
+    Respond with a complete analysis and critique of the prompt. Be concise and direct. Less than 350 words.
+  </output-format>
+
+  <human-msg>
+    <cp caption="Prompt">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="Sample Runs of the Prompts (Historical Messages and Rewards)">
+      <cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }}">
+        <cp caption="Overall Status">
+          This run has {{ experiment.status }}. The final score is {{ experiment.final_reward }}.
+        </cp>
+        <cp caption="Messages">
+          <object data="{{ experiment.messages }}" />
+        </cp>
+      </cp>
+    </cp>
+  </human-msg>
+</poml>
@@ -0,0 +1,162 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import inspect
+import weakref
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Optional,
+    Union,
+)
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.client import AgentLightningClient
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset, NamedResources
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.trainer import Trainer
+
+
+class Algorithm:
+    """Algorithm is the strategy, or tuner to train the agent."""
+
+    _trainer_ref: weakref.ReferenceType[Trainer] | None = None
+    _llm_proxy_ref: weakref.ReferenceType["LLMProxy"] | None = None
+    _store: LightningStore | None = None
+    _initial_resources: NamedResources | None = None
+    _adapter_ref: weakref.ReferenceType[TraceAdapter[Any]] | None = None
+
+    def is_async(self) -> bool:
+        """Return True if the algorithm is asynchronous."""
+        return inspect.iscoroutinefunction(self.run)
+
+    def set_trainer(self, trainer: Trainer) -> None:
+        """
+        Set the trainer for this algorithm.
+
+        Args:
+            trainer: The Trainer instance that will handle training and validation.
+        """
+        self._trainer_ref = weakref.ref(trainer)
+
+    def get_trainer(self) -> Trainer:
+        """
+        Get the trainer for this algorithm.
+
+        Returns:
+            The Trainer instance associated with this agent.
+        """
+        if self._trainer_ref is None:
+            raise ValueError("Trainer has not been set for this agent.")
+        trainer = self._trainer_ref()
+        if trainer is None:
+            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
+        return trainer
+
+    def set_llm_proxy(self, llm_proxy: LLMProxy | None) -> None:
+        """
+        Set the LLM proxy for this algorithm to reuse when available.
+
+        Args:
+            llm_proxy: The LLMProxy instance configured by the trainer, if any.
+        """
+        self._llm_proxy_ref = weakref.ref(llm_proxy) if llm_proxy is not None else None
+
+    def get_llm_proxy(self) -> Optional[LLMProxy]:
+        """
+        Retrieve the configured LLM proxy instance, if one has been set.
+
+        Returns:
+            The active LLMProxy instance or None when not configured.
+        """
+        if self._llm_proxy_ref is None:
+            return None
+
+        llm_proxy = self._llm_proxy_ref()
+        if llm_proxy is None:
+            raise ValueError("LLM proxy reference is no longer valid (object has been garbage collected).")
+
+        return llm_proxy
+
+    def set_adapter(self, adapter: TraceAdapter[Any]) -> None:
+        """
+        Set the adapter for this algorithm to collect and convert traces.
+        """
+        self._adapter_ref = weakref.ref(adapter)
+
+    def get_adapter(self) -> TraceAdapter[Any]:
+        """
+        Retrieve the adapter for this algorithm to communicate with the runners.
+        """
+        if self._adapter_ref is None:
+            raise ValueError("Adapter has not been set for this algorithm.")
+        adapter = self._adapter_ref()
+        if adapter is None:
+            raise ValueError("Adapter reference is no longer valid (object has been garbage collected).")
+        return adapter
+
+    def set_store(self, store: LightningStore) -> None:
+        """
+        Set the store for this algorithm to communicate with the runners.
+
+        Store is set directly instead of using weakref because its copy is meant to be
+        maintained throughout the algorithm's lifecycle.
+        """
+        self._store = store
+
+    def get_store(self) -> LightningStore:
+        """
+        Retrieve the store for this algorithm to communicate with the runners.
+        """
+        if self._store is None:
+            raise ValueError("Store has not been set for this algorithm.")
+        return self._store
+
+    def get_initial_resources(self) -> Optional[NamedResources]:
+        """
+        Get the initial resources for this algorithm.
+        """
+        return self._initial_resources
+
+    def set_initial_resources(self, resources: NamedResources) -> None:
+        """
+        Set the initial resources for this algorithm.
+        """
+        self._initial_resources = resources
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.run(*args, **kwargs)
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Union[None, Awaitable[None]]:
+        """Subclasses should implement this method to implement the algorithm.
+
+        Args:
+            train_dataset: The dataset to train on. Not all algorithms require a training dataset.
+            val_dataset: The dataset to validate on. Not all algorithms require a validation dataset.
+
+        Returns:
+            Algorithm should refrain from returning anything. It should just run the algorithm.
+        """
+        raise NotImplementedError("Subclasses must implement run().")
+
+    def get_client(self) -> AgentLightningClient:
+        """Get the client to communicate with the algorithm.
+
+        If the algorithm does not require a server-client communication, it can also create a mock client
+        that never communicates with itself.
+
+        Deprecated and will be removed in a future version.
+
+        Returns:
+            The AgentLightningClient instance associated with this algorithm.
+        """
+        raise NotImplementedError("Subclasses must implement get_client().")
@@ -0,0 +1,264 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Dict,
+    Generic,
+    Literal,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset, NamedResources
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+
+from .base import Algorithm
+
+# Algorithm function signature types
+# We've missed a lot of combinations here.
+# Let's add them in future.
+
+
+class AlgorithmFuncSyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyDataset(Protocol):
+    def __call__(self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None: ...
+
+
+class AlgorithmFuncAsyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyDataset(Protocol):
+    def __call__(
+        self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]
+    ) -> Awaitable[None]: ...
+
+
+AlgorithmFuncAsync = Union[AlgorithmFuncAsyncOnlyStore, AlgorithmFuncAsyncOnlyDataset, AlgorithmFuncAsyncFull]
+
+AlgorithmFuncSync = Union[AlgorithmFuncSyncOnlyStore, AlgorithmFuncSyncOnlyDataset, AlgorithmFuncSyncFull]
+
+
+class AlgorithmFuncSyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+class AlgorithmFuncAsyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]: ...
+
+
+AlgorithmFuncSyncLike = Union[AlgorithmFuncSync, AlgorithmFuncSyncFallback]
+AlgorithmFuncAsyncLike = Union[AlgorithmFuncAsync, AlgorithmFuncAsyncFallback]
+
+AlgorithmFunc = Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]
+
+
+AsyncFlag = Literal[True, False]
+AF = TypeVar("AF", bound=AsyncFlag)
+
+
+class FunctionalAlgorithm(Algorithm, Generic[AF]):
+    """An algorithm wrapper built from a callable implementation.
+
+    Functional algorithms let you provide an ordinary function instead of
+    subclassing [`Algorithm`][agentlightning.Algorithm]. The wrapper inspects
+    the callable signature to supply optional dependencies
+    such as the store, adapter, and LLM proxy.
+    """
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[False]]", algorithm_func: AlgorithmFuncSyncLike) -> None: ...
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...
+
+    def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
+        """Wrap a function that implements algorithm behaviour.
+
+        Args:
+            algorithm_func: Sync or async callable implementing the algorithm
+                contract. Arguments are detected automatically based on the
+                function signature.
+        """
+        super().__init__()
+        self._algorithm_func = algorithm_func
+        self._sig = inspect.signature(algorithm_func)
+        self._is_async = inspect.iscoroutinefunction(algorithm_func)
+
+        # Copy function metadata to preserve type hints and other attributes
+        functools.update_wrapper(self, algorithm_func)  # type: ignore
+
+    def is_async(self) -> bool:
+        return self._is_async
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[False]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None: ...
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[True]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Awaitable[None]: ...
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._algorithm_func(*args, **kwargs)  # type: ignore
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Union[None, Awaitable[None]]:
+        """Execute the wrapped function with injected dependencies.
+
+        Args:
+            train_dataset: Optional training dataset passed through when the
+                callable declares a `train_dataset` parameter.
+            val_dataset: Optional validation dataset passed through when the
+                callable declares a `val_dataset` parameter.
+
+        Returns:
+            None for sync callables or an awaitable when the callable is async.
+
+        Raises:
+            TypeError: If a dataset is provided but the function signature does
+                not accept the corresponding argument.
+        """
+        kwargs: Dict[str, Any] = {}
+        if "store" in self._sig.parameters:
+            kwargs["store"] = self.get_store()
+        if "adapter" in self._sig.parameters:
+            kwargs["adapter"] = self.get_adapter()
+        if "llm_proxy" in self._sig.parameters:
+            kwargs["llm_proxy"] = self.get_llm_proxy()
+        if "initial_resources" in self._sig.parameters:
+            kwargs["initial_resources"] = self.get_initial_resources()
+        if "train_dataset" in self._sig.parameters:
+            kwargs["train_dataset"] = train_dataset
+        elif train_dataset is not None:
+            raise TypeError(
+                f"train_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        if "val_dataset" in self._sig.parameters:
+            kwargs["val_dataset"] = val_dataset
+        elif val_dataset is not None:
+            raise TypeError(
+                f"val_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        # both sync and async functions can be called with the same signature
+        result = self._algorithm_func(**kwargs)  # type: ignore[misc]
+        if self._is_async:
+            return cast(Awaitable[None], result)
+        return None
+
+
+@overload
+def algo(func: AlgorithmFuncAsync) -> FunctionalAlgorithm[Literal[True]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncAsyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSync) -> FunctionalAlgorithm[Literal[False]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+def algo(
+    func: Union[
+        AlgorithmFuncSync,
+        AlgorithmFuncAsync,
+        AlgorithmFuncSyncFallback,
+        AlgorithmFuncAsyncFallback,
+    ],
+) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
+    """Convert a callable into a [`FunctionalAlgorithm`][agentlightning.algorithm.decorator.FunctionalAlgorithm].
+
+    The decorator inspects the callable signature to decide which dependencies
+    to inject at runtime, enabling concise algorithm definitions that still
+    leverage the full training runtime.
+
+    Args:
+        func: Function implementing the algorithm logic. May be synchronous or
+            asynchronous. The function can expect all of, or a subset of the following parameters:
+
+            - `store`: [`LightningStore`][agentlightning.store.base.LightningStore],
+            - `train_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `val_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `llm_proxy`: [`LLMProxy`][agentlightning.LLMProxy],
+            - `adapter`: [`TraceAdapter`][agentlightning.TraceAdapter],
+            - `initial_resources`: [`NamedResources`][agentlightning.NamedResources],
+
+            If the function does not expect a parameter, the wrapper will not inject it into the call.
+            Using `*args` and `**kwargs` will not work and no parameters will be injected.
+
+    Returns:
+        FunctionalAlgorithm that proxies the callable while exposing the
+        `Algorithm` interface.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.decorator import algo
+
+        @algo
+        def batching_algorithm(*, store, train_dataset, val_dataset):
+            for sample in train_dataset:
+                store.enqueue_rollout(input=sample, mode="train")
+
+        @algo
+        async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
+            await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
+        ```
+    """
+    return FunctionalAlgorithm(func)
@@ -0,0 +1,241 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import Any, List, Literal, Optional
+
+from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
+
+from .base import Algorithm
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["FastAlgorithm", "Baseline"]
+
+
+class FastAlgorithm(Algorithm):
+    """Base class for lightweight algorithms optimised for developer workflows.
+
+    Fast algorithms prioritise short feedback loops so an agent developer can run
+    small-scale experiments without waiting for long-running training jobs to
+    finish.
+    """
+
+
+def _timestamp_to_iso_str(timestamp: float) -> str:
+    return datetime.fromtimestamp(timestamp).isoformat()
+
+
+class Baseline(FastAlgorithm):
+    """Reference implementation that streams the full dataset through the rollout queue.
+
+    The baseline algorithm batches task submissions, waits for each rollout to
+    finish, and logs every collected span and reward. It is primarily useful as
+    a smoke test for the platform plumbing rather than a performant trainer.
+
+    Args:
+        n_epochs: Number of dataset passes to execute for both the train and val
+            splits during developer experiments.
+        train_split: Fraction of the concatenated dataset to treat as training
+            data. Must be strictly between 0 and 1.
+        polling_interval: Interval, in seconds, to poll the store for queue
+            depth and rollout completion.
+        max_queue_length: Number of rollouts allowed to wait in the queue before
+            throttling additional submissions.
+        span_verbosity: Level of detail to include when logging span metadata.
+
+    Raises:
+        ValueError: If `train_split` falls outside the `(0, 1)` interval.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.fast import Baseline
+
+        algorithm = Baseline(n_epochs=2, train_split=0.8, span_verbosity="key_values")
+        trainer.fit(algorithm, train_dataset=my_train, val_dataset=my_val)
+        ```
+    """
+
+    def __init__(
+        self,
+        *,
+        n_epochs: int = 1,
+        train_split: float = 0.5,
+        polling_interval: float = 5.0,
+        max_queue_length: int = 4,
+        span_verbosity: Literal["keys", "key_values", "none"] = "keys",
+    ) -> None:
+        super().__init__()
+        self.n_epochs = n_epochs
+        self.train_split = train_split
+        self.polling_interval = polling_interval
+        self.max_queue_length = max_queue_length
+        self.span_verbosity = span_verbosity
+        if not (0.0 < self.train_split < 1.0):
+            raise ValueError("train_split must be between 0 and 1.")
+
+        self._finished_rollout_count = 0
+
+    def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str:
+        """Format a span for logging based on the configured verbosity."""
+        if self.span_verbosity == "none":
+            return ""
+
+        prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
+        elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
+
+        msg = (
+            prefix_msg
+            + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
+            + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
+            + f"{elapsed} seconds. "
+        )
+        if self.span_verbosity == "key_values":
+            msg += f"Attributes: {span.attributes}"
+        else:
+            msg += f"Attribute keys: {list(span.attributes.keys())}"
+        return msg
+
+    async def _handle_rollout_finish(self, rollout: Rollout) -> None:
+        """Log attempt metadata and emit adapted traces when a rollout ends."""
+        store = self.get_store()
+
+        rollout_id = rollout.rollout_id
+        rollout_end_time = rollout.end_time or asyncio.get_event_loop().time()
+        logger.info(
+            f"[Rollout {rollout_id}] Finished with status {rollout.status} in {rollout_end_time - rollout.start_time:.2f} seconds."
+        )
+
+        # Logs all the attempts and their corresponding spans
+        attempts = await store.query_attempts(rollout_id)
+        for attempt in attempts:
+            logger.info(
+                "[Rollout %s | Attempt %s] ID: %s. Status: %s. Worker: %s",
+                rollout_id,
+                attempt.sequence_id,
+                attempt.attempt_id,
+                attempt.status,
+                attempt.worker_id,
+            )
+            spans = await store.query_spans(rollout_id=rollout_id)
+            for span in spans:
+                if self.span_verbosity != "none":
+                    logger.info(self._span_to_string(rollout.rollout_id, attempt, span))
+
+        # Attempts to adapt the spans using the adapter if provided
+        try:
+            adapter = self.get_adapter()
+        except ValueError:
+            logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
+            adapter = None
+        if adapter is not None:
+            spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+            transformed_data = adapter.adapt(spans)
+            logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}")
+
+    async def _enqueue_rollouts(
+        self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
+    ) -> None:
+        """Submit rollouts while respecting the maximum queue length."""
+        store = self.get_store()
+
+        for index in train_indices + val_indices:
+            queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+            if len(queuing_rollouts) <= 1:
+                # Only enqueue a new rollout when there is at most 1 rollout in the queue.
+                sample = dataset[index]
+                mode = "train" if index in train_indices else "val"
+                rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                logger.info(f"[Rollout {rollout.rollout_id}] Enqueued in {mode} mode with sample: {sample}")
+            await asyncio.sleep(self.polling_interval)
+
+    async def _harvest_rollout_spans(self, rollout_id: str):
+        """Poll rollout status updates until completion and log transitions."""
+        store = self.get_store()
+        last_status: Optional[RolloutStatus] = None
+
+        while True:
+            rollout = await store.get_rollout_by_id(rollout_id)
+            if rollout is not None:
+                if rollout.status in ["succeeded", "failed", "cancelled"]:
+                    # Rollout is finished, log all the data.
+                    await self._handle_rollout_finish(rollout)
+                    # We are done here.
+                    self._finished_rollout_count += 1
+                    logger.info(f"Finished {self._finished_rollout_count} rollouts.")
+                    break
+
+                if last_status != rollout.status:
+                    if last_status is not None:
+                        logger.info(f"[Rollout {rollout_id}] Status changed to {rollout.status}.")
+                    else:
+                        logger.info(f"[Rollout {rollout_id}] Status is initialized to {rollout.status}.")
+                    last_status = rollout.status
+
+                else:
+                    logger.debug(f"[Rollout {rollout_id}] Status is still {rollout.status}.")
+
+            await asyncio.sleep(self.polling_interval)
+
+    async def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None:
+        """Execute the baseline loop across the provided datasets."""
+        train_dataset_length = len(train_dataset) if train_dataset is not None else 0
+        val_dataset_length = len(val_dataset) if val_dataset is not None else 0
+        if train_dataset_length == 0 and val_dataset_length == 0:
+            logger.error(
+                "MockAlgorithm requires at least one dataset. Provide train_dataset or val_dataset before running."
+            )
+            return
+
+        concatenated_dataset = [train_dataset[i] for i in range(train_dataset_length) if train_dataset is not None] + [
+            val_dataset[i] for i in range(val_dataset_length) if val_dataset is not None
+        ]
+        train_indices = list(range(0, train_dataset_length))
+        val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length))
+        logger.debug(f"Train indices: {train_indices}")
+        logger.debug(f"Val indices: {val_indices}")
+
+        store = self.get_store()
+
+        # Currently we only supports a single resource update at the start.
+        initial_resources = self.get_initial_resources()
+        if initial_resources is not None:
+            resource_update = await store.update_resources("default", initial_resources)
+            resources_id = resource_update.resources_id
+            logger.info(f"Initial resources set: {initial_resources}")
+        else:
+            logger.warning("No initial resources provided. Skip initializing resources.")
+            resources_id = None
+
+        for epoch in range(self.n_epochs):
+            harvest_tasks: List[asyncio.Task[None]] = []
+            logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.")
+            for index in train_indices + val_indices:
+                logger.info(
+                    f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
+                )
+                while True:
+                    queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+                    if len(queuing_rollouts) <= self.max_queue_length:
+                        # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
+                        sample = concatenated_dataset[index]
+                        mode = "train" if index in train_indices else "val"
+                        rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                        harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
+                        logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
+                        break
+                    else:
+                        # Sleep a bit and try again later.
+                        await asyncio.sleep(self.polling_interval)
+
+            # Wait for all harvest tasks to complete
+            logger.info(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...")
+            if len(harvest_tasks) > 0:
+                await asyncio.gather(*harvest_tasks)
@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import random
+from typing import Iterator, List, Sequence, TypeVar
+
+from agentlightning.types import Dataset
+
+T_task = TypeVar("T_task")
+
+
+def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterator[Sequence[T_task]]:
+    """
+    Create an infinite iterator that yields batches from the dataset.
+
+    When batch_size >= dataset size, yields the entire shuffled dataset repeatedly.
+    When batch_size < dataset size, yields batches of the specified size, reshuffling
+    after each complete pass through the dataset.
+
+    Args:
+        dataset: The dataset to iterate over.
+        batch_size: The desired batch size.
+
+    Yields:
+        Sequences of tasks from the dataset. Each task appears at most once per epoch.
+    """
+    if batch_size >= len(dataset):
+        while True:
+            dataset_copy = [dataset[i] for i in range(len(dataset))]
+            random.shuffle(dataset_copy)
+            yield dataset_copy
+
+    else:
+        current_batch: List[int] = []
+        while True:
+            indices = list(range(len(dataset)))
+            random.shuffle(indices)
+            for index in indices:
+                if index in current_batch:
+                    continue
+                current_batch.append(index)
+                if len(current_batch) == batch_size:
+                    yield [dataset[index] for index in current_batch]
+                    current_batch = []
@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .interface import VERL
+
+__all__ = ["VERL"]
@@ -0,0 +1,154 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any, Optional
+
+from hydra import compose, initialize
+from omegaconf import OmegaConf
+
+from agentlightning.algorithm.base import Algorithm
+from agentlightning.client import AgentLightningClient
+from agentlightning.types import Dataset
+from agentlightning.verl.entrypoint import run_ppo  # type: ignore
+
+
+class VERL(Algorithm):
+    """VERL-powered algorithm that delegates training to the VERL PPO runner.
+
+    !!! warning
+        Advanced customisation currently requires copying the VERL source and
+        modifying it directly. Native hooks for overriding training behaviour
+        will land in a future release.
+
+    Args:
+        config: Dictionary mirroring the overrides passed to the VERL CLI. The
+            overrides are merged with VERL's packaged defaults via Hydra before
+            launching training.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.verl import VERL
+
+        algorithm = VERL(
+            config={
+                "algorithm": {
+                    "adv_estimator": "grpo",
+                    "use_kl_in_reward": False,
+                },
+                "data": {
+                    "train_batch_size": 32,
+                    "max_prompt_length": 4096,
+                    "max_response_length": 2048,
+                },
+                "actor_rollout_ref": {
+                    "rollout": {
+                        "tensor_model_parallel_size": 1,
+                        "n": 4,
+                        "log_prob_micro_batch_size_per_gpu": 4,
+                        "multi_turn": {"format": "hermes"},
+                        "name": "vllm",
+                        "gpu_memory_utilization": 0.6,
+                    },
+                    "actor": {
+                        "ppo_mini_batch_size": 32,
+                        "ppo_micro_batch_size_per_gpu": 4,
+                        "optim": {"lr": 1e-6},
+                        "use_kl_loss": False,
+                        "kl_loss_coef": 0.0,
+                        "entropy_coeff": 0,
+                        "clip_ratio_low": 0.2,
+                        "clip_ratio_high": 0.3,
+                        "fsdp_config": {
+                            "param_offload": True,
+                            "optimizer_offload": True,
+                        },
+                    },
+                    "ref": {
+                        "log_prob_micro_batch_size_per_gpu": 8,
+                        "fsdp_config": {"param_offload": True},
+                    },
+                    "model": {
+                        "path": "Qwen/Qwen2.5-1.5B-Instruct",
+                        "use_remove_padding": True,
+                        "enable_gradient_checkpointing": True,
+                    },
+                },
+                "trainer": {
+                    "n_gpus_per_node": 1,
+                    "val_before_train": True,
+                    "critic_warmup": 0,
+                    "logger": ["console", "wandb"],
+                    "project_name": "AgentLightning",
+                    "experiment_name": "calc_x",
+                    "nnodes": 1,
+                    "save_freq": 64,
+                    "test_freq": 32,
+                    "total_epochs": 2,
+                },
+            }
+        )
+        trainer.fit(algorithm, train_dataset=my_train_dataset)
+        ```
+    """
+
+    def __init__(self, config: dict[str, Any]):
+        super().__init__()
+
+        # Compose the base config exactly like your decorator:
+        with initialize(version_base=None, config_path="pkg://agentlightning/verl"):
+            base_cfg = compose(config_name="config")
+
+        # Merge your dict overrides
+        override_conf = OmegaConf.create(config)
+        # Allow adding new fields
+        OmegaConf.set_struct(base_cfg, False)
+        self.config = OmegaConf.merge(base_cfg, override_conf)
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None:
+        """Launch the VERL PPO entrypoint with the configured runtime context.
+
+        Args:
+            train_dataset: Optional dataset forwarded to VERL for training.
+            val_dataset: Optional dataset forwarded to VERL for evaluation.
+
+        Raises:
+            ValueError: If required dependencies such as the store, LLM proxy, or
+                adapter have been garbage-collected when using the V1 execution
+                mode.
+        """
+        try:
+            store = self.get_store()
+        except Exception:
+            print("Store is not set. Assuming v0 execution mode.")
+            run_ppo(
+                self.config,
+                train_dataset=train_dataset,
+                val_dataset=val_dataset,
+                store=None,
+                llm_proxy=None,
+                adapter=None,
+            )
+        else:
+            print("Store is set. Assuming v1 execution mode.")
+            llm_proxy = self.get_llm_proxy()
+            adapter = self.get_adapter()
+            run_ppo(
+                self.config,
+                train_dataset=train_dataset,
+                val_dataset=val_dataset,
+                store=store,
+                llm_proxy=llm_proxy,
+                adapter=adapter,
+            )
+
+    def get_client(self) -> AgentLightningClient:
+        """Create a client bound to the VERL-managed Agent Lightning server.
+
+        Deprecated:
+            Since v0.2.
+        """
+        port = self.config.agentlightning.port
+        return AgentLightningClient(endpoint=f"http://localhost:{port}")
@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent Lightning command line interface entry point."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from typing import Dict, Iterable, Tuple
+
+_SUBCOMMANDS: Dict[str, Tuple[str, str]] = {
+    "vllm": ("agentlightning.cli.vllm", "Run the vLLM CLI with Agent Lightning instrumentation."),
+    "store": ("agentlightning.cli.store", "Run a LightningStore server."),
+    "agentops": ("agentlightning.cli.agentops_server", "Start the AgentOps server manager."),
+}
+
+_DESCRIPTION = "Agent Lightning CLI entry point.\n\nAvailable subcommands:\n" + "\n".join(
+    f"  {name:<10}{desc}" for name, (_, desc) in _SUBCOMMANDS.items()
+)
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    """Dispatch to the requested Agent Lightning subcommand."""
+    parser = argparse.ArgumentParser(
+        prog="agl",
+        description=_DESCRIPTION,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("subcommand", choices=_SUBCOMMANDS.keys(), help="Subcommand to run.")
+    parser.add_argument("args", nargs=argparse.REMAINDER, help=argparse.SUPPRESS)
+
+    parsed = parser.parse_args(list(argv) if argv is not None else None)
+    module_name, _ = _SUBCOMMANDS[parsed.subcommand]
+    module = importlib.import_module(module_name)
+
+    entry_point = getattr(module, "main", None)
+    if entry_point is None:
+        parser.error(f"Subcommand '{parsed.subcommand}' does not define a callable 'main'")
+
+    dispatch_args = parsed.args
+    original_argv = sys.argv
+    sys.argv = [f"{parser.prog} {parsed.subcommand}", *dispatch_args]
+    try:
+        result = entry_point(dispatch_args or None)
+    finally:
+        sys.argv = original_argv
+
+    if isinstance(result, int):
+        return result
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
@@ -1,20 +0,0 @@
-import time
-from agentlightning.instrumentation.agentops import AgentOpsServerManager
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Start AgentOps server")
-    parser.add_argument("--daemon", action="store_true", help="Run server as a daemon")
-    parser.add_argument("--port", type=int, default=8002, help="Port to run the server on")
-    args = parser.parse_args()
-
-    manager = AgentOpsServerManager(daemon=args.daemon, port=args.port)
-    try:
-        manager.start()
-        # Wait forever
-        while True:
-            time.sleep(1)
-    except KeyboardInterrupt:
-        manager.stop()
@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Run a LightningStore server for persistent access from multiple processes."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+from typing import Iterable
+
+from agentlightning.logging import configure_logger
+from agentlightning.store.client_server import LightningStoreServer
+from agentlightning.store.memory import InMemoryLightningStore
+
+logger = logging.getLogger(__name__)
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Run a LightningStore server")
+    parser.add_argument("--port", type=int, default=4747, help="Port to run the server on")
+    parser.add_argument(
+        "--cors-origin",
+        dest="cors_origins",
+        action="append",
+        help="Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.",
+    )
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    configure_logger()
+
+    store = InMemoryLightningStore()
+    server = LightningStoreServer(
+        store,
+        host="0.0.0.0",
+        port=args.port,
+        cors_allow_origins=args.cors_origins,
+    )
+    try:
+        asyncio.run(server.run_forever())
+    except RuntimeError as exc:
+        logger.error("LightningStore server failed to start: %s", exc, exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
@@ -1,10 +1,29 @@
-from typing import List
+# Copyright (c) Microsoft. All rights reserved.

-from vllm.entrypoints.cli.main import main
+from __future__ import annotations

-from agentlightning.instrumentation.vllm import instrument_vllm
+from typing import Iterable
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    import sys
+
+    from vllm.entrypoints.cli.main import main as vllm_main
+
+    from agentlightning.instrumentation.vllm import instrument_vllm
+
+    instrument_vllm()
+    if argv is not None:
+        original_argv = sys.argv
+        sys.argv = [original_argv[0], *list(argv)]
+        try:
+            vllm_main()
+        finally:
+            sys.argv = original_argv
+    else:
+        vllm_main()
+    return 0


 if __name__ == "__main__":
-    instrument_vllm()
-    main()
+    raise SystemExit(main())
@@ -1,26 +1,47 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utilities for interacting with legacy Agent Lightning servers.
+
+This module contains compatibility shims that speak the deprecated HTTP
+interface used by older Agent Lightning deployments. Modern code should prefer
+the store-based APIs exposed by `agentlightning.store`, but keeping these
+clients available makes it easier to migrate existing workflows incrementally.
+"""
+
 import asyncio
 import logging
 import time
 import urllib.parse
-from typing import Any, Dict, Optional, List, Union
+import warnings
+from typing import Any, Dict, List, Optional, Union

 import aiohttp
 import requests

-from .types import Rollout, Task, TaskInput, TaskIfAny, ResourcesUpdate, NamedResources
-
+from .types import NamedResources, ResourcesUpdate, RolloutLegacy, Task, TaskIfAny, TaskInput

 logger = logging.getLogger(__name__)


 class AgentLightningClient:
-    """
-    Client for interacting with a version-aware Agent Lightning Server.
+    """Client wrapper for the legacy version-aware Agent Lightning server.

-    This client handles polling for tasks, fetching specific versions of resources
-    (like model configurations), and posting completed rollouts back to the server.
-    It provides both synchronous and asynchronous methods for these operations and
-    includes a cache for resources.
+    The client exposes synchronous and asynchronous helpers for polling tasks,
+    retrieving resource bundles, and submitting rollouts. It also maintains a
+    simple in-memory cache keyed by the server-provided resource identifier to
+    avoid redundant network requests.
+
+    !!! warning "Deprecated"
+        [`AgentLightningClient`][agentlightning.client.AgentLightningClient] is part of
+        the legacy client/server stack. New code should rely on the store-based APIs
+        implemented in `agentlightning.store`.
+
+    Attributes:
+        endpoint: Base URL of the Agent Lightning server.
+        poll_interval: Delay in seconds between polling attempts when no task is
+            available.
+        timeout: Timeout in seconds applied to HTTP requests.
+        task_count: Number of tasks claimed during the lifetime of this client.
    """

    _next_task_uri = "/task"
@@ -29,13 +50,16 @@ class AgentLightningClient:
    _report_rollout_uri = "/rollout"

    def __init__(self, endpoint: str, poll_interval: float = 5.0, timeout: float = 10.0):
-        """Initializes the AgentLightningClient.
+        """Initialize the client.

        Args:
-            endpoint: The root URL of the Agent Lightning server.
-            poll_interval: The interval in seconds to wait between polling for new tasks.
-            timeout: The timeout in seconds for HTTP requests.
+            endpoint: Root URL of the Agent Lightning server.
+            poll_interval: Seconds to wait between polling attempts.
+            timeout: Seconds before a request to the server is considered timed out.
        """
+        warnings.warn(
+            "AgentLightningClient is deprecated. Please use LightningStoreClient instead.", DeprecationWarning
+        )
        self.endpoint = endpoint
        self.task_count = 0
        self.poll_interval = poll_interval
@@ -44,13 +68,13 @@ class AgentLightningClient:
        self._default_headers = {"X-AgentLightning-Client": "true"}

    async def _request_json_async(self, url: str) -> Optional[Dict[str, Any]]:
-        """Makes an async GET request to the specified URL and returns the JSON response.
+        """Perform an asynchronous ``GET`` request and parse the JSON payload.

        Args:
-            url: The URL to request.
+            url: Fully qualified URL to query.

        Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
        """
        timeout = aiohttp.ClientTimeout(total=self.timeout)
        async with aiohttp.ClientSession(timeout=timeout) as session:
@@ -63,14 +87,14 @@ class AgentLightningClient:
                return None

    async def _post_json_async(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        """Makes an async POST request with a JSON payload.
+        """Perform an asynchronous ``POST`` request with a JSON body.

        Args:
-            url: The URL to post to.
-            payload: The dictionary data to send as JSON.
+            url: Fully qualified URL that accepts the payload.
+            payload: Dictionary that will be serialized and sent as JSON.

        Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
        """
        timeout = aiohttp.ClientTimeout(total=self.timeout)
        async with aiohttp.ClientSession(timeout=timeout) as session:
@@ -82,11 +106,12 @@ class AgentLightningClient:
                logger.debug(f"Async POST request failed for {url}: {e}")
                return None

-    async def poll_next_task_async(self) -> Task:
-        """Polls the server asynchronously for the next task until one is available.
+    async def poll_next_task_async(self) -> Optional[Task]:
+        """Poll the server asynchronously until a task becomes available.

        Returns:
-            A Task object containing the task details.
+            The next [`Task`][agentlightning.Task] exposed by the server,
+            or ``None`` if polling fails.
        """
        url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
        while True:
@@ -101,13 +126,15 @@ class AgentLightningClient:
            await asyncio.sleep(self.poll_interval)

    async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
-        """Fetches a specific version of resources by its ID, using a cache.
+        """Fetch a specific resource bundle by identifier.

        Args:
-            resource_id: The ID of the resources to fetch, usually from a Task's metadata.
+            resource_id: Identifier sourced from the task metadata.

        Returns:
-            A ResourcesUpdate object containing the versioned resources, or None if not found.
+            Cached or freshly downloaded
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
+            ``None`` when the server returns an error.
        """
        if resource_id in self._resource_cache:
            logger.debug(f"Found resources '{resource_id}' in cache.")
@@ -123,10 +150,11 @@ class AgentLightningClient:
        return None

    async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
-        """Fetches the latest available resources from the server.
+        """Fetch the most recent resource bundle advertised by the server.

        Returns:
-            A ResourcesUpdate object containing the latest resources.
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
+            newest version, or ``None`` when unavailable.
        """
        url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
        response = await self._request_json_async(url)
@@ -137,27 +165,27 @@ class AgentLightningClient:
            return resources_update
        return None

-    async def post_rollout_async(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
-        """Posts a completed rollout to the server asynchronously.
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
+        """Submit a completed rollout back to the server.

        Args:
-            rollout: A Rollout object containing the results of a task.
+            rollout: Legacy rollout payload produced by the executor.

        Returns:
-            The server's JSON response as a dictionary.
+            Parsed JSON response returned by the server, or ``None`` when the request fails.
        """
        url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
        payload = rollout.model_dump(mode="json")
        return await self._post_json_async(url, payload)

    def _request_json(self, url: str) -> Optional[Dict[str, Any]]:
-        """Makes a sync GET request to the specified URL and returns the JSON response.
+        """Perform a blocking ``GET`` request and parse the JSON payload.

        Args:
-            url: The URL to request.
+            url: Fully qualified URL to query.

        Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
        """
        try:
            response = requests.get(url, timeout=self.timeout, headers=self._default_headers)
@@ -168,14 +196,14 @@ class AgentLightningClient:
            return None

    def _post_json(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        """Makes a sync POST request with a JSON payload.
+        """Perform a blocking ``POST`` request with a JSON payload.

        Args:
-            url: The URL to post to.
-            payload: The dictionary data to send as JSON.
+            url: Fully qualified URL that accepts the payload.
+            payload: Dictionary that will be serialized and sent as JSON.

        Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
        """
        try:
            response = requests.post(url, json=payload, timeout=self.timeout, headers=self._default_headers)
@@ -185,11 +213,12 @@ class AgentLightningClient:
            logger.debug(f"Sync POST request failed for {url}: {e}")
            return None

-    def poll_next_task(self) -> Task:
-        """Polls the server synchronously for the next task until one is available.
+    def poll_next_task(self) -> Optional[Task]:
+        """Poll the server synchronously until a task becomes available.

        Returns:
-            A Task object containing the task details, including the required `resources_id`.
+            The next [`Task`][agentlightning.Task] available for execution, or
+            ``None`` if polling fails.
        """
        url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
        while True:
@@ -204,13 +233,15 @@ class AgentLightningClient:
            time.sleep(self.poll_interval)

    def get_resources_by_id(self, resource_id: str) -> Optional[ResourcesUpdate]:
-        """Fetches a specific version of resources by its ID synchronously, using a cache.
+        """Fetch a specific resource bundle by identifier.

        Args:
-            resource_id: The ID of the resources to fetch, usually from a Task's metadata.
+            resource_id: Identifier sourced from the task metadata.

        Returns:
-            A ResourcesUpdate object containing the versioned resources, or None if not found.
+            Cached or freshly downloaded
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
+            ``None`` when the server returns an error.
        """
        if resource_id in self._resource_cache:
            logger.debug(f"Found resources '{resource_id}' in cache.")
@@ -226,10 +257,11 @@ class AgentLightningClient:
        return None

    def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        """Fetches the latest available resources from the server synchronously.
+        """Fetch the most recent resource bundle advertised by the server.

        Returns:
-            A ResourcesUpdate object containing the latest resources.
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
+            newest version, or ``None`` when unavailable.
        """
        url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
        response = self._request_json(url)
@@ -239,14 +271,14 @@ class AgentLightningClient:
            return resources_update
        return None

-    def post_rollout(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
-        """Posts a completed rollout to the server synchronously.
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
+        """Submit a completed rollout back to the server.

        Args:
-            rollout: A Rollout object containing the results of a task.
+            rollout: Legacy rollout payload produced by the executor.

        Returns:
-            The server's JSON response as a dictionary.
+            Parsed JSON response returned by the server, or ``None`` when the request fails.
        """
        url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
        payload = rollout.model_dump(mode="json")
@@ -254,14 +286,16 @@ class AgentLightningClient:


 class DevTaskLoader(AgentLightningClient):
-    """A local task manager for development that provides sample tasks and resources.
+    """In-memory task loader used for development and integration tests.

-    This client mocks the server APIs by maintaining a local queue of tasks and resources
-    within the same process. It's designed for development, testing, and scenarios where
-    a full Agent Lightning server is not needed.
+    The loader mimics the behavior of the legacy HTTP server by storing tasks and
+    resources locally. Polling methods simply iterate over the provided collection,
+    allowing rapid iteration without provisioning any external infrastructure.

-    The DevTaskLoader overrides the polling and resource fetching methods to return data
-    from local collections instead of making HTTP requests to a remote server.
+    !!! warning "Deprecated"
+
+        [`DevTaskLoader`][agentlightning.client.DevTaskLoader] is a compatibility shim.
+        Prefer [`Trainer.dev`][agentlightning.Trainer.dev] for new code.
    """

    def __init__(
@@ -270,13 +304,19 @@ class DevTaskLoader(AgentLightningClient):
        resources: Union[NamedResources, ResourcesUpdate],
        **kwargs: Any,
    ):
-        """Initializes the DevTaskLoader with pre-defined tasks and resources.
+        """Initialize the loader with predefined tasks and resources.

        Args:
-            tasks: Either a List of TaskInput objects or a List of Task objects.
-            resources: Either NamedResources or ResourcesUpdate object.
-            **kwargs: Additional arguments passed to the parent AgentLightningClient.
+            tasks: Sequence of task inputs or preconstructed tasks that will be served in
+                order.
+            resources: Static resources returned for any `resources_id` query.
+            **kwargs: Additional keyword arguments forwarded to the parent client.
+
+        Raises:
+            ValueError: If no tasks are provided or both [`Task`][agentlightning.Task]
+                and [`TaskInput`][agentlightning.TaskInput] instances are mixed.
        """
+        warnings.warn("DevTaskLoader is deprecated. Please use Trainer.dev instead.", DeprecationWarning)
        super().__init__(endpoint="local://", **kwargs)
        self._tasks = tasks.copy()
        if len(self._tasks) == 0:
@@ -292,24 +332,27 @@ class DevTaskLoader(AgentLightningClient):
        if isinstance(resources, ResourcesUpdate):
            self._resources_update = resources
        else:
-            self._resources_update = ResourcesUpdate(resources_id="local", resources=resources)
+            self._resources_update = ResourcesUpdate(
+                resources_id="local", resources=resources, create_time=time.time(), update_time=time.time(), version=1
+            )

        # Store rollouts posted back to the loader for easy debugging of local runs
-        self._rollouts: List[Rollout] = []
+        self._rollouts: List[RolloutLegacy] = []

    @property
-    def rollouts(self) -> List[Rollout]:
-        """Return rollouts that have been posted back to the loader."""
+    def rollouts(self) -> List[RolloutLegacy]:
+        """Return the rollouts posted back to the loader during development runs."""
        return self._rollouts

-    def poll_next_task(self) -> Task:
-        """Returns the next task from the local queue.
+    def poll_next_task(self) -> Optional[Task]:
+        """Return the next task from the local queue.

-        If tasks are TaskInput objects, assembles them into Task objects.
-        If tasks are already Task objects, returns them directly.
+        If [`TaskInput`][agentlightning.TaskInput] instances were provided,
+        they are converted into [`Task`][agentlightning.Task] objects on the
+        fly. Otherwise, the preconstructed tasks are returned in sequence.

        Returns:
-            The next Task object from the local task list.
+            Next task to execute.
        """
        if self._task_index >= len(self._tasks):
            self._task_index = 0
@@ -344,12 +387,12 @@ class DevTaskLoader(AgentLightningClient):
        logger.debug("DevTaskLoader returning latest resources.")
        return self._resources_update

-    def post_rollout(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
        logger.debug(f"DevTaskLoader received rollout for task: {rollout.rollout_id}")
        self._rollouts.append(rollout)
        return {"status": "received", "rollout_id": rollout.rollout_id}

-    async def poll_next_task_async(self) -> Task:
+    async def poll_next_task_async(self) -> Optional[Task]:
        return self.poll_next_task()

    async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
@@ -358,7 +401,7 @@ class DevTaskLoader(AgentLightningClient):
    async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
        return self.get_latest_resources()

-    async def post_rollout_async(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
        return self.post_rollout(rollout)

    def __repr__(self):
@@ -1,3 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 """
 This file is not carefully reviewed.
 It might contain unintentional bugs and issues.
@@ -9,26 +11,28 @@ from __future__ import annotations
 import argparse
 import inspect
 import logging
+from typing import _GenericAlias  # type: ignore
 from typing import (
    Any,
+    Callable,
+    Dict,
    List,
+    Tuple,
    Type,
    TypeVar,
    Union,
-    _GenericAlias,  # type: ignore
-    get_origin,
    get_args,
-    Tuple,
-    Callable,
-    overload,
-    Dict,
+    get_origin,
    get_type_hints,
+    overload,
 )

 CliConfigurable = Any

 logger = logging.getLogger(__name__)

+__all__ = ["lightning_cli"]
+
 # TypeVars for precise return type hinting with overloads
 _C = TypeVar("_C", bound=CliConfigurable)
 _C1 = TypeVar("_C1", bound=CliConfigurable)
@@ -67,8 +71,8 @@ def nullable_float(value: str) -> float | None:

 def _str_to_bool(v: str) -> bool:
    """Converts common string representations of bool to Python bool (case-insensitive)."""
-    if isinstance(v, bool):  # Allow passing bools directly if used programmatically
-        return v
+    if isinstance(v, bool):  # type: ignore
+        return v  # Allow passing bools directly if used programmatically
    lowered_v = v.lower()
    if lowered_v in ("yes", "true", "t", "y", "1"):
        return True
@@ -79,12 +83,17 @@ def _str_to_bool(v: str) -> bool:


 def _get_param_type_details(param_annotation: Any) -> Tuple[Any, bool, bool]:
-    """
-    Determines the core type, if it's Optional, and if it's a List.
-    Returns: (core_type, is_optional, is_list)
-    - For Optional[T]: (T, True, is_list_status_of_T)
-    - For List[T]: (List[T], is_optional_status_of_List, True)
-    - For Optional[List[T]]: (List[T], True, True)
+    """Normalize an annotation into its core type, optionality, and list status.
+
+    Args:
+        param_annotation: The annotation to inspect.
+
+    Returns:
+        A tuple ``(core_type, is_optional, is_list)`` describing the normalized type.
+
+        - For ``Optional[T]`` → ``(T, True, is_list_status_of_T)``
+        - For ``List[T]`` → ``(List[T], is_optional_status_of_List, True)``
+        - For ``Optional[List[T]]`` → ``(List[T], True, True)``
    """
    is_optional = False
    is_list = False
@@ -305,7 +314,10 @@ def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3], cls4: Type[
 def lightning_cli(*classes: Type[CliConfigurable]) -> Tuple[CliConfigurable, ...]: ...


-def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:
+# FIXME: lightning_cli needs to be fixed to comply with the latest trainer implementation.
+
+
+def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:  # type: ignore
    """
    Parses command-line arguments to configure and instantiate provided CliConfigurable classes.

@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .exception import emit_exception
+from .message import emit_message
+from .object import emit_object
+from .reward import (
+    emit_reward,
+    find_final_reward,
+    find_reward_spans,
+    get_reward_value,
+    is_reward_span,
+    reward,
+)
+
+__all__ = [
+    "reward",
+    "emit_reward",
+    "get_reward_value",
+    "is_reward_span",
+    "find_reward_spans",
+    "find_final_reward",
+    "emit_message",
+    "emit_object",
+    "emit_exception",
+]
@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import traceback
+
+from opentelemetry.semconv.attributes import exception_attributes
+
+from agentlightning.types import SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def emit_exception(exception: BaseException) -> None:
+    """Record an exception with OpenTelemetry metadata.
+
+    Args:
+        exception: Raised exception instance to serialize into telemetry attributes.
+
+    !!! note
+        The helper validates its input. Non-exception values are ignored to prevent
+        noisy telemetry and indicate programming mistakes via the logger.
+    """
+    if not isinstance(exception, BaseException):  # type: ignore
+        logger.error(f"Expected an BaseException instance, got: {type(exception)}. Skip emit_exception.")
+        return
+
+    tracer = get_tracer()
+    stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
+    attributes = {
+        exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
+        exception_attributes.EXCEPTION_MESSAGE: str(exception),
+        exception_attributes.EXCEPTION_ESCAPED: True,
+    }
+    if stacktrace.strip():
+        attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+
+    span = tracer.start_span(
+        SpanNames.EXCEPTION.value,
+        attributes=attributes,
+    )
+    logger.debug("Emitting exception span for %s", type(exception).__name__)
+    with span:
+        span.record_exception(exception)
+        # We don't set the status of the span here. They have other semantics.
@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+
+from agentlightning.types import SpanAttributeNames, SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def emit_message(message: str) -> None:
+    """Emit a textual message as an OpenTelemetry span.
+
+    Args:
+        message: Human readable message to attach as a span attribute.
+
+    !!! note
+        OpenTelemetry distinguishes between logs and spans. Emitting the message as a
+        span keeps all Agent Lightning telemetry in a single data store for analysis.
+    """
+    if not isinstance(message, str):  # type: ignore
+        logger.error(f"Message must be a string, got: {type(message)}. Skip emit_message.")
+        return
+
+    tracer = get_tracer()
+    span = tracer.start_span(
+        SpanNames.MESSAGE.value,
+        attributes={SpanAttributeNames.MESSAGE.value: message},
+    )
+    logger.debug("Emitting message span with message: %s", message)
+    with span:
+        pass
@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+from typing import Any
+
+from agentlightning.types import SpanAttributeNames, SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def emit_object(object: Any) -> None:
+    """Emit an object's serialized representation as an OpenTelemetry span.
+
+    Args:
+        object: Data structure to encode as JSON and attach to the span payload.
+
+    !!! note
+        The payload must be JSON serializable. Non-serializable objects are ignored and
+        an error is logged to aid debugging.
+    """
+    try:
+        serialized = json.dumps(object)
+    except (TypeError, ValueError):
+        logger.error(f"Object must be JSON serializable, got: {type(object)}. Skip emit_object.")
+        return
+
+    tracer = get_tracer()
+    span = tracer.start_span(
+        SpanNames.OBJECT.value,
+        attributes={SpanAttributeNames.OBJECT.value: serialized},
+    )
+    logger.debug("Emitting object span with payload size %d characters", len(serialized))
+    with span:
+        pass
@@ -0,0 +1,238 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Helpers for emitting reward spans and integrating with AgentOps telemetry."""
+
+import asyncio
+import inspect
+import json
+import logging
+import warnings
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    TypedDict,
+    TypeVar,
+    cast,
+)
+
+import agentops
+from agentops.sdk.decorators import operation
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import SpanLike, SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "reward",
+    "emit_reward",
+    "get_reward_value",
+    "is_reward_span",
+    "find_reward_spans",
+    "find_final_reward",
+]
+
+
+class RewardSpanData(TypedDict):
+    type: Literal["reward"]
+    value: Optional[float]
+
+
+FnType = TypeVar("FnType", bound=Callable[..., Any])
+
+
+def _agentops_initialized() -> bool:
+    """Return `True` when the AgentOps client has been configured."""
+    return agentops.get_client().initialized
+
+
+def reward(fn: FnType) -> FnType:
+    """Decorate a reward function so its outputs are tracked as spans.
+
+    The decorator integrates with AgentOps when it is available and falls back to
+    the built-in telemetry otherwise. Both synchronous and asynchronous functions
+    are supported transparently.
+
+    Deprecated:
+        This decorator is deprecated. Use [`emit_reward`][agentlightning.emit_reward] instead.
+
+    Args:
+        fn: Callable that produces a numeric reward.
+
+    Returns:
+        Wrapped callable that preserves the original signature.
+    """
+
+    def wrap_result(result: Optional[float]) -> RewardSpanData:
+        """Normalize the reward value into the span payload format."""
+        if result is None:
+            return {"type": "reward", "value": None}
+        if not isinstance(result, (float, int)):  # type: ignore
+            warnings.warn(f"Reward is ignored because it is not a number: {result}")
+            return {"type": "reward", "value": None}
+        return {"type": "reward", "value": float(result)}
+
+    # Check if the function is async
+    is_async = asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn)
+
+    if is_async:
+
+        async def wrapper_async(*args: Any, **kwargs: Any) -> Any:
+            if not _agentops_initialized():
+                # Track the reward without AgentOps
+                result = await fn(*args, **kwargs)
+                emit_reward(cast(float, result))
+                return result
+
+            result: Optional[float] = None
+
+            @operation
+            async def agentops_reward_operation() -> RewardSpanData:
+                # The reward function we are interested in tracing
+                # It takes zero inputs and return a formatted dict
+                nonlocal result
+                result = await fn(*args, **kwargs)
+                return wrap_result(result)
+
+            await agentops_reward_operation()
+            return result
+
+        return wrapper_async  # type: ignore
+
+    else:
+
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if not _agentops_initialized():
+                # Track the reward without AgentOps
+                result = fn(*args, **kwargs)
+                emit_reward(cast(float, result))
+                return result
+
+            result: Optional[float] = None
+
+            @operation
+            def agentops_reward_operation() -> RewardSpanData:
+                nonlocal result
+                result = fn(*args, **kwargs)
+                return wrap_result(result)
+
+            agentops_reward_operation()
+            return result
+
+        return wrapper  # type: ignore
+
+
+def emit_reward(reward: float) -> 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.
+
+    Returns:
+        Readable span capturing the recorded reward.
+
+    Raises:
+        ValueError: If the provided reward cannot be interpreted as a float or the
+            resulting span is not a [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) instance.
+    """
+    logger.debug(f"Emitting reward: {reward}")
+    if isinstance(reward, (int, bool)):
+        reward = float(reward)
+    if not isinstance(reward, float):
+        raise ValueError(f"Reward must be a number, got: {type(reward)}")
+
+    # TODO: This should use the tracer from current context by tracer
+    tracer = get_tracer()
+    span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
+    # Do nothing; it's just a number
+    with span:
+        pass
+    if not isinstance(span, ReadableSpan):
+        raise ValueError(f"Span is not a ReadableSpan: {span}")
+    return span
+
+
+def get_reward_value(span: SpanLike) -> Optional[float]:
+    """Extract the reward value from a span, if available.
+
+    Args:
+        span: Span object produced by AgentOps or Agent Lightning emitters.
+
+    Returns:
+        The reward encoded in the span or `None` when the span does not represent a reward.
+    """
+    for key in [
+        "agentops.task.output",  # newer versions of agentops
+        "agentops.entity.output",
+    ]:
+        reward_dict: Dict[str, Any] | None = None
+        if span.attributes:
+            output = span.attributes.get(key)
+            if output:
+                if isinstance(output, dict):
+                    reward_dict = cast(Dict[str, Any], output)
+                elif isinstance(output, str):
+                    try:
+                        reward_dict = cast(Dict[str, Any], json.loads(output))
+                    except json.JSONDecodeError:
+                        reward_dict = None
+
+        if reward_dict and reward_dict.get("type") == "reward":
+            reward_value = reward_dict.get("value", None)
+            if reward_value is None:
+                return None
+            if not isinstance(reward_value, float):
+                logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+            return cast(float, reward_value)
+
+    # Latest emit reward format
+    if span.name == SpanNames.REWARD.value and span.attributes:
+        reward_value = span.attributes.get("reward", None)
+        if reward_value is None:
+            return None
+        if not isinstance(reward_value, float):
+            logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+        return cast(float, reward_value)
+    return None
+
+
+def is_reward_span(span: SpanLike) -> bool:
+    """Return ``True`` when the provided span encodes a reward value."""
+    maybe_reward = get_reward_value(span)
+    return maybe_reward is not None
+
+
+def find_reward_spans(spans: Sequence[SpanLike]) -> List[SpanLike]:
+    """Return all reward spans in the provided sequence.
+
+    Args:
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
+
+    Returns:
+        List of spans that could be parsed as rewards.
+    """
+    return [span for span in spans if is_reward_span(span)]
+
+
+def find_final_reward(spans: Sequence[SpanLike]) -> Optional[float]:
+    """Return the last reward value present in the provided spans.
+
+    Args:
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
+
+    Returns:
+        Reward value from the latest reward span, or `None` when none are found.
+    """
+    for span in reversed(spans):
+        reward = get_reward_value(span)
+        if reward is not None:
+            return reward
+    return None
@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utilities shared across emitter implementations."""
+
+import opentelemetry.trace as trace_api
+from opentelemetry.trace import get_tracer_provider
+
+
+def get_tracer() -> trace_api.Tracer:
+    """Resolve the OpenTelemetry tracer configured for Agent Lightning.
+
+    Returns:
+        OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.
+
+    Raises:
+        RuntimeError: If OpenTelemetry was not initialized before calling this helper.
+    """
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+
+    tracer_provider = get_tracer_provider()
+    return tracer_provider.get_tracer("agentlightning")
@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import ExecutionStrategy
+from .client_server import ClientServerExecutionStrategy
+from .events import ExecutionEvent, MultiprocessingEvent, ThreadingEvent
+from .shared_memory import SharedMemoryExecutionStrategy
+
+__all__ = [
+    "ExecutionStrategy",
+    "ClientServerExecutionStrategy",
+    "ExecutionEvent",
+    "ThreadingEvent",
+    "MultiprocessingEvent",
+    "SharedMemoryExecutionStrategy",
+]
@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Protocol
+
+from agentlightning.store.base import LightningStore
+
+from .events import ExecutionEvent
+
+logger = logging.getLogger(__name__)
+
+
+_TRUTHY_VALUES = {"1", "true", "yes", "on"}
+_FALSY_VALUES = {"0", "false", "no", "off"}
+
+
+def resolve_managed_store_flag(value: bool | None) -> bool:
+    """Determine whether execution helpers should wrap the provided store.
+
+    The helper first honours an explicit `value`. When `None` it falls back
+    to the `AGL_MANAGED_STORE` environment variable, accepting a variety
+    of truthy and falsy spellings. Missing environment configuration defaults to
+    `True` so that higher-level strategies create the appropriate client or
+    server wrappers automatically.
+
+    Args:
+        value: Optional override supplied by the caller.
+
+    Returns:
+        `True` when a managed store should be created around the provided
+        instance, otherwise `False`.
+
+    Raises:
+        ValueError: If `AGL_MANAGED_STORE` is set to an unsupported
+            value.
+    """
+
+    if value is not None:
+        return value
+
+    env_value = os.getenv("AGL_MANAGED_STORE")
+    if env_value is None:
+        return True
+
+    normalized = env_value.strip().lower()
+    if normalized in _TRUTHY_VALUES:
+        return True
+    if normalized in _FALSY_VALUES:
+        return False
+
+    raise ValueError("AGL_MANAGED_STORE must be one of 1, 0, true, false, yes, no, on, or off")
+
+
+class AlgorithmBundle(Protocol):
+    """Callable bundle produced by [`Trainer`][agentlightning.Trainer].
+
+    Execution strategies treat the returned coroutine as opaque, only providing
+    the shared store instance and cooperative stop event. Bundles typically
+    encapsulate algorithm setup plus adapter and LLM proxy, etc.
+    """
+
+    async def __call__(self, store: LightningStore, event: ExecutionEvent) -> None:
+        """Execute algorithm logic using ``store`` until completion or stop."""
+
+
+class RunnerBundle(Protocol):
+    """Callable bundle wrapping runner setup and the worker loop, as opposed to the
+    [`AlgorithmBundle`][agentlightning.AlgorithmBundle]."""
+
+    async def __call__(self, store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
+        """Execute runner logic for ``worker_id`` using ``store`` and ``event``."""
+
+
+class ExecutionStrategy:
+    """Coordinate algorithm and runner bundles within a single process abstraction.
+
+    Strategies decide how many worker bundles to launch, whether to communicate
+    through shared memory or an HTTP boundary, and how to react to shutdown
+    signals. They intentionally avoid inspecting the bundle internals; instead,
+    each bundle remains responsible for its own scheduling semantics.
+
+    !!! note
+        Implementations must honor the [execute()][agentlightning.ExecutionStrategy.execute]
+        contract by propagating `KeyboardInterrupt` and ensuring resources are
+        released when an error occurs on either side of the algorithm/runner
+        pair.
+    """
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        """Run the provided bundles using the configured orchestration model.
+
+        Args:
+            algorithm: Callable bundle responsible for algorithm execution.
+            runner: Callable bundle for runner workers.
+            store: Concrete [`LightningStore`][agentlightning.LightningStore]
+                shared across bundles.
+
+        Raises:
+            NotImplementedError: Subclasses must provide the orchestration
+                implementation.
+        """
+
+        raise NotImplementedError()
@@ -0,0 +1,433 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import multiprocessing
+import os
+import signal
+import time
+from multiprocessing.context import BaseContext
+from typing import Callable, Iterable, Literal, cast
+
+from agentlightning.store.base import LightningStore
+from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
+
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
+from .events import ExecutionEvent, MultiprocessingEvent
+
+logger = logging.getLogger(__name__)
+
+
+class ClientServerExecutionStrategy(ExecutionStrategy):
+    """Run algorithm and runner bundles as separate processes over HTTP.
+
+    Execution Roles:
+
+    - `"algorithm"`: Start [`LightningStoreServer`][agentlightning.LightningStoreServer]
+      in-process and execute the algorithm bundle against it.
+    - `"runner"`: Connect to an existing server with
+      [`LightningStoreClient`][agentlightning.LightningStoreClient] and run the
+      runner bundle locally (spawning multiple processes when requested).
+    - `"both"`: Spawn runner processes first, then execute the algorithm and
+      server on the same machine. This mode orchestrates the full loop locally.
+
+    When `role == "both"` you may choose which side runs on the main process
+    via `main_process`. The runner-on-main option is limited to
+    `n_runners == 1` because each additional runner requires its own event
+    loop and process.
+
+    !!! warning
+        When `main_process == "runner"` the algorithm and HTTP server execute
+        in a child process. Store mutations remain isolated inside that process,
+        so the original store instance passed to
+        [execute()][agentlightning.ExecutionStrategy.execute] is not updated.
+
+    Abort Model (four-step escalation):
+
+    1. Cooperative stop. Every bundle receives a shared
+       [`MultiprocessingEvent`][agentlightning.MultiprocessingEvent] (`stop_evt`).
+       Any failure flips the event so peers can exit cleanly. Ctrl+C on the main
+       process also sets the flag.
+    2. KeyboardInterrupt synthesis. Remaining subprocesses receive ``SIGINT`` to
+       trigger `KeyboardInterrupt` handlers.
+    3. Termination. Stubborn processes are asked to ``terminate()``
+       (`SIGTERM` on POSIX).
+    4. Kill. As a last resort `kill()` is invoked (`SIGKILL` on POSIX).
+
+    This mirrors the semantics implemented in
+    [`SharedMemoryExecutionStrategy`][agentlightning.SharedMemoryExecutionStrategy]
+    but adapts them to multiple processes and the HTTP client/server boundary.
+    """
+
+    alias: str = "cs"
+
+    def __init__(
+        self,
+        role: Literal["algorithm", "runner", "both"] | None = None,
+        server_host: str | None = None,
+        server_port: int | None = None,
+        n_runners: int = 1,
+        graceful_timeout: float = 5.0,
+        terminate_timeout: float = 5.0,
+        main_process: Literal["algorithm", "runner"] = "algorithm",
+        managed_store: bool | None = None,
+    ) -> None:
+        """Configure the strategy.
+
+        Args:
+            role: Which side(s) to run in this process. When omitted, the
+                `AGL_CURRENT_ROLE` environment variable is used.
+            server_host: Interface the HTTP server binds to when running the
+                algorithm bundle locally. Defaults to `AGL_SERVER_HOST`
+                or `"localhost"` if unset.
+            server_port: Port for the HTTP server in "algorithm"/"both" modes.
+                Defaults to `AGL_SERVER_PORT` or `4747` if unset.
+            n_runners: Number of runner processes to spawn in "runner"/"both".
+            graceful_timeout: How long to wait (seconds) after setting the stop
+                event before escalating to signals.
+            terminate_timeout: How long to wait between escalation steps beyond
+                the cooperative phase (re-used for SIGINT, terminate, and kill).
+            main_process: Which bundle runs on the main process when
+                `role == "both"`. `"runner"` requires `n_runners == 1` and is
+                primarily intended for debugging.
+            managed_store: When `True` (default) the strategy constructs
+                LightningStore client/server wrappers automatically. When
+                `False` the provided `store` is passed directly to the
+                bundles, allowing callers to manage store wrappers manually.
+        """
+        if role is None:
+            role_env = os.getenv("AGL_CURRENT_ROLE")
+            if role_env is None:
+                # Use both if not specified via env var or argument
+                role = "both"
+            elif role_env not in ("algorithm", "runner", "both"):
+                raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
+            else:
+                role = role_env
+
+        if server_host is None:
+            server_host = os.getenv("AGL_SERVER_HOST", "localhost")
+
+        if server_port is None:
+            server_port_env = os.getenv("AGL_SERVER_PORT")
+            if server_port_env is None:
+                server_port = 4747
+            else:
+                try:
+                    server_port = int(server_port_env)
+                except ValueError as exc:
+                    raise ValueError("AGL_SERVER_PORT must be an integer") from exc
+
+        self.role = role
+        self.n_runners = n_runners
+        self.server_host = server_host
+        self.server_port = server_port
+        self.graceful_timeout = graceful_timeout
+        self.terminate_timeout = terminate_timeout
+        if main_process not in ("algorithm", "runner"):
+            raise ValueError("main_process must be 'algorithm' or 'runner'")
+        if main_process == "runner":
+            if role != "both":
+                raise ValueError("main_process='runner' is only supported when role='both'")
+            if n_runners != 1:
+                raise ValueError("main_process='runner' requires n_runners to be 1")
+        self.main_process = main_process
+        self.managed_store = resolve_managed_store_flag(managed_store)
+
+    async def _execute_algorithm(
+        self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent
+    ) -> None:
+        wrapper_store: LightningStore | None = None
+        if self.managed_store:
+            logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
+            wrapper_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
+            server_started = False
+        else:
+            wrapper_store = store
+            server_started = False
+
+        try:
+            if self.managed_store and isinstance(wrapper_store, LightningStoreServer):
+                await wrapper_store.start()
+                server_started = True
+                logger.debug("Algorithm bundle starting against endpoint %s", wrapper_store.endpoint)
+            await algorithm(wrapper_store, stop_evt)
+            logger.debug("Algorithm bundle completed successfully")
+        except KeyboardInterrupt:
+            logger.warning("Algorithm received KeyboardInterrupt; signaling stop event")
+            stop_evt.set()
+            raise
+        except BaseException:
+            logger.exception("Algorithm bundle crashed; signaling stop event")
+            stop_evt.set()
+            raise
+        finally:
+            if self.managed_store and isinstance(wrapper_store, LightningStoreServer) and server_started:
+                try:
+                    await wrapper_store.stop()
+                except Exception:
+                    logger.exception("Error stopping LightningStore server")
+                else:
+                    logger.debug("LightningStore server shutdown completed")
+
+    async def _execute_runner(
+        self,
+        runner: RunnerBundle,
+        worker_id: int,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+    ) -> None:
+        if self.managed_store:
+            # If managed, we actually do not use the provided store
+            client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
+        else:
+            client_store = store
+        try:
+            if self.managed_store:
+                logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
+            else:
+                logger.debug("Runner %s executing with provided store", worker_id)
+            await runner(client_store, worker_id, stop_evt)
+            logger.debug("Runner %s completed successfully", worker_id)
+        except KeyboardInterrupt:
+            logger.warning("Runner %s received KeyboardInterrupt; signaling stop event", worker_id)
+            stop_evt.set()
+            raise
+        except BaseException:
+            logger.exception("Runner %s crashed; signaling stop event", worker_id)
+            stop_evt.set()
+            raise
+        finally:
+            if self.managed_store and isinstance(client_store, LightningStoreClient):
+                try:
+                    await client_store.close()
+                except Exception:
+                    logger.exception("Error closing LightningStore client for runner %s", worker_id)
+                else:
+                    logger.debug("Runner %s closed LightningStore client", worker_id)
+
+    def _spawn_runners(
+        self,
+        runner: RunnerBundle,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+        *,
+        ctx: BaseContext,
+    ) -> list[multiprocessing.Process]:
+        """Used when `role == "runner"` or `role == "both"` and `n_runners > 1`."""
+        processes: list[multiprocessing.Process] = []
+
+        def _runner_sync(runner: RunnerBundle, worker_id: int, store: LightningStore, stop_evt: ExecutionEvent) -> None:
+            # Runners are executed in child processes; each process owns its own
+            # event loop to keep the asyncio scheduler isolated.
+            asyncio.run(self._execute_runner(runner, worker_id, store, stop_evt))
+
+        for i in range(self.n_runners):
+            process = cast(
+                multiprocessing.Process,
+                ctx.Process(target=_runner_sync, args=(runner, i, store, stop_evt), name=f"runner-{i}"),  # type: ignore
+            )
+            process.start()
+            logger.debug("Spawned runner process %s (pid=%s)", process.name, process.pid)
+            processes.append(process)
+
+        return processes
+
+    def _spawn_algorithm_process(
+        self,
+        algorithm: AlgorithmBundle,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+        *,
+        ctx: BaseContext,
+    ) -> multiprocessing.Process:
+        """Used when `main_process == "runner"`."""
+
+        def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent) -> None:
+            asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+
+        process = cast(
+            multiprocessing.Process,
+            ctx.Process(target=_algorithm_sync, args=(algorithm, store, stop_evt), name="algorithm"),  # type: ignore
+        )
+        process.start()
+        logger.debug("Spawned algorithm process %s (pid=%s)", process.name, process.pid)
+        return process
+
+    def _join_until_deadline(
+        self,
+        processes: Iterable[multiprocessing.Process],
+        timeout: float,
+    ) -> list[multiprocessing.Process]:
+        """Join ``processes`` until ``timeout`` elapses, returning those still alive."""
+        deadline = time.monotonic() + timeout
+        still_alive: list[multiprocessing.Process] = []
+        for process in processes:
+            remaining = deadline - time.monotonic()
+            if remaining > 0:
+                process.join(remaining)
+            else:
+                process.join(0)
+            if process.is_alive():
+                still_alive.append(process)
+        return still_alive
+
+    def _signal_processes(
+        self,
+        processes: Iterable[multiprocessing.Process],
+        action: Callable[[multiprocessing.Process], None],
+    ) -> None:
+        """Invoke ``action`` on each process while suppressing individual failures."""
+        for process in processes:
+            try:
+                action(process)
+            except Exception:
+                logger.exception("Error signaling process %s (pid=%s)", process.name, process.pid)
+
+    def _shutdown_processes(
+        self,
+        processes: list[multiprocessing.Process],
+        stop_evt: ExecutionEvent,
+    ) -> None:
+        """4-step escalation shutdown of ``processes``."""
+        if not processes:
+            logger.debug("No subprocesses to shutdown")
+            return
+
+        if not stop_evt.is_set():
+            logger.debug("Sending cooperative stop signal to subprocesses")
+            stop_evt.set()
+        else:
+            logger.debug("Stop event already set; waiting for subprocesses to exit")
+
+        alive = self._join_until_deadline(processes, self.graceful_timeout)
+        if not alive:
+            return
+
+        logger.warning(
+            "Subprocesses still alive after cooperative wait; sending SIGINT to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        # SIGINT is not reliable on Windows, but we do not consider such case yet.
+        self._signal_processes(alive, lambda p: os.kill(cast(int, p.pid), signal.SIGINT))
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+        if not alive:
+            return
+
+        logger.warning(
+            "Subprocesses still alive after SIGINT wait; sending terminate() to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        self._signal_processes(alive, lambda p: p.terminate())
+
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+        if not alive:
+            return
+
+        logger.error(
+            "Subprocesses still alive after terminate(); sending kill() to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        self._signal_processes(alive, lambda p: p.kill())
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+
+        if alive:
+            logger.error(
+                "Subprocesses failed to exit even after kill(): %s", ", ".join(p.name or str(p.pid) for p in alive)
+            )
+
+    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)]
+        if failed:
+            formatted = ", ".join(f"{p.name or p.pid} (exitcode={p.exitcode})" for p in failed)
+            raise RuntimeError(f"Subprocesses failed: {formatted}")
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        logger.info(
+            "Starting client-server execution with %d runner(s) [role=%s, main_process=%s]",
+            self.n_runners,
+            self.role,
+            self.main_process,
+        )
+
+        # Re-use the active multiprocessing context so the event and processes
+        # agree on the start method (fork/spawn/forkserver).
+        ctx = multiprocessing.get_context()
+        stop_evt = MultiprocessingEvent(ctx=ctx)
+        # Track spawned processes so we can enforce termination ordering and
+        # surface non-zero exit codes back to the caller.
+        processes: list[multiprocessing.Process] = []
+
+        exception: BaseException | None = None
+        keyboard_interrupt = False
+
+        try:
+            if self.role == "algorithm":
+                logger.info("Running algorithm solely...")
+                asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+            elif self.role == "runner":
+                if self.n_runners == 1:
+                    logger.info("Running runner solely...")
+                    asyncio.run(self._execute_runner(runner, 0, store, stop_evt))
+                else:
+                    logger.info("Spawning runner processes...")
+                    processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
+                    # Wait for the processes to finish naturally.
+                    for process in processes:
+                        process.join()
+                    self._check_process_exitcodes(processes)
+            elif self.role == "both":
+                if self.main_process == "algorithm":
+                    logger.info("Spawning runner processes...")
+                    processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
+                    try:
+                        logger.info("Running algorithm...")
+                        asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+                    finally:
+                        # Always request the runner side to unwind once the
+                        # algorithm/server portion finishes (successfully or not).
+                        stop_evt.set()
+                else:  # main_process == "runner"
+                    if self.n_runners > 1:
+                        raise ValueError("main_process='runner' requires n_runners to be 1")
+
+                    logger.info("Spawning algorithm process...")
+                    algorithm_process = self._spawn_algorithm_process(algorithm, store, stop_evt, ctx=ctx)
+                    processes = [algorithm_process]
+
+                    # Run the lone runner cooperatively in-process so users can
+                    # attach a debugger. The algorithm + HTTP server live in
+                    # the background process spawned above (the provided
+                    # store must therefore be picklable when using spawn).
+                    logger.info("Running runner...")
+                    asyncio.run(self._execute_runner(runner, 0, store, stop_evt))
+
+                    # Wait for the algorithm process to finish.
+                    algorithm_process.join()
+            else:
+                raise ValueError(f"Unknown role: {self.role}")
+        except KeyboardInterrupt:
+            logger.warning("KeyboardInterrupt received; initiating shutdown")
+            stop_evt.set()
+            keyboard_interrupt = True
+        except BaseException as exc:
+            logger.exception("Unhandled exception in execute method")
+            stop_evt.set()
+            # Preserve the original exception so we can avoid masking it during
+            # the cleanup phase.
+            exception = exc
+            raise
+        finally:
+            logger.info("Shutting down subprocesses")
+            self._shutdown_processes(processes, stop_evt)
+            if processes:
+                try:
+                    self._check_process_exitcodes(processes)
+                except RuntimeError as err:
+                    if exception is not None or keyboard_interrupt:
+                        # We already propagate/handled a different failure, so
+                        # emit a warning instead of raising a secondary error.
+                        logger.warning("Subprocesses ended abnormally during shutdown: %s", err)
+                    else:
+                        raise
@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import multiprocessing as mp
+import threading
+from multiprocessing.context import BaseContext
+from typing import Optional, Protocol
+
+
+class ExecutionEvent(Protocol):
+    """Protocol capturing the cooperative stop contract shared by strategies.
+
+    Implementations mirror the API of ``threading.Event`` and
+    ``multiprocessing.Event`` so the rest of the execution layer can remain
+    agnostic to the underlying concurrency primitive.
+
+    Methods:
+
+        set: Signal cancellation. The call must be idempotent.
+        clear: Reset the event to the unsignaled state.
+        is_set: Return ``True`` when cancellation has been requested.
+        wait: Block until the event is signaled or an optional timeout elapses.
+    """
+
+    def set(self) -> None: ...
+    def clear(self) -> None: ...
+    def is_set(self) -> bool: ...
+    def wait(self, timeout: Optional[float] = None) -> bool: ...
+
+
+class ThreadingEvent:
+    """Thread-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""
+
+    __slots__ = ("_evt",)
+
+    def __init__(self) -> None:
+        self._evt = threading.Event()
+
+    def set(self) -> None:
+        self._evt.set()
+
+    def clear(self) -> None:
+        self._evt.clear()
+
+    def is_set(self) -> bool:
+        return self._evt.is_set()
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        return self._evt.wait(timeout)
+
+
+class MultiprocessingEvent:
+    """Process-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""
+
+    __slots__ = ("_evt",)
+
+    def __init__(self, *, ctx: Optional[BaseContext] = None) -> None:
+        self._evt = (ctx or mp).Event()
+
+    def set(self) -> None:
+        self._evt.set()
+
+    def clear(self) -> None:
+        self._evt.clear()
+
+    def is_set(self) -> bool:
+        return self._evt.is_set()
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        return self._evt.wait(timeout)
@@ -0,0 +1,16 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import ExecutionStrategy
+
+
+class InterProcessExecutionStrategy(ExecutionStrategy):
+    """Placeholder strategy for future inter-process primitives.
+
+    The class exists to reserve the `ipc` alias and make the planned
+    implementation discoverable. Attempting to use it today will raise
+    `NotImplementedError` once the execution contract is finalized.
+    """
+
+    alias: str = "ipc"
+
+    # TODO: to be implemented
@@ -0,0 +1,279 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import threading
+from contextlib import suppress
+from queue import SimpleQueue
+from typing import Any, Awaitable, Callable, List, Literal, Optional, Tuple
+
+from agentlightning.store.base import LightningStore
+from agentlightning.store.threading import LightningStoreThreaded
+
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
+from .events import ExecutionEvent, ThreadingEvent
+
+logger = logging.getLogger(__name__)
+
+
+class SharedMemoryExecutionStrategy(ExecutionStrategy):
+    """Execute bundles in a single process with cooperative worker threads.
+
+    Stop Model:
+
+    - All bundles share one [`ThreadingEvent`][agentlightning.ThreadingEvent]
+      named `stop_evt`.
+    - Only the main thread receives `KeyboardInterrupt`. When Ctrl+C occurs we
+      set `stop_evt`.
+    - Any exception raised inside a bundle sets `stop_evt` so other threads can
+      unwind cooperatively.
+    - Once the bundle running on the main thread exits successfully the
+      treatment depends on `main_thread`:
+        - `"algorithm"`: the runners are asked to stop by setting `stop_evt`.
+        - `"runner"`: the algorithm keeps running until it exits naturally.
+    - Background threads are marked as daemons. We join them briefly and log any
+      stragglers before shutting down.
+
+    !!! note
+        Signals other than `SIGINT` (such as `SIGTERM`) are not intercepted;
+        Python's default behavior for those signals is preserved.
+    """
+
+    alias: str = "shm"
+
+    def __init__(
+        self,
+        n_runners: int = 1,
+        main_thread: Literal["algorithm", "runner"] = "runner",
+        join_timeout: float = 15.0,
+        graceful_delay: float = 5.0,
+        poll_interval: float = 0.05,
+        managed_store: bool | None = None,
+    ) -> None:
+        if main_thread not in ("algorithm", "runner"):
+            raise ValueError("main_thread must be 'algorithm' or 'runner'")
+        if main_thread == "runner" and n_runners != 1:
+            raise ValueError(
+                "When main_thread is 'runner', n_runners must be 1. "
+                "Either use 'algorithm' on the main thread or set n_runners to 1."
+            )
+        self.n_runners = n_runners
+        self.main_thread = main_thread
+        self.join_timeout = join_timeout
+        self.graceful_delay = graceful_delay
+        self.poll_interval = poll_interval
+        self.managed_store = resolve_managed_store_flag(managed_store)
+
+    async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
+        """Run `coro` until it finishes or a cooperative stop is requested.
+
+        Control flow:
+
+        1. Start the bundle coroutine as `task`.
+        2. Launch a watcher that polls `stop_evt` without blocking the loop.
+        3. When the stop event flips:
+            a. Give the bundle `graceful_delay` seconds to finish on its own,
+               because well-behaved bundles will check the event and return.
+            b. Cancel the bundle task if it is still running after the grace
+               period.
+        4. Await both tasks and swallow `CancelledError` where appropriate.
+
+        This is a *backup* mechanism for bundles that might not poll the event
+        frequently; cooperative shutdown (checking `stop_evt` inside the
+        bundle) remains the preferred approach.
+        """
+        task: asyncio.Task[Any] = asyncio.create_task(coro)  # type: ignore
+        task_exception: Optional[BaseException] = None
+
+        async def watcher() -> None:
+            # Poll the threading event without blocking the event loop. Using a
+            # background thread via ``asyncio.to_thread`` makes cancellation
+            # difficult because ``ThreadingEvent.wait`` is not interruptible.
+            # Instead we cooperatively check the flag from the loop so the
+            # watcher task stays cancellable and tests don't hang when the
+            # bundle finishes naturally before the stop event is set.
+            while not stop_evt.is_set():
+                await asyncio.sleep(self.poll_interval)
+
+            # Grace period: let a cooperative bundle exit on its own.
+            try:
+                # At this point of waiting, the main task should already see the stop event.
+                await asyncio.wait_for(asyncio.shield(task), timeout=self.graceful_delay)  # type: ignore
+                logger.debug("Bundle finished by itself during grace period.")
+                return  # bundle finished by itself during grace period
+            except asyncio.TimeoutError:
+                # Still running after the grace window.
+                pass
+            except asyncio.CancelledError:
+                # If someone else canceled the task already, we're done.
+                logger.debug("Bundle already canceled by someone else; exiting watcher.")
+                return
+
+            # Still running after the grace window: cancel it.
+            if not task.done():
+                logger.debug("Graceful delay elapsed; canceling bundle task...")
+                task.cancel()
+
+        watcher_task = asyncio.create_task(watcher())
+        result: Any = None
+
+        try:
+            # We don't wait on FIRST_COMPLETED here, because we want the watcher
+            # to be able to grant a grace window after stop_evt flips.
+            await asyncio.wait(
+                {task, watcher_task}, return_when=asyncio.FIRST_COMPLETED
+            )  # pyright: ignore[reportUnknownArgumentType]
+        finally:
+            # If the main task hasn't completed yet (e.g., watcher scheduled cancel),
+            # finish the cancellation handshake.
+            if not task.done():
+                try:
+                    await asyncio.wait_for(task, timeout=self.graceful_delay)  # second chance
+                except asyncio.TimeoutError:
+                    logger.error(
+                        "Bundle task did not stop after cancellation; abandoning task."
+                        "This thread could live until the process exits."
+                    )
+                    # We return without awaiting it. asyncio.run will still try to cancel
+                    # pending tasks on loop close; if the task ignores cancellation, this
+                    # thread may still stick. It's the best we can do in Python.
+                    # We don't raise an exception here, but the thread could be a zombie.
+                    return result
+            else:
+                # Task completed naturally; retrieve result.
+                try:
+                    result = await task  # type: ignore
+                except asyncio.CancelledError:
+                    pass
+                except BaseException as exc:
+                    task_exception = exc
+
+            watcher_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await watcher_task
+
+        if task_exception is not None:
+            raise task_exception
+
+        return result  # type: ignore
+
+    def _run_algorithm(
+        self,
+        algorithm: AlgorithmBundle,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+        thread_exceptions: Optional[SimpleQueue[BaseException]],
+    ) -> None:
+        try:
+            asyncio.run(self._run_until_completed_or_canceled(algorithm(store, stop_evt), stop_evt))
+        except asyncio.CancelledError:
+            logger.info("Algorithm bundle canceled due to stop signal.")
+        except BaseException as exc:
+            logger.exception("Algorithm bundle crashed; signaling stop to others.")
+            if thread_exceptions is not None:
+                thread_exceptions.put(exc)
+            stop_evt.set()
+            raise
+
+    def _run_runner(
+        self,
+        runner: RunnerBundle,
+        store: LightningStore,
+        worker_id: int,
+        stop_evt: ExecutionEvent,
+        thread_exceptions: Optional[SimpleQueue[BaseException]],
+    ) -> None:
+        try:
+            asyncio.run(self._run_until_completed_or_canceled(runner(store, worker_id, stop_evt), stop_evt))
+        except asyncio.CancelledError:
+            logger.info("Runner bundle (worker_id=%s) canceled due to stop signal.", worker_id)
+        except BaseException as exc:
+            logger.exception("Runner bundle crashed (worker_id=%s); signaling stop to others.", worker_id)
+            if thread_exceptions is not None:
+                thread_exceptions.put(exc)
+            stop_evt.set()
+            raise
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        logger.info(
+            "Starting shm execution with %d runner(s); main thread runs '%s'",
+            self.n_runners,
+            self.main_thread,
+        )
+
+        # Create stop event and thread-safe store.
+        stop_evt = ThreadingEvent()
+        if self.managed_store:
+            thread_safe_store = LightningStoreThreaded(store)
+        else:
+            thread_safe_store = store
+
+        thread_exceptions: SimpleQueue[BaseException] = SimpleQueue()
+        raised_from_thread: Optional[BaseException] = None
+
+        def make_thread(name: str, target: Callable[..., Any], args: Tuple[Any, ...]) -> threading.Thread:
+            t = threading.Thread(name=name, target=target, args=args, daemon=True)
+            t.start()
+            return t
+
+        threads: List[threading.Thread] = []
+
+        try:
+            if self.main_thread == "algorithm":
+                # Start runner threads; algorithm runs on main thread.
+                for i in range(self.n_runners):
+                    thread = make_thread(
+                        name=f"runner-{i}",
+                        target=self._run_runner,
+                        args=(runner, thread_safe_store, i, stop_evt, thread_exceptions),
+                    )
+                    threads.append(thread)
+
+                # Ctrl+C here raises KeyboardInterrupt on this stack.
+                # Main thread doesn't need to collect exceptions.
+                self._run_algorithm(algorithm, thread_safe_store, stop_evt, None)
+
+                # If algo finishes naturally, request runners to stop.
+                stop_evt.set()
+
+            else:  # main_thread == "runner"
+                # Start algorithm in background; runner runs on main thread.
+                thread = make_thread(
+                    name="algorithm",
+                    target=self._run_algorithm,
+                    args=(algorithm, thread_safe_store, stop_evt, thread_exceptions),
+                )
+                threads.append(thread)
+
+                # Ctrl+C here raises KeyboardInterrupt on this stack.
+                # Main thread doesn't need to collect exceptions.
+                self._run_runner(runner, thread_safe_store, 0, stop_evt, None)
+
+                # If runner finishes naturally, WAIT FOR ALGORITHM TO FINISH.
+                thread.join()
+
+            if not thread_exceptions.empty():
+                raised_from_thread = thread_exceptions.get()
+
+        except KeyboardInterrupt:
+            logger.warning("KeyboardInterrupt received on main thread; initiating cooperative shutdown...")
+            stop_evt.set()
+        finally:
+            # Attempt a clean join; if some threads don't comply, log and move on.
+            for t in threads:
+                logger.debug("Joining thread %s...", t.name)
+                t.join(timeout=self.join_timeout)
+
+            alive = [t.name for t in threads if t.is_alive()]
+            if alive:
+                logger.error(
+                    "Threads still alive after %.1fs: %s. They are daemons; continuing shutdown.",
+                    self.join_timeout,
+                    ", ".join(alive),
+                )
+
+            if raised_from_thread is None and not thread_exceptions.empty():
+                raised_from_thread = thread_exceptions.get()
+
+        if raised_from_thread is not None:
+            raise raised_from_thread
@@ -1,21 +1,23 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 import warnings

-AGENTOPS_INSTALLED = False
-AGENTOPS_LANGCHAIN_INSTALLED = False
-LITELLM_INSTALLED = False
-VLLM_INSTALLED = False
+AGENTOPS_INSTALLED: bool = False
+AGENTOPS_LANGCHAIN_INSTALLED: bool = False
+LITELLM_INSTALLED: bool = False
+VLLM_INSTALLED: bool = False

 try:
-    from . import agentops
+    from . import agentops  # type: ignore

-    AGENTOPS_INSTALLED = True
+    AGENTOPS_INSTALLED = True  # type: ignore
 except ImportError:
    pass

 try:
-    from . import litellm
+    from . import litellm  # type: ignore

-    LITELLM_INSTALLED = True
+    LITELLM_INSTALLED = True  # type: ignore
 except ImportError:
    pass

@@ -30,14 +32,15 @@ except ImportError:


 try:
-    from . import agentops_langchain
+    from . import agentops_langchain  # type: ignore

-    AGENTOPS_LANGCHAIN_INSTALLED = True
+    AGENTOPS_LANGCHAIN_INSTALLED = True  # type: ignore
 except ImportError:
    pass


 def instrument_all():
+    """Instrument all the instrumentation libraries."""
    if AGENTOPS_INSTALLED:
        from .agentops import instrument_agentops

@@ -68,6 +71,7 @@ def instrument_all():


 def uninstrument_all():
+    """Uninstrument all the instrumentation libraries."""
    if AGENTOPS_INSTALLED:
        try:
            from .agentops import uninstrument_agentops
@@ -1,23 +1,83 @@
-import logging
-import multiprocessing
-import signal
-import socket
-import time
+# Copyright (c) Microsoft. All rights reserved.

-import flask
-import setproctitle
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, Callable, no_type_check
+
+import requests
+from agentops.client.api import V3Client, V4Client
+from agentops.client.api.types import AuthTokenResponse
+from agentops.sdk.exporters import AuthenticatedOTLPExporter
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.metrics.export import MetricExportResult
+from opentelemetry.sdk.trace.export import SpanExportResult

 logger = logging.getLogger(__name__)

+__all__ = [
+    "instrument_agentops",
+    "uninstrument_agentops",
+]
+
 # Module-level storage for originals
-_original_handle_chat_attributes = None
-_original_handle_response = None
+_original_handle_chat_attributes: Callable[..., Any] | None = None
+_original_handle_response: Callable[..., Any] | None = None
+_agentops_service_enabled = False
+
+
+def enable_agentops_service(enabled: bool = True) -> None:
+    """
+    Enable or disable communication with the AgentOps service.
+
+    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.
+    """
+    global _agentops_service_enabled
+    _agentops_service_enabled = enabled
+    logger.info(f"Switch set to {enabled} for exporters and clients.")
+
+
+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.client.api.V3Client = BypassableV3Client
+    agentops.client.api.V4Client = BypassableV4Client
+
+
+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.client.api.V3Client = V3Client
+    agentops.client.api.V4Client = V4Client
+
+
+def _unwrap_legacy_response(response: Any) -> Any:
+    if hasattr(response, "parse") and callable(response.parse):
+        return response.parse()
+    return response


 def _patch_new_agentops():
-    import agentops.instrumentation.providers.openai.wrappers.chat
    import agentops.instrumentation.providers.openai.stream_wrapper
-    from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes
+    import agentops.instrumentation.providers.openai.wrappers.chat
+    from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes  # type: ignore

    global _original_handle_chat_attributes

@@ -25,23 +85,60 @@ def _patch_new_agentops():
        logger.warning("AgentOps already patched. Skipping.")
        return True

-    _original_handle_chat_attributes = handle_chat_attributes
+    _original_handle_chat_attributes = handle_chat_attributes  # type: ignore

-    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):
+    @no_type_check
+    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):  # type: ignore
        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
-        if hasattr(return_value, "prompt_token_ids"):
+
+        # In some cases, response is a openai._legacy_response.LegacyAPIResponse (e.g., LiteLLM, or LangChain),
+        # This is created by client.with_raw_response.create()
+        return_value = _unwrap_legacy_response(return_value)
+
+        if (
+            return_value is not None
+            and hasattr(return_value, "prompt_token_ids")
+            and return_value.prompt_token_ids is not None
+        ):
            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
-        if hasattr(return_value, "response_token_ids"):
+        if (
+            return_value is not None
+            and hasattr(return_value, "response_token_ids")
+            and return_value.response_token_ids is not None
+        ):
            attributes["response_token_ids"] = list(return_value.response_token_ids[0])

-        # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
-        if hasattr(return_value, "http_response") and hasattr(return_value.http_response, "json"):
-            json_data = return_value.http_response.json()
-            if isinstance(json_data, dict):
-                if "prompt_token_ids" in json_data:
-                    attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])
-                if "response_token_ids" in json_data:
-                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])
+        # For LiteLLM Proxy (v0.2) with vLLM return_token_ids, response_token_ids now lives in choices
+        if (
+            return_value is not None
+            and hasattr(return_value, "choices")
+            and return_value.choices
+            and isinstance(return_value.choices, list)
+            and len(return_value.choices) > 0
+        ):
+            first_choice = return_value.choices[0]
+            # Token IDs from "choices[0].token_ids"
+            if "response_token_ids" not in attributes:
+                if hasattr(first_choice, "token_ids") and first_choice.token_ids is not None:
+                    attributes["response_token_ids"] = list(first_choice.token_ids)
+                # newer versions of OpenAI client SDK
+                elif (
+                    hasattr(first_choice, "provider_specific_fields")
+                    and first_choice.provider_specific_fields.get("token_ids") is not None
+                ):
+                    attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])
+
+            # log probability
+            # This is temporary. We need a unified convention for classifying and naming logprobs.
+            if hasattr(first_choice, "logprobs") and first_choice.logprobs is not None:
+                if hasattr(first_choice.logprobs, "content") and first_choice.logprobs.content is not None:
+                    attributes["logprobs.content"] = json.dumps(
+                        [logprob.model_dump() for logprob in first_choice.logprobs.content]
+                    )
+                if hasattr(first_choice.logprobs, "refusal") and first_choice.logprobs.refusal is not None:
+                    attributes["logprobs.refusal"] = json.dumps(
+                        [logprob.model_dump() for logprob in first_choice.logprobs.refusal]
+                    )

        return attributes

@@ -54,8 +151,8 @@ def _patch_new_agentops():


 def _unpatch_new_agentops():
-    import agentops.instrumentation.providers.openai.wrappers.chat
    import agentops.instrumentation.providers.openai.stream_wrapper
+    import agentops.instrumentation.providers.openai.wrappers.chat

    global _original_handle_chat_attributes
    if _original_handle_chat_attributes is not None:
@@ -70,40 +167,40 @@ def _unpatch_new_agentops():


 def _patch_old_agentops():
-    import opentelemetry.instrumentation.openai.shared.chat_wrappers
-    from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore
+    from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw  # type: ignore

    global _original_handle_response
-    _original_handle_response = _handle_response
+    _original_handle_response = _handle_response  # type: ignore

-    @dont_throw
-    def _handle_response_with_tokens(response, span, *args, **kwargs):
-        _original_handle_response(response, span, *args, **kwargs)
-        if hasattr(response, "prompt_token_ids"):
-            span.set_attribute("prompt_token_ids", list(response.prompt_token_ids))
-        if hasattr(response, "response_token_ids"):
-            span.set_attribute("response_token_ids", list(response.response_token_ids[0]))
+    @dont_throw  # type: ignore
+    def _handle_response_with_tokens(response, span, *args, **kwargs):  # type: ignore
+        _original_handle_response(response, span, *args, **kwargs)  # type: ignore
+        if hasattr(response, "prompt_token_ids"):  # type: ignore
+            span.set_attribute("prompt_token_ids", list(response.prompt_token_ids))  # type: ignore
+        if hasattr(response, "response_token_ids"):  # type: ignore
+            span.set_attribute("response_token_ids", list(response.response_token_ids[0]))  # type: ignore

        # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
-        if hasattr(response, "http_response") and hasattr(response.http_response, "json"):
-            json_data = response.http_response.json()
+        if hasattr(response, "http_response") and hasattr(response.http_response, "json"):  # type: ignore
+            json_data = response.http_response.json()  # type: ignore
            if isinstance(json_data, dict):
                if "prompt_token_ids" in json_data:
-                    span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"]))
+                    span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"]))  # type: ignore
                if "response_token_ids" in json_data:
-                    span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0]))
+                    span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0]))  # type: ignore

-    opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens
+    opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens  # type: ignore
    logger.info("Patched earlier version of agentops using _handle_response")
    return True


 def _unpatch_old_agentops():
-    import opentelemetry.instrumentation.openai.shared.chat_wrappers
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore

    global _original_handle_response
    if _original_handle_response is not None:
-        opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response
+        opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response  # type: ignore
        _original_handle_response = None
        logger.info("Unpatched earlier version of agentops using _handle_response")

@@ -113,6 +210,8 @@ def instrument_agentops():
    Instrument agentops to capture token IDs.
    Automatically detects and uses the appropriate patching method based on the installed agentops version.
    """
+    _patch_exporters()
+
    # Try newest version first (tested for 0.4.16)
    try:
        return _patch_new_agentops()
@@ -131,6 +230,9 @@ def instrument_agentops():


 def uninstrument_agentops():
+    """Uninstrument agentops to stop capturing token IDs."""
+    _unpatch_exporters()
+
    try:
        _unpatch_new_agentops()
    except Exception:
@@ -141,100 +243,75 @@ def uninstrument_agentops():
        pass


-def agentops_local_server():
+class BypassableAuthenticatedOTLPExporter(AuthenticatedOTLPExporter):
    """
-    Returns a Flask app that can be used to test agentops integration.
-    This server provides endpoints for token fetching and a catch-all endpoint.
+    AuthenticatedOTLPExporter with switchable service control.
+    When `_agentops_service_enabled` is False, skip export and return success.
    """
-    app = flask.Flask(__name__)

-    @app.route("/v3/auth/token", methods=["POST"])
-    def fetch_token():
-        return {"token": "dummy", "project_id": "dummy"}
-
-    @app.route("/", defaults={"path": ""}, methods=["GET", "POST"])
-    @app.route("/<path:path>", methods=["GET", "POST"])
-    def catch_all(path):
-        return {"path": path}
-
-    return app
-
-
-def _run_server(**kwargs):
-    """
-    Internal function to run the Flask server.
-    This is used to avoid issues with multiprocessing and Flask's reloader.
-    """
-    signal.signal(signal.SIGINT, signal.SIG_IGN)  # Ignore SIGINT in worker processes
-    setproctitle.setproctitle(multiprocessing.current_process().name)
-    app = agentops_local_server()
-    app.run(**kwargs)
-
-
-class AgentOpsServerManager:
-    def __init__(self, daemon: bool = True, port: int | None = None):
-        self.server_process: multiprocessing.Process | None = None
-        self.server_port = port
-        self.daemon = daemon
-        logger.info("AgentOpsServerManager initialized.")
-
-    def _find_available_port(self) -> int:
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.bind(("", 0))
-            return s.getsockname()[1]
-
-    def start(self):
-        if self.server_process and self.server_process.is_alive():
-            logger.warning("AgentOps server process appears to be already running.")
-            return
-
-        if self.server_port is None:
-            self.server_port = self._find_available_port()
-
-        logger.info(f"Starting AgentOps local server on port {self.server_port}...")
-
-        self.server_process = multiprocessing.Process(
-            target=_run_server,
-            kwargs={"host": "127.0.0.1", "port": self.server_port, "use_reloader": False, "debug": False},
-            daemon=self.daemon,
-            name="AgentLightning-AgentOpsServer",
-        )
-        self.server_process.start()
-        logger.info(
-            f"AgentOps local server process (PID: {self.server_process.pid}) started, targeting port {self.server_port}."
-        )
-        time.sleep(0.5)  # Brief wait for server to start up
-        if not self.server_process.is_alive():
-            logger.error(f"AgentOps local server failed to start or exited prematurely.")
-
-    def is_alive(self) -> bool:
-        if self.server_process and self.server_process.is_alive():
-            return True
-        return False
-
-    def stop(self):
-        if self.is_alive():
-            logger.info(f"Stopping AgentOps local server (PID: {self.server_process.pid})...")
-            self.server_process.terminate()  # Send SIGTERM
-            self.server_process.join(timeout=5)  # Wait for clean exit
-            if self.server_process.is_alive():
-                logger.warning(
-                    f"AgentOps server (PID: {self.server_process.pid}) did not terminate gracefully, killing..."
-                )
-                self.server_process.kill()  # Force kill
-                self.server_process.join(timeout=10)  # Wait for kill
-            self.server_process = None
-            logger.info(f"AgentOps local server stopped.")
+    def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
+        if _agentops_service_enabled:
+            return super().export(*args, **kwargs)
        else:
-            logger.info("AgentOps local server was not running or already stopped.")
+            logger.debug("SwitchableAuthenticatedOTLPExporter is switched off, skipping export.")
+            return SpanExportResult.SUCCESS

-    def get_port(self) -> int | None:
-        # Check liveness again in case it died since start()
-        if self.is_alive() and self.server_port is not None:
-            return self.server_port
-        # If called after server stopped or failed, port might be stale or None
-        if self.server_port is not None and (self.server_process is None or not self.server_process.is_alive()):
-            logger.warning(
-                f"AgentOps server port {self.server_port} is stored, but server process is not alive. Returning stored port."
-            )
-        return self.server_port
+
+class BypassableOTLPMetricExporter(OTLPMetricExporter):
+    """
+    OTLPMetricExporter with switchable service control.
+    When `_agentops_service_enabled` is False, skip export and return success.
+    """
+
+    def export(self, *args: Any, **kwargs: Any) -> MetricExportResult:
+        if _agentops_service_enabled:
+            return super().export(*args, **kwargs)  # type: ignore[reportUnknownMemberType]
+        else:
+            logger.debug("SwitchableOTLPMetricExporter is switched off, skipping export.")
+            return MetricExportResult.SUCCESS
+
+
+class BypassableOTLPSpanExporter(OTLPSpanExporter):
+    """
+    OTLPSpanExporter 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("SwitchableOTLPSpanExporter is switched off, skipping export.")
+            return SpanExportResult.SUCCESS
+
+
+class BypassableV3Client(V3Client):
+    """
+    V3Client with toggleable authentication calls.
+    Returns dummy auth response when `_agentops_service_enabled` is False.
+    """
+
+    # Temporary synchronous override of fetch_auth_token for mock purposes.
+    def fetch_auth_token(self, *args: Any, **kwargs: Any) -> AuthTokenResponse:  # type: ignore[override]
+        if _agentops_service_enabled:
+            return super().fetch_auth_token(*args, **kwargs)  # type: ignore[override]
+        else:
+            logger.debug("SwitchableV3Client is switched off, skipping fetch_auth_token request.")
+            return AuthTokenResponse(token="dummy", project_id="dummy")
+
+
+class BypassableV4Client(V4Client):
+    """
+    V4Client with toggleable post requests.
+    Returns dummy response when `_agentops_service_enabled` is False.
+    """
+
+    def post(self, *args: Any, **kwargs: Any) -> requests.Response:
+        if _agentops_service_enabled:
+            return super().post(*args, **kwargs)
+        else:
+            logger.debug("SwitchableV4Client is switched off, skipping post request.")
+            response = requests.Response()
+            response.status_code = 200
+            response._content = b"{}"
+            return response
@@ -1,20 +1,27 @@
-from typing import Dict, Any
-from agentops.integration.callbacks.langchain import LangchainCallbackHandler
-from agentops import instrumentation
+# Copyright (c) Microsoft. All rights reserved.

+from typing import Any, Dict
+
+from agentops import instrumentation
+from agentops.integration.callbacks.langchain import LangchainCallbackHandler

 original_on_chain_start = LangchainCallbackHandler.on_chain_start
 langgraph_entry = None

+__all__ = [
+    "instrument_agentops_langchain",
+    "uninstrument_agentops_langchain",
+]

-def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
+
+def on_chain_start(self: Any, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
    if "name" in kwargs:
-        if serialized is None:
+        if serialized is None:  # type: ignore
            serialized = {}
        serialized = serialized.copy()
        serialized["name"] = kwargs["name"]
    if "run_id" in kwargs:
-        if serialized is None:
+        if serialized is None:  # type: ignore
            serialized = {}
        serialized = serialized.copy()
        if "id" not in serialized:
@@ -23,12 +30,14 @@ def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **k


 def instrument_agentops_langchain():
+    """Bypass AgentOp's native support for Langchain."""
    global langgraph_entry
    langgraph_entry = instrumentation.AGENTIC_LIBRARIES.pop("langgraph", None)
    LangchainCallbackHandler.on_chain_start = on_chain_start


 def uninstrument_agentops_langchain():
+    """Restore AgentOp's native support for Langchain."""
    global langgraph_entry
    if langgraph_entry is not None:
        instrumentation.AGENTIC_LIBRARIES["langgraph"] = langgraph_entry
@@ -1,26 +1,39 @@
-from typing import Optional, Any
+# Copyright (c) Microsoft. All rights reserved.
+
+"""LiteLLM instrumentations.
+
+It's unclear whether or not this file is useful.
+It seems that LiteLLM owns its own telemetry from their own entrance
+
+[Related documentation](https://docs.litellm.ai/docs/observability/agentops_integration).
+"""
+
+from typing import Any, Optional

 from litellm.integrations.opentelemetry import OpenTelemetry

-# It's unclear whether or not this file is useful
-# It seems that LiteLLM owns its own telemetry from their own entrance
-# https://docs.litellm.ai/docs/observability/agentops_integration
+__all__ = [
+    "instrument_litellm",
+    "uninstrument_litellm",
+]

-original_set_attributes = OpenTelemetry.set_attributes
+original_set_attributes = OpenTelemetry.set_attributes  # type: ignore


-def patched_set_attributes(self, span: Any, kwargs, response_obj: Optional[Any]):
+def patched_set_attributes(self: Any, span: Any, kwargs: Any, response_obj: Optional[Any]):
    original_set_attributes(self, span, kwargs, response_obj)
    # Add custom attributes
-    if response_obj.get("prompt_token_ids"):
+    if response_obj is not None and response_obj.get("prompt_token_ids"):
        span.set_attribute("prompt_token_ids", list(response_obj.get("prompt_token_ids")))
-    if response_obj.get("response_token_ids"):
+    if response_obj is not None and response_obj.get("response_token_ids"):
        span.set_attribute("response_token_ids", list(response_obj.get("response_token_ids")[0]))


 def instrument_litellm():
+    """Instrument litellm to capture token IDs."""
    OpenTelemetry.set_attributes = patched_set_attributes


 def uninstrument_litellm():
+    """Uninstrument litellm to stop capturing token IDs."""
    OpenTelemetry.set_attributes = original_set_attributes
@@ -1,148 +0,0 @@
-# type: ignore
-
-# https://github.com/volcengine/verl/blob/bd94bd61fe4193e56f2845dc794004afbef7f818/examples/ppo_trainer/naive_chat_scheduler.py
-# This file is part of VERL example. It should be included in the VERL package but it's not currently.
-
-# Copyright 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import asyncio
-from typing import Any, Dict, List
-
-import torch
-from openai.types.chat.chat_completion import ChatCompletion
-from tensordict import TensorDict
-
-from verl.protocol import DataProto
-from verl.workers.rollout.async_server import ChatCompletionScheduler
-
-
-class NaiveChatCompletionScheduler(ChatCompletionScheduler):
-    """
-    A very naive implementation of ChatCompletionScheduler for demo purpose,
-    only do single-turn chat completion.
-    """
-
-    async def generate_sequences(self, batch: DataProto, **sampling_params) -> DataProto:
-        kwargs = dict(
-            n=self.config.n,
-            max_completion_tokens=self.config.response_length,
-            temperature=self.config.temperature,
-            top_p=self.config.top_p,
-        )
-
-        do_sample = batch.meta_info.get("do_sample", True)
-        is_validate = batch.meta_info.get("validate", False)
-        if not do_sample or is_validate:
-            kwargs["n"] = 1
-            kwargs["temperature"] = 0
-
-        kwargs.update(sampling_params)
-        print(f"[NaiveChatCompletionScheduler] generate_sequences sampling params: {kwargs}")
-
-        async def callback(completions: ChatCompletion, info: Dict[str, Any], exception: Exception):
-            assert exception is None, f"exception: {exception}"
-            conversation, batch_conversations, batch_index = (
-                info["conversation"],
-                info["batch_conversations"],
-                info["batch_index"],
-            )
-
-            conversations = []
-            for choice in completions.choices:
-                chat = conversation.copy()
-                chat.append({"role": choice.message.role, "content": choice.message.content})
-                conversations.append(chat)
-            batch_conversations[batch_index] = conversations
-
-            # NOTE: we can call tools and resubmit chat completions here.
-            # call_tools(completions, info)
-            # await self.submit_chat_completions(callback2, ...)
-
-        # TODO: we may need to control max concurrent requests here, or it will harm prefix cache hit rate.
-        tasks, batch_conversations = [], [None] * len(batch)
-        for batch_index, conversation in enumerate(batch.non_tensor_batch["raw_prompt"]):
-            # raw_prompt: [{"role": "user", "content": ""}, ["role": "assistant", "content"], ...]
-            tasks.append(
-                asyncio.create_task(
-                    self.submit_chat_completions(
-                        callback=callback,
-                        callback_additional_info={
-                            "batch_conversations": batch_conversations,
-                            "batch_index": batch_index,
-                            "conversation": list(conversation),
-                        },
-                        model=self.model_name,
-                        messages=conversation.tolist(),
-                        **kwargs,
-                    )
-                )
-            )
-        await asyncio.gather(*tasks)
-        print("[NaiveChatCompletionScheduler] generate_sequences done")
-
-        return self._postprocess(batch, batch_conversations, kwargs["n"])
-
-    def _postprocess(
-        self, batch: DataProto, batch_conversations: List[List[List[Dict[str, str]]]], n: int
-    ) -> DataProto:
-        # NOTE: consistent with batch version of generate_sequences in vllm_rollout_spmd.py
-        # prompts: left pad
-        # responses: right pad
-        # input_ids: prompt + response
-        # attention_mask: [0,0,0,0,1,1,1,1, | 1,1,1,0,0,0,0,0]
-        # position_ids:   [0,0,0,0,0,1,2,3, | 4,5,6,7,8,9,10,11]
-
-        # prompts: [prompt] from input dataset
-        prompts = [
-            self.tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=False)
-            for prompt in batch.non_tensor_batch["raw_prompt"]
-        ]
-
-        # flatten batch_conversations if n > 1
-        assert len(batch_conversations) == len(prompts)
-        batch_conversations = [conversation for conversations in batch_conversations for conversation in conversations]
-        assert len(batch_conversations) == len(prompts) * n
-
-        # sequences: [prompt + response]
-        sequences = [
-            self.tokenizer.apply_chat_template(conversation, add_generation_prompt=False, tokenize=False)
-            for conversation in batch_conversations
-        ]
-
-        # responses: [response]
-        # TODO: mask out tools calling tokens?
-        responses = [sequence[len(prompts[i // n]) :] for i, sequence in enumerate(sequences)]
-
-        prompts = self.tokenizer(prompts, return_tensors="pt", padding="longest", padding_side="left")
-        responses = self.tokenizer(responses, return_tensors="pt", padding="longest", padding_side="right")
-        if n > 1:
-            prompts["input_ids"] = prompts["input_ids"].repeat_interleave(n, dim=0)
-            prompts["attention_mask"] = prompts["attention_mask"].repeat_interleave(n, dim=0)
-
-        input_ids = torch.cat([prompts["input_ids"], responses["input_ids"]], dim=1)
-        attention_mask = torch.cat([prompts["attention_mask"], responses["attention_mask"]], dim=1)
-        position_ids = (attention_mask.cumsum(dim=1) - 1) * attention_mask
-
-        batch = TensorDict(
-            {
-                "prompts": prompts["input_ids"],
-                "responses": responses["input_ids"],
-                "input_ids": input_ids,
-                "attention_mask": attention_mask,
-                "position_ids": position_ids,
-            },
-            batch_size=len(input_ids),
-        )
-
-        return DataProto(batch=batch)
@@ -1,12 +1,19 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 from __future__ import annotations

 import warnings
-from typing import List
+from typing import Any, List

-from vllm.entrypoints.openai.protocol import ChatCompletionResponse
 import vllm.entrypoints.openai.protocol
+from vllm.entrypoints.openai.protocol import ChatCompletionResponse
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat

+__all__ = [
+    "instrument_vllm",
+    "uninstrument_vllm",
+]
+

 class ChatCompletionResponsePatched(ChatCompletionResponse):
    prompt_token_ids: List[int] | None = None
@@ -17,15 +24,15 @@ original_chat_completion_full_generator = OpenAIServingChat.chat_completion_full


 async def chat_completion_full_generator(
-    self,
-    request,
-    result_generator,
+    self: Any,
+    request: Any,
+    result_generator: Any,
    request_id: str,
    model_name: str,
-    conversation,
-    tokenizer,
-    request_metadata,
-):
+    conversation: Any,
+    tokenizer: Any,
+    request_metadata: Any,
+) -> Any:
    prompt_token_ids: List[int] | None = None
    response_token_ids: List[List[int]] | None = None

@@ -57,6 +64,10 @@ async def chat_completion_full_generator(


 def instrument_vllm():
+    """Instrument vLLM to capture token IDs generated by engine.
+
+    This instrumentation has been merged to upstream vLLM since v0.10.2.
+    """
    if vllm.entrypoints.openai.protocol.ChatCompletionResponse is ChatCompletionResponsePatched:
        warnings.warn("vllm is already instrumented. Skip the instrumentation.")
        return
@@ -66,4 +77,5 @@ def instrument_vllm():


 def uninstrument_vllm():
+    """Uninstrument vLLM to stop capturing token IDs generated by engine."""
    OpenAIServingChat.chat_completion_full_generator = original_chat_completion_full_generator
@@ -1,204 +0,0 @@
-from __future__ import annotations
-
-import logging
-import weakref
-from typing import Any, List, Dict, Union, Optional, TYPE_CHECKING
-
-from .types import NamedResources, Rollout, Task, TaskInput, Triplet, RolloutRawResult
-
-if TYPE_CHECKING:
-    from .trainer import Trainer
-    from .runner import AgentRunner
-    from .tracer import BaseTracer
-
-
-logger = logging.getLogger(__name__)
-
-
-class LitAgent:
-    """Base class for the training and validation logic of an agent.
-
-    Developers should subclass this class and implement the rollout methods
-    to define the agent's behavior for a single task. The agent's logic
-    is completely decoupled from the server communication and training
-    infrastructure.
-    """
-
-    def __init__(self, *, trained_agents: Optional[str] = None) -> None:  # FIXME: str | None won't work for cli
-        """
-        Initialize the LitAgent.
-
-        Args:
-            trained_agents: Optional string representing the trained agents.
-                            This can be used to track which agents have been trained by this instance.
-        """
-        self.trained_agents = trained_agents
-        self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
-        self._runner_ref: weakref.ReferenceType[AgentRunner] | None = None
-
-    def set_trainer(self, trainer: Trainer) -> None:
-        """
-        Set the trainer for this agent.
-
-        Args:
-            trainer: The Trainer instance that will handle training and validation.
-        """
-        self._trainer_ref = weakref.ref(trainer)
-
-    @property
-    def trainer(self) -> Trainer:
-        """
-        Get the trainer for this agent.
-
-        Returns:
-            The Trainer instance associated with this agent.
-        """
-        if self._trainer_ref is None:
-            raise ValueError("Trainer has not been set for this agent.")
-        trainer = self._trainer_ref()
-        if trainer is None:
-            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
-        return trainer
-
-    @property
-    def tracer(self) -> BaseTracer:
-        """
-        Get the tracer for this agent.
-
-        Returns:
-            The BaseTracer instance associated with this agent.
-        """
-        return self.trainer.tracer
-
-    def set_runner(self, runner: AgentRunner) -> None:
-        """
-        Set the runner for this agent.
-
-        Args:
-            runner: The AgentRunner instance that will handle the execution of rollouts.
-        """
-        self._runner_ref = weakref.ref(runner)
-
-    @property
-    def runner(self) -> AgentRunner:
-        """
-        Get the runner for this agent.
-
-        Returns:
-            The AgentRunner instance associated with this agent.
-        """
-        if self._runner_ref is None:
-            raise ValueError("Runner has not been set for this agent.")
-        runner = self._runner_ref()
-        if runner is None:
-            raise ValueError("Runner reference is no longer valid (object has been garbage collected).")
-        return runner
-
-    def on_rollout_start(self, task: Task, runner: AgentRunner, tracer: BaseTracer) -> None:
-        """Hook called immediately before a rollout begins.
-
-        Args:
-            task: The :class:`Task` object that will be processed.
-            runner: The :class:`AgentRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
-
-        Subclasses can override this method to implement custom logic such as
-        logging, metric collection, or resource setup. By default, this is a
-        no-op.
-        """
-
-    def on_rollout_end(self, task: Task, rollout: Rollout, runner: AgentRunner, tracer: BaseTracer) -> None:
-        """Hook called after a rollout completes.
-
-        Args:
-            task: The :class:`Task` object that was processed.
-            rollout: The resulting :class:`Rollout` object.
-            runner: The :class:`AgentRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
-
-        Subclasses can override this method for cleanup or additional
-        logging. By default, this is a no-op.
-        """
-
-    def training_rollout(self, task: TaskInput, rollout_id: str, resources: NamedResources) -> RolloutRawResult:
-        """Defines the agent's behavior for a single training task.
-
-        This method should contain the logic for how the agent processes an
-        input, uses the provided resources (like LLMs or prompts), and
-        produces a result.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            rollout_id: A unique identifier for the rollout, used for tracking
-                        and reporting purposes.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-
-        Returns:
-            The result of the rollout, which can be one of:
-            - None. The tracing should be handled by the agent runner.
-            - A float representing the final reward.
-            - A list of `Triplet` objects for detailed, step-by-step feedback.
-            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
-            - A list of dictionaries for any trace spans.
-            - A complete `Rollout` object for full control over reporting.
-        """
-        raise NotImplementedError("Subclasses must implement the `training_rollout` method.")
-
-    def validation_rollout(self, task: TaskInput, rollout_id: str, resources: NamedResources) -> RolloutRawResult:
-        """Defines the agent's behavior for a single validation task.
-
-        By default, this method redirects to `training_rollout`. Override it
-        if the agent should behave differently during validation.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            rollout_id: A unique identifier for the validation rollout,
-                        used for tracking and reporting purposes.
-            resources: A dictionary of named resources for the agent to use.
-
-        Returns:
-            The result of the validation rollout. See `training_rollout` for
-            possible return types.
-        """
-        return self.training_rollout(task, rollout_id, resources)
-
-    async def training_rollout_async(
-        self, task: TaskInput, rollout_id: str, resources: NamedResources
-    ) -> RolloutRawResult:
-        """Asynchronous version of `training_rollout`.
-
-        This method should be implemented by agents that perform asynchronous
-        operations (e.g., non-blocking I/O, concurrent API calls).
-
-        Args:
-            task: The task object received from the server.
-            rollout_id: A unique identifier for the training rollout,
-                        used for tracking and reporting purposes.
-            resources: A dictionary of named resources for the agent to use.
-
-        Returns:
-            The result of the asynchronous training rollout.
-        """
-        raise NotImplementedError("Async agents must implement the `training_rollout_async` method.")
-
-    async def validation_rollout_async(
-        self, task: TaskInput, rollout_id: str, resources: NamedResources
-    ) -> RolloutRawResult:
-        """Asynchronous version of `validation_rollout`.
-
-        By default, this method redirects to `training_rollout_async`.
-        Override it for different asynchronous validation behavior.
-
-        Args:
-            task: The task object received from the server.
-            rollout_id: A unique identifier for the validation rollout,
-                        used for tracking and reporting purposes.
-            resources: A dictionary of named resources for the agent to use.
-
-        Returns:
-            The result of the asynchronous validation rollout.
-        """
-        return await self.training_rollout_async(task, rollout_id, resources)
@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .decorator import *
+from .litagent import *
+
+__all__ = [
+    "LitAgent",
+    "llm_rollout",
+    "prompt_rollout",
+    "rollout",
+]
@@ -0,0 +1,536 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Convenience decorators for building lightweight `LitAgent` implementations."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from typing import Any, Awaitable, Callable, Dict, Protocol, TypeGuard, TypeVar, Union, overload
+
+from agentlightning.types import (
+    LLM,
+    AttemptedRollout,
+    NamedResources,
+    PromptTemplate,
+    ProxyLLM,
+    Rollout,
+    RolloutRawResult,
+)
+
+from .litagent import LitAgent
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+__all__ = [
+    "llm_rollout",
+    "prompt_rollout",
+    "rollout",
+]
+
+
+T_contra = TypeVar("T_contra", contravariant=True)
+
+
+class LlmRolloutFuncSync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM) -> RolloutRawResult: ...
+
+
+class LlmRolloutFuncSync3(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> RolloutRawResult: ...
+
+
+class LlmRolloutFuncAsync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM) -> Awaitable[RolloutRawResult]: ...
+
+
+class LlmRolloutFuncAsync3(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> Awaitable[RolloutRawResult]: ...
+
+
+LlmRolloutFunc = Union[
+    LlmRolloutFuncSync2[T_contra],
+    LlmRolloutFuncSync3[T_contra],
+    LlmRolloutFuncAsync2[T_contra],
+    LlmRolloutFuncAsync3[T_contra],
+]
+
+
+class PromptRolloutFuncSync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> RolloutRawResult: ...
+
+
+class PromptRolloutFuncAsync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> Awaitable[RolloutRawResult]: ...
+
+
+class PromptRolloutFuncSync3(Protocol[T_contra]):
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout) -> RolloutRawResult: ...
+
+
+class PromptRolloutFuncAsync3(Protocol[T_contra]):
+    def __call__(
+        self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout
+    ) -> Awaitable[RolloutRawResult]: ...
+
+
+PromptRolloutFunc = Union[
+    PromptRolloutFuncSync2[T_contra],
+    PromptRolloutFuncSync3[T_contra],
+    PromptRolloutFuncAsync2[T_contra],
+    PromptRolloutFuncAsync3[T_contra],
+]
+
+
+class FunctionalLitAgentFunc(Protocol[T_contra]):
+    def __call__(
+        self, task: T_contra, *args: Any, **kwargs: Any
+    ) -> Union[RolloutRawResult, Awaitable[RolloutRawResult]]: ...
+
+
+class FunctionalLitAgent(LitAgent[T]):
+    """Adapter that turns plain rollout functions into [`LitAgent`][agentlightning.LitAgent] instances.
+
+    The helper inspects the wrapped function to determine which resources to
+    inject, allowing both synchronous and asynchronous callables to participate
+    in the training loop without writing a dedicated subclass.
+    """
+
+    def __init__(self, rollout_func: FunctionalLitAgentFunc[T], *, strip_proxy: bool = True) -> None:
+        """Initialize the wrapper around a rollout function.
+
+        Args:
+            rollout_func: Callable that implements the rollout. It may be synchronous
+                or asynchronous and can optionally receive a
+                [`Rollout`][agentlightning.Rollout] alongside resources such as
+                `llm` or `prompt_template`.
+            strip_proxy: When ``True``, convert
+                [`ProxyLLM`][agentlightning.ProxyLLM] inputs into
+                [`LLM`][agentlightning.LLM] instances before calling the
+                rollout function. Defaults to `True`.
+        """
+        super().__init__()
+        self._rollout_func = rollout_func
+        self._strip_proxy = strip_proxy
+        self._is_async = inspect.iscoroutinefunction(rollout_func)
+        self._sig = inspect.signature(rollout_func)
+
+        # Copy function metadata to preserve type hints and other attributes
+        functools.update_wrapper(self, rollout_func)  # type: ignore
+
+    def _accepts_rollout(self) -> bool:
+        return "rollout" in self._sig.parameters
+
+    def _accepts_llm(self) -> bool:
+        return "llm" in self._sig.parameters
+
+    def _accepts_prompt_template(self) -> bool:
+        return "prompt_template" in self._sig.parameters
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Make the agent instance callable, preserving the original function behavior."""
+        return self._rollout_func(*args, **kwargs)  # type: ignore
+
+    def is_async(self) -> bool:
+        return self._is_async
+
+    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Execute a synchronous rollout using the wrapped function.
+
+        Args:
+            task: Task input data.
+            resources: Mapping of named resources available to the agent.
+            rollout: Rollout metadata provided by the runtime.
+
+        Returns:
+            Result produced by the wrapped rollout function.
+
+        Raises:
+            RuntimeError: If the wrapped function is asynchronous.
+        """
+        if self._is_async:
+            raise RuntimeError(f"{self._rollout_func} is asynchronous. Use rollout_async instead.")
+
+        kwargs = self._get_kwargs(resources, rollout)
+        return self._rollout_func(task, **kwargs)  # type: ignore
+
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Execute an asynchronous rollout using the wrapped function.
+
+        Args:
+            task: Task input data.
+            resources: Mapping of named resources available to the agent.
+            rollout: Rollout metadata provided by the runtime.
+
+        Returns:
+            Result produced by the wrapped rollout coroutine.
+
+        Raises:
+            RuntimeError: If the wrapped function is synchronous.
+        """
+        if not self._is_async:
+            raise RuntimeError(f"{self._rollout_func} is synchronous. Use rollout instead.")
+
+        kwargs = self._get_kwargs(resources, rollout)
+        return await self._rollout_func(task, **kwargs)  # type: ignore
+
+    def _get_kwargs(self, resources: NamedResources, rollout: Rollout) -> Dict[str, Any]:
+        """Prepare keyword arguments expected by the wrapped rollout function.
+
+
+        It dynamically builds the `kwargs` dictionary by inspecting the function signature and
+        including only the parameters the function accepts. This allows flexible function
+        signatures that can request any combination of: rollout, llm, and/or prompt_template.
+
+        Args:
+            resources: Mapping of named resources available for the rollout.
+            rollout: Rollout metadata provided by the runtime.
+
+        Returns:
+            Dictionary of keyword arguments to forward to the rollout function.
+        """
+
+        kwargs: Dict[str, Any] = {}
+        if self._accepts_rollout():
+            kwargs["rollout"] = rollout
+        if self._accepts_llm():
+            kwargs["llm"] = self._get_llm_resource(resources, rollout)
+        if self._accepts_prompt_template():
+            kwargs["prompt_template"] = self._get_prompt_template_resource(resources, rollout)
+
+        return kwargs
+
+    def _get_llm_resource(self, resources: NamedResources, rollout: Rollout) -> LLM:
+        """Retrieve the first LLM resource from the available resources.
+
+        Strip the ProxyLLM resource into a LLM resource if needed.
+
+        Args:
+            resources: Mapping of named resources.
+            rollout: Rollout metadata used when stripping proxy endpoints.
+
+        Returns:
+            First [`LLM`][agentlightning.LLM] resource encountered.
+
+        Raises:
+            ValueError: If no LLM resource is present.
+        """
+        resource_found: LLM | None = None
+        for name, resource in resources.items():
+            if isinstance(resource, LLM):
+                if resource_found is not None:
+                    logger.warning(f"Multiple LLM resources found in resources. Using the first one: '{name}'.")
+                    break
+                resource_found = resource
+
+        if resource_found is None:
+            raise ValueError("No LLM resource found in the provided resources.")
+
+        if self._strip_proxy:
+            resource_found = self._strip_proxy_helper(resource_found, rollout)
+
+        return resource_found
+
+    def _get_prompt_template_resource(self, resources: NamedResources, rollout: Rollout) -> PromptTemplate:
+        """Retrieve the first prompt template resource from the available resources.
+
+        Args:
+            resources: Mapping of named resources.
+            rollout: Rollout metadata (unused).
+
+        Returns:
+            First [`PromptTemplate`][agentlightning.PromptTemplate] resource encountered.
+
+        Raises:
+            ValueError: If no prompt template resource is present.
+        """
+        resource_found: PromptTemplate | None = None
+        for name, resource in resources.items():
+            if isinstance(resource, PromptTemplate):
+                if resource_found is not None:
+                    logger.warning(
+                        f"Multiple prompt template resources found in resources. Using the first one: '{name}'."
+                    )
+                    break
+                resource_found = resource
+
+        if resource_found is None:
+            raise ValueError("No prompt template resource found in the provided resources.")
+
+        return resource_found
+
+    def _strip_proxy_helper(self, proxy_llm: LLM, rollout: Rollout) -> LLM:
+        """Convert [`ProxyLLM`][agentlightning.ProxyLLM] instances into concrete LLMs.
+
+        It resolves ProxyLLM instances to their concrete LLM implementation
+        by attaching the attempted rollout context. This is only used when the function
+        signature accepts an `llm` parameter and strip_proxy is True.
+
+        Args:
+            proxy_llm: Candidate LLM resource.
+            rollout: Rollout metadata that provides rollout and attempt identifiers.
+
+        Returns:
+            [`LLM`][agentlightning.LLM] with rollout context baked into the endpoint.
+
+        Raises:
+            ValueError: If the rollout is not an
+                [`AttemptedRollout`][agentlightning.AttemptedRollout].
+        """
+
+        if not isinstance(proxy_llm, ProxyLLM):
+            # Not a ProxyLLM, nothing to strip here.
+            return proxy_llm
+
+        # Rollout is still a Rollout here because API is not stabilized yet.
+        # In practice, it must be an AttemptedRollout.
+        if not isinstance(rollout, AttemptedRollout):
+            raise ValueError("Rollout is not an AttemptedRollout.")
+
+        return proxy_llm.with_attempted_rollout(rollout)
+
+
+@overload
+def llm_rollout(func: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]: ...
+
+
+@overload
+def llm_rollout(*, strip_proxy: bool = True) -> Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]: ...
+
+
+def llm_rollout(
+    func: LlmRolloutFunc[T] | None = None, *, strip_proxy: bool = True
+) -> FunctionalLitAgent[T] | Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]:
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for LLM-based rollouts.
+
+    Args:
+        func: Callable defining the agent's behaviour. Supported signatures include:
+
+            * `(task, llm) -> result`
+            * `(task, llm, rollout) -> result`
+            * `async (task, llm) -> result`
+            * `async (task, llm, rollout) -> result`
+
+        strip_proxy: When `True`, convert proxy resources into concrete
+            [`LLM`][agentlightning.LLM] instances before calling the
+            function. Defaults to `True`.
+
+    Returns:
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.
+
+    Examples:
+        ```python
+        @llm_rollout
+        def my_agent(task, llm):
+            return llm.endpoint
+
+        @llm_rollout(strip_proxy=False)
+        def my_agent_no_strip(task, llm):
+            return llm.model
+
+        result = my_agent(task, llm)
+        result = my_agent.rollout(task, resources, rollout)
+        ```
+    """
+
+    def decorator(f: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]:
+        _validate_llm_rollout_func(f)
+        return FunctionalLitAgent(f, strip_proxy=strip_proxy)
+
+    if func is None:
+        # Called with arguments: @llm_rollout(strip_proxy=False)
+        return decorator
+    else:
+        # Called without arguments: @llm_rollout
+        return decorator(func)
+
+
+def _validate_llm_rollout_func(func: Any) -> TypeGuard[LlmRolloutFunc[Any]]:
+    """Validate the function signature of an LLM rollout function.
+
+    Ensures the function follows the expected pattern for LLM-based rollouts:
+
+    - Must have at least 2 parameters
+    - First parameter must be named 'task'
+    - Must have a parameter named 'llm'
+    - Optionally can have a 'rollout' parameter
+
+    Args:
+        func: Function to inspect.
+
+    Returns:
+        `True` when the signature matches the supported patterns.
+
+    Raises:
+        ValueError: If the function signature does not match the expected pattern.
+    """
+    sig = inspect.signature(func)
+    params = list(sig.parameters.keys())
+    if len(params) < 2:
+        raise ValueError(f"Function {func} must have at least 2 parameters.")
+    if params[0] != "task":
+        raise ValueError(f"Function {func} must be a positional parameter called 'task'.")
+    if "llm" not in params:
+        raise ValueError(f"Function {func} must have a positional parameter called 'llm'.")
+
+    return True
+
+
+@overload
+def prompt_rollout(func: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]: ...
+
+
+@overload
+def prompt_rollout() -> Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]: ...
+
+
+def prompt_rollout(
+    func: PromptRolloutFunc[T] | None = None,
+) -> FunctionalLitAgent[T] | Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]:
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for prompt-based rollouts.
+
+    This decorator is designed for agents that work with tunable prompt templates. It enables
+    a workflow where algorithms manage and optimize the prompt template, while agents consume
+    the template to perform rollouts. This is particularly useful for prompt optimization scenarios.
+
+    Args:
+        func: Callable defining the agent's behavior. Supported signatures include:
+
+            * `(task, prompt_template) -> result`
+            * `(task, prompt_template, rollout) -> result`
+            * `async (task, prompt_template) -> result`
+            * `async (task, prompt_template, rollout) -> result`
+
+    Returns:
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.
+
+    Examples:
+        ```python
+        @prompt_rollout
+        def my_agent(task, prompt_template):
+            messages = prompt_template.format(task=task.input)
+            return messages
+
+        result = my_agent(task, prompt_template)
+        result = my_agent.rollout(task, resources, rollout)
+        ```
+    """
+
+    def decorator(f: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]:
+        _validate_prompt_rollout_func(f)
+        return FunctionalLitAgent(f)
+
+    if func is None:
+        return decorator
+    else:
+        return decorator(func)
+
+
+def _validate_prompt_rollout_func(func: Any) -> TypeGuard[PromptRolloutFunc[Any]]:
+    """Validate the function signature of a prompt rollout function.
+
+    Ensures the function follows the expected pattern for prompt-template-based rollouts:
+
+    - Must have at least 2 parameters
+    - First parameter must be named 'task'
+    - Must have a parameter named 'prompt_template'
+    - Optionally can have a 'rollout' parameter
+
+    Args:
+        func: Function to inspect.
+
+    Returns:
+        `True` when the signature matches the supported patterns.
+
+    Raises:
+        ValueError: If the function signature does not match the expected pattern.
+    """
+    sig = inspect.signature(func)
+    params = list(sig.parameters.keys())
+    if len(params) < 2:
+        raise ValueError(f"Function {func} must have at least 2 parameters.")
+    if params[0] != "task":
+        raise ValueError(f"Function {func} must be a positional parameter called 'task'.")
+    if "prompt_template" not in params:
+        raise ValueError(f"Function {func} must have a positional parameter called 'prompt_template'.")
+
+    return True
+
+
+def rollout(func: Union[LlmRolloutFunc[T], PromptRolloutFunc[T], Callable[..., Any]]) -> FunctionalLitAgent[T]:
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] from an arbitrary rollout function.
+
+    This function inspects the provided callable and creates the appropriate
+    agent type based on its signature. It supports both LLM-based and prompt-template-based
+    agents. The returned agent instance is callable, preserving the original function's
+    behavior and type hints.
+
+    See [`llm_rollout`][agentlightning.litagent.decorator.llm_rollout] and
+    [`prompt_rollout`][agentlightning.litagent.decorator.prompt_rollout] for more details.
+
+    Args:
+        func: Callable that implements the rollout. Supported signatures:
+
+            - `[async ](task, llm[, rollout])` for LLM-based agents
+            - `[async ](task, prompt_template[, rollout])` for prompt-template-based agents
+
+            The supported output types of `func` is same as the return type of [`rollout`][agentlightning.LitAgent.rollout].
+
+    Returns:
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.
+
+    Examples:
+        ```python
+        # LLM-based agent
+        @rollout
+        def my_llm_agent(task, llm):
+            client = OpenAI(base_url=llm.endpoint)
+            response = client.chat.completions.create(
+                model=llm.model,
+                messages=[{"role": "user", "content": task.input}],
+            )
+            return response
+
+        # Prompt-template-based agent
+        @rollout
+        def my_prompt_agent(task, prompt_template):
+            messages = prompt_template.format(task=task.input)
+            # ... perform rollout with the formatted prompt
+            return response
+
+        # Function is still callable with original behavior
+        result = my_llm_agent(task, llm)
+
+        # Agent methods are also available
+        result = my_llm_agent.rollout(task, resources, rollout)
+        ```
+
+    Raises:
+        NotImplementedError: If the function signature doesn't match any known patterns.
+    """
+    # Check if it matches the LLM rollout API pattern
+    sig = inspect.signature(func)
+
+    try:
+        if _validate_llm_rollout_func(func):
+            return llm_rollout(func)
+    except ValueError:
+        pass
+
+    try:
+        if _validate_prompt_rollout_func(func):
+            return prompt_rollout(func)
+    except ValueError:
+        pass
+
+    raise NotImplementedError(
+        f"Function signature {sig} does not match any known agent patterns. "
+        "Expected signatures: (task, llm[, rollout]) or (task, prompt_template[, rollout]). "
+        "Functions can be sync or async."
+    )
@@ -0,0 +1,251 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Base abstractions for building agents that plug into Agent Lightning."""
+
+from __future__ import annotations
+
+import inspect
+import logging
+import warnings
+import weakref
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar
+
+from agentlightning.types import NamedResources, Rollout, RolloutRawResult, Task
+
+if TYPE_CHECKING:
+    from agentlightning.runner import Runner
+    from agentlightning.tracer import Tracer
+    from agentlightning.trainer import Trainer
+
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+__all__ = [
+    "LitAgent",
+]
+
+
+def is_v0_1_rollout_api(func: Callable[..., Any]) -> bool:
+    """Return `True` when the rollout function uses the deprecated v0.1 signature.
+
+    The helper inspects the callable's signature to detect whether a `rollout_id`
+    parameter is present, which indicates the legacy API.
+
+    Args:
+        func: Function to analyze.
+
+    Returns:
+        `True` if the callable exposes a `rollout_id` parameter.
+    """
+    return "rollout_id" in inspect.signature(func).parameters
+
+
+class LitAgent(Generic[T]):
+    """Base class for implementing agent rollouts.
+
+    Subclasses override the rollout methods to process tasks while the trainer and
+    runner infrastructure manages orchestration, tracing, and persistence.
+    """
+
+    def __init__(self, *, trained_agents: Optional[str] = None) -> None:  # FIXME: str | None won't work for cli
+        """Initialize the agent instance.
+
+        Args:
+            trained_agents: Optional identifier used by legacy tooling to mark trained
+                agents.
+
+        !!! warning "Deprecated"
+            The `trained_agents` flag is deprecated. Configure `agent_match` in the adapter
+            layer instead. See [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet]
+            for more details.
+        """
+        if trained_agents is not None:
+            warnings.warn(
+                "`trained_agents` is deprecated. Configure `agent_match` in adapter instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        self.trained_agents = trained_agents
+
+        self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
+        self._runner_ref: weakref.ReferenceType[Runner[T]] | None = None
+
+    def is_async(self) -> bool:
+        """Return `True` when the agent overrides any asynchronous rollout methods.
+
+        Override this method for customized async detection logic.
+        """
+        return (
+            (
+                hasattr(self, "training_rollout_async")
+                and self.__class__.training_rollout_async is not LitAgent.training_rollout_async  # type: ignore
+            )
+            or (
+                hasattr(self, "validation_rollout_async")
+                and self.__class__.validation_rollout_async is not LitAgent.validation_rollout_async  # type: ignore
+            )
+            or (hasattr(self, "rollout_async") and self.__class__.rollout_async is not LitAgent.rollout_async)  # type: ignore
+        )
+
+    def set_trainer(self, trainer: Trainer) -> None:
+        """Attach the trainer responsible for orchestration.
+
+        Args:
+            trainer: [`Trainer`][agentlightning.Trainer] that manages the agent.
+        """
+        self._trainer_ref = weakref.ref(trainer)
+
+    def get_trainer(self) -> Trainer:
+        """Return the trainer associated with this agent."""
+        if self._trainer_ref is None:
+            raise ValueError("Trainer has not been set for this agent.")
+        trainer = self._trainer_ref()
+        if trainer is None:
+            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
+        return trainer
+
+    @property
+    def trainer(self) -> Trainer:
+        """Return the trainer associated with this agent."""
+        return self.get_trainer()
+
+    def get_tracer(self) -> Tracer:
+        """Return the tracer configured for this agent."""
+        if hasattr(self.runner, "tracer"):
+            return self.runner.tracer  # type: ignore
+        else:
+            return self.trainer.tracer
+
+    @property
+    def tracer(self) -> Tracer:
+        """Return the tracer configured for this agent."""
+        return self.get_tracer()
+
+    def set_runner(self, runner: Runner[T]) -> None:
+        """Attach the runner responsible for executing rollouts.
+
+        Args:
+            runner: [`Runner`][agentlightning.Runner] coordinating execution.
+        """
+        self._runner_ref = weakref.ref(runner)
+
+    def get_runner(self) -> Runner[T]:
+        """Return the runner responsible for executing rollouts."""
+        if self._runner_ref is None:
+            raise ValueError("Runner has not been set for this agent.")
+        runner = self._runner_ref()
+        if runner is None:
+            raise ValueError("Runner reference is no longer valid (object has been garbage collected).")
+        return runner
+
+    @property
+    def runner(self) -> Runner[T]:
+        """Return the runner responsible for executing rollouts."""
+        return self.get_runner()
+
+    def on_rollout_start(self, task: Task, runner: Runner[T], tracer: Tracer) -> None:
+        """Hook invoked immediately before a rollout begins.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource setup. The default implementation is a no-op.
+
+        Args:
+            task: [`Task`][agentlightning.Task] that will be processed.
+            runner: [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.
+
+        !!! warning "Deprecated"
+            Override [`Hook.on_rollout_start`][agentlightning.Hook.on_rollout_start]
+            instead of this method when extending agents.
+        """
+
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: Runner[T], tracer: Tracer) -> None:
+        """Hook invoked after a rollout completes.
+
+        Subclasses can override this method for cleanup or additional logging. The default
+        implementation is a no-op.
+
+        Args:
+            task: [`Task`][agentlightning.Task] that was processed.
+            rollout: Resulting [`Rollout`][agentlightning.Rollout].
+            runner: [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.
+
+        !!! warning "Deprecated"
+            Override [`Hook.on_rollout_end`][agentlightning.Hook.on_rollout_end]
+            instead of this method when extending agents.
+        """
+
+    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Execute a rollout synchronously.
+
+
+        If you don't wish to implement both training rollout and validation
+        rollout separately, you can just implement `rollout` which will work for both.
+
+        Args:
+            task: Task payload provided by the scheduler.
+            resources: Mapping of named resources (for example LLMs or prompt templates).
+            rollout: Rollout metadata. Avoid mutating this object directly unless a
+                subclass needs to override defaults.
+
+        Returns:
+            One of the following values:
+
+            * `None` when tracing is handled by the runner.
+            * `float` representing the final reward.
+            * `List[ReadableSpan]` with OpenTelemetry spans.
+            * `List[Span]` with Agent Lightning spans.
+        """
+        raise NotImplementedError("Agents must implement the `rollout` method.")
+
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Execute a rollout asynchronously.
+
+        Args:
+            task: Task payload provided by the scheduler.
+            resources: Mapping of named resources (for example LLMs or prompt templates).
+            rollout: Rollout metadata. Avoid mutating this object directly unless a
+                subclass needs to override defaults.
+
+        Returns:
+            Same possible return values as
+            [`rollout`][agentlightning.LitAgent.rollout].
+        """
+        raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")
+
+    def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Process a single training task synchronously.
+
+        By default, this method delegates to
+        [`rollout`][agentlightning.LitAgent.rollout].
+        """
+        return self.rollout(task, resources, rollout)
+
+    def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Process a single validation task synchronously.
+
+        Override this method when validation should differ from training. The default
+        implementation delegates to
+        [`training_rollout`][agentlightning.LitAgent.training_rollout].
+        """
+        return self.rollout(task, resources, rollout)
+
+    async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Process a single training task asynchronously.
+
+        By default, this method delegates to
+        [`rollout_async`][agentlightning.LitAgent.rollout_async].
+        """
+        return await self.rollout_async(task, resources, rollout)
+
+    async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Process a single validation task asynchronously.
+
+        Override this method when validation should differ from training. The default
+        implementation delegates to
+        [`training_rollout_async`][agentlightning.LitAgent.training_rollout_async].
+        """
+        return await self.rollout_async(task, resources, rollout)
@@ -0,0 +1,913 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import ast
+import asyncio
+import logging
+import os
+import re
+import tempfile
+import threading
+from contextlib import asynccontextmanager
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    TypedDict,
+    Union,
+    cast,
+)
+
+import litellm
+import opentelemetry.trace as trace_api
+import yaml
+from fastapi import Request, Response
+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 opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from agentlightning.types import LLM, ProxyLLM
+from agentlightning.utils.server_launcher import (
+    LaunchMode,
+    PythonServerLauncher,
+    PythonServerLauncherArgs,
+    noop_context,
+)
+
+from .store.base import LightningStore
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "LLMProxy",
+]
+
+
+class ModelConfig(TypedDict):
+    """LiteLLM model registration entry.
+
+    This mirrors the items in LiteLLM's `model_list` section.
+
+    Attributes:
+        model_name: Logical model name exposed by the proxy.
+        litellm_params: Parameters passed to LiteLLM for this model
+            (e.g., backend model id, api_base, additional options).
+    """  # Google style kept concise.
+
+    model_name: str
+    litellm_params: Dict[str, Any]
+
+
+def _get_pre_call_data(args: Any, kwargs: Any) -> Dict[str, Any]:
+    """Extract LiteLLM request payload from hook args.
+
+    The LiteLLM logger hooks receive `(*args, **kwargs)` whose third positional
+    argument or `data=` kwarg contains the request payload.
+
+    Args:
+        args: Positional arguments from the hook.
+        kwargs: Keyword arguments from the hook.
+
+    Returns:
+        The request payload dict.
+
+    Raises:
+        ValueError: If the payload cannot be located or is not a dict.
+    """
+    if kwargs.get("data"):
+        data = kwargs["data"]
+    elif len(args) >= 3:
+        data = args[2]
+    else:
+        raise ValueError(f"Unable to get request data from args or kwargs: {args}, {kwargs}")
+    if not isinstance(data, dict):
+        raise ValueError(f"Request data is not a dictionary: {data}")
+    return cast(Dict[str, Any], data)
+
+
+def _reset_litellm_logging_worker() -> None:
+    """Reset LiteLLM's global logging worker to the current event loop.
+
+    LiteLLM keeps a module-level ``GLOBAL_LOGGING_WORKER`` singleton that owns an
+    ``asyncio.Queue``. The queue is bound to the event loop where it was created.
+    When the proxy is restarted, Uvicorn spins up a brand new event loop in a new
+    thread. If the existing logging worker (and its queue) are reused, LiteLLM
+    raises ``RuntimeError: <Queue ...> is bound to a different event loop`` the
+    next time it tries to log. Recreating the worker ensures that LiteLLM will
+    lazily initialise a fresh queue on the new loop.
+    """
+
+    # ``GLOBAL_LOGGING_WORKER`` is imported in a few LiteLLM modules at runtime.
+    # Update any already-imported references so future calls use the fresh worker.
+    try:
+        import litellm.utils as litellm_utils
+        from litellm.litellm_core_utils import logging_worker as litellm_logging_worker
+
+        litellm_logging_worker.GLOBAL_LOGGING_WORKER = litellm_logging_worker.LoggingWorker()
+        litellm_utils.GLOBAL_LOGGING_WORKER = litellm_logging_worker.GLOBAL_LOGGING_WORKER  # type: ignore[reportAttributeAccessIssue]
+    except Exception:  # pragma: no cover - best-effort hygiene
+        logger.warning("Unable to propagate LiteLLM logging worker reset.", exc_info=True)
+
+
+def _reset_litellm_logging_callback_manager() -> None:
+    """Reset LiteLLM's global callback manager.
+
+    To get rid of the warning message: "Cannot add callback - would exceed MAX_CALLBACKS limit of 30."
+    when litellm is restarted multiple times in the same process.
+
+    It does not respect existing input/output callbacks.
+    """
+
+    try:
+        litellm.logging_callback_manager._reset_all_callbacks()  # pyright: ignore[reportPrivateUsage]
+    except Exception:  # pragma: no cover - best-effort hygiene
+        logger.warning("Unable to reset LiteLLM logging callback manager.", exc_info=True)
+
+
+class AddReturnTokenIds(CustomLogger):
+    """LiteLLM logger hook to request token ids from vLLM.
+
+    This mutates the outgoing request payload to include `return_token_ids=True`
+    for backends that support token id return (e.g., vLLM).
+
+    See also:
+        [vLLM PR #22587](https://github.com/vllm-project/vllm/pull/22587)
+    """
+
+    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
+        """Async pre-call hook to adjust request payload.
+
+        Args:
+            args: Positional args from LiteLLM.
+            kwargs: Keyword args from LiteLLM.
+
+        Returns:
+            Either an updated payload dict or an Exception to short-circuit.
+        """
+        try:
+            data = _get_pre_call_data(args, kwargs)
+        except Exception as e:
+            return e
+
+        # Ensure token ids are requested from the backend when supported.
+        return {**data, "return_token_ids": True}
+
+
+class LightningSpanExporter(SpanExporter):
+    """Buffered OTEL span exporter with subtree flushing and training-store sink.
+
+    Design:
+
+    * Spans are buffered until a root span's entire subtree is available.
+    * A private event loop on a daemon thread runs async flush logic.
+    * Rollout/attempt/sequence metadata is reconstructed by merging headers
+      from any span within a subtree.
+
+    Thread-safety:
+
+    * Buffer access is protected by a re-entrant lock.
+    * Export is synchronous to the caller yet schedules an async flush on the
+      internal loop, then waits for completion.
+    """
+
+    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._loop_lock_pid: Optional[int] = None
+
+        # Single dedicated event loop running in a daemon thread.
+        # This decouples OTEL SDK threads from our async store I/O.
+        # Deferred creation until first use.
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[threading.Thread] = None
+
+    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
+        """Lazily initialize the event loop and thread on first use.
+
+        Returns:
+            asyncio.AbstractEventLoop: The initialized event loop.
+        """
+        self._clear_loop_and_lock()
+        if self._loop is None:
+            self._loop = asyncio.new_event_loop()
+            self._loop_thread = threading.Thread(target=self._run_loop, name="LightningSpanExporterLoop", daemon=True)
+            self._loop_thread.start()
+        return self._loop
+
+    def _ensure_lock(self) -> threading.RLock:
+        """Lazily initialize the lock on first use.
+
+        Returns:
+            threading.RLock: The initialized lock.
+        """
+        self._clear_loop_and_lock()
+        if self._lock is None:
+            self._lock = threading.RLock()
+        return self._lock
+
+    def _clear_loop_and_lock(self) -> None:
+        """Clear the loop and lock.
+        This happens if the exporter was used in a process then used in another process.
+
+        This should only happen in CI.
+        """
+        if os.getpid() != self._loop_lock_pid:
+            logger.warning("Loop and lock are not owned by the current process. Clearing them.")
+            self._loop = None
+            self._loop_thread = None
+            self._lock = None
+            self._loop_lock_pid = os.getpid()
+        elif self._loop_lock_pid is None:
+            self._loop_lock_pid = os.getpid()
+
+    def _run_loop(self) -> None:
+        """Run the private asyncio loop forever on the exporter thread."""
+        assert self._loop is not None, "Loop should be initialized before thread starts"
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def shutdown(self) -> None:
+        """Shut down the exporter event loop.
+
+        Safe to call at process exit.
+
+        """
+        if self._loop is None:
+            return
+
+        try:
+
+            def _stop():
+                assert self._loop is not None
+                self._loop.stop()
+
+            self._loop.call_soon_threadsafe(_stop)
+            if self._loop_thread is not None:
+                self._loop_thread.join(timeout=2.0)
+            self._loop.close()
+        except Exception:
+            logger.exception("Error during exporter shutdown")
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Export spans via buffered subtree flush.
+
+        Appends spans to the internal buffer, then triggers an async flush on the
+        private event loop. Blocks until that flush completes.
+
+        Args:
+            spans: Sequence of spans to export.
+
+        Returns:
+            SpanExportResult: SUCCESS on flush success, else FAILURE.
+        """
+        # Buffer append under lock to protect against concurrent exporters.
+        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
+
+        return SpanExportResult.SUCCESS
+
+    async def _maybe_flush(self):
+        """Flush ready subtrees from the buffer.
+
+        Strategy:
+            We consider a subtree "ready" if we can identify a root span. We
+            then take that root and all its descendants out of the buffer and
+            try to reconstruct rollout/attempt/sequence headers by merging any
+            span's `metadata.requester_custom_headers` within the subtree.
+
+        Required headers:
+            `x-rollout-id` (str), `x-attempt-id` (str), `x-sequence-id` (str of int)
+
+        Raises:
+            None directly. Logs and skips malformed spans.
+
+        """
+        # Iterate over current roots. Each iteration pops a whole subtree.
+        for root_span_id in self._get_root_span_ids():
+            subtree_spans = self._pop_subtrees(root_span_id)
+            if not subtree_spans:
+                continue
+
+            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
+
+            # Merge all custom headers found in the subtree.
+            headers_merged: Dict[str, Any] = {}
+
+            for span in subtree_spans:
+                if span.attributes is None:
+                    continue
+                headers_str = span.attributes.get("metadata.requester_custom_headers")
+                if headers_str is None:
+                    continue
+                if not isinstance(headers_str, str):
+                    logger.error(
+                        f"metadata.requester_custom_headers is not stored as a string: {headers_str}. Skipping the span."
+                    )
+                    continue
+                if not headers_str.strip():
+                    logger.warning("metadata.requester_custom_headers is an empty string. Skipping the span.")
+                    continue
+                try:
+                    # Use literal_eval to parse the stringified dict safely.
+                    headers = ast.literal_eval(headers_str)
+                except Exception as e:
+                    logger.error(
+                        f"Failed to parse metadata.requester_custom_headers: {headers_str}, error: {e}. Skipping the span."
+                    )
+                    continue
+                if not isinstance(headers, dict):
+                    logger.error(
+                        f"metadata.requester_custom_headers is not parsed as a dict: {headers}. Skipping the span."
+                    )
+                    continue
+                headers_merged.update(cast(Dict[str, Any], headers))
+
+            if not headers_merged:
+                logger.warning(f"No headers found in {len(subtree_spans)} subtree spans. Cannot log to store.")
+                continue
+
+            # Validate and normalize required header fields.
+            rollout_id = headers_merged.get("x-rollout-id")
+            attempt_id = headers_merged.get("x-attempt-id")
+            sequence_id = headers_merged.get("x-sequence-id")
+            if not rollout_id or not attempt_id or not sequence_id or not sequence_id.isdigit():
+                logger.warning(
+                    f"Missing or invalid rollout_id, attempt_id, or sequence_id in headers: {headers_merged}. Cannot log to store."
+                )
+                continue
+            if not isinstance(rollout_id, str) or not isinstance(attempt_id, str):
+                logger.warning(
+                    f"rollout_id or attempt_id is not a string: {rollout_id}, {attempt_id}. Cannot log to store."
+                )
+                continue
+            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
+                )
+
+    def _get_root_span_ids(self) -> Iterable[int]:
+        """Yield span_ids for root spans currently in the buffer.
+
+        A root span is defined as one with `parent is None`.
+
+        Yields:
+            int: Span id for each root span found.
+        """
+        for span in self._buffer:
+            if span.parent is None:
+                span_context = span.get_span_context()
+                if span_context is not None:
+                    yield span_context.span_id
+
+    def _get_subtrees(self, root_span_id: int) -> Iterable[int]:
+        """Yield span_ids in the subtree rooted at `root_span_id`.
+
+        Depth-first traversal over the current buffer.
+
+        Args:
+            root_span_id: The span id of the root.
+
+        Yields:
+            int: Span ids including the root and all descendants found.
+        """
+        # Yield the root span id first.
+        yield root_span_id
+        for span in self._buffer:
+            # Check whether the span's parent is the root_span_id.
+            if span.parent is not None and span.parent.span_id == root_span_id:
+                span_context = span.get_span_context()
+                if span_context is not None:
+                    # Recursively get child spans.
+                    yield from self._get_subtrees(span_context.span_id)
+
+    def _pop_subtrees(self, root_span_id: int) -> List[ReadableSpan]:
+        """Remove and return the subtree for a particular root from the buffer.
+
+        Args:
+            root_span_id: Root span id identifying the subtree.
+
+        Returns:
+            list[ReadableSpan]: Spans that were part of the subtree. Order follows buffer order.
+        """
+        subtree_span_ids = set(self._get_subtrees(root_span_id))
+        subtree_spans: List[ReadableSpan] = []
+        new_buffer: List[ReadableSpan] = []
+        for span in self._buffer:
+            span_context = span.get_span_context()
+            if span_context is not None and span_context.span_id in subtree_span_ids:
+                subtree_spans.append(span)
+            else:
+                new_buffer.append(span)
+        # Replace buffer with remaining spans to avoid re-processing.
+        self._buffer = new_buffer
+        return subtree_spans
+
+
+class LightningOpenTelemetry(OpenTelemetry):
+    """OpenTelemetry integration that exports spans to the Lightning store.
+
+    Responsibilities:
+
+    * Ensures each request is annotated with a per-attempt sequence id so spans
+      are ordered deterministically even with clock skew across nodes.
+    * Uses [`LightningSpanExporter`][agentlightning.llm_proxy.LightningSpanExporter] to persist spans for analytics and training.
+    """
+
+    def __init__(self):
+        config = OpenTelemetryConfig(exporter=LightningSpanExporter())
+
+        # Check for tracer initialization
+        if _check_tracer_provider():
+            logger.error("Tracer is already initialized. OpenTelemetry may not work as expected.")
+
+        super().__init__(config=config)  # pyright: ignore[reportUnknownMemberType]
+
+
+class RolloutAttemptMiddleware(BaseHTTPMiddleware):
+    """
+    Rewrites /rollout/{rid}/attempt/{aid}/... -> /...
+    and injects x-rollout-id, x-attempt-id, x-sequence-id headers.
+
+    LLMProxy can update store later without rebuilding middleware.
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        # Decode rollout and attempt from the URL prefix. Example:
+        #   /rollout/r123/attempt/a456/v1/chat/completions
+        # becomes
+        #   /v1/chat/completions
+        # while adding request-scoped headers for trace attribution.
+        path = request.url.path
+
+        match = re.match(r"^/rollout/([^/]+)/attempt/([^/]+)(/.*)?$", path)
+        if match:
+            rollout_id = match.group(1)
+            attempt_id = match.group(2)
+            new_path = match.group(3) if match.group(3) is not None else "/"
+
+            # Rewrite the ASGI scope path so downstream sees a clean OpenAI path.
+            request.scope["path"] = new_path
+            request.scope["raw_path"] = new_path.encode()
+
+            store = get_active_llm_proxy().get_store()
+            if store is not None:
+                # Allocate a monotonic sequence id per (rollout, attempt).
+                sequence_id = await store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+                # Inject headers so downstream components and exporters can retrieve them.
+                request.scope["headers"] = list(request.scope["headers"]) + [
+                    (b"x-rollout-id", rollout_id.encode()),
+                    (b"x-attempt-id", attempt_id.encode()),
+                    (b"x-sequence-id", str(sequence_id).encode()),
+                ]
+            else:
+                logger.warning("Store is not set. Skipping sequence id allocation and header injection.")
+
+        response = await call_next(request)
+        return response
+
+
+class LLMProxy:
+    """Host a LiteLLM OpenAI-compatible proxy bound to a LightningStore.
+
+    The proxy:
+
+    * Serves an OpenAI-compatible API via uvicorn.
+    * Adds rollout/attempt routing and headers via middleware.
+    * Registers OTEL export and token-id callbacks.
+    * Writes a LiteLLM worker config file with `model_list` and settings.
+
+    Lifecycle:
+
+    * [`start()`][agentlightning.LLMProxy.start] writes config, starts uvicorn server in a thread, and waits until ready.
+    * [`stop()`][agentlightning.LLMProxy.stop] tears down the server and removes the temp config file.
+    * [`restart()`][agentlightning.LLMProxy.restart] convenience wrapper to stop then start.
+
+    !!! note
+
+        As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
+        process from the main runner (i.e., tracer from agents). See `launch_mode` for how to change that.
+
+    !!! warning
+
+        The LLM Proxy does support streaming, but the tracing is still problematic when streaming is enabled.
+
+    !!! danger
+
+        Do not run LLM proxy in the same process as the main runner. It's easy to cause conflicts in the tracer provider
+        with tracers like [`AgentOpsTracer`][agentlightning.AgentOpsTracer].
+
+    Args:
+        port: TCP port to bind. Will bind to a random port if not provided.
+        model_list: LiteLLM `model_list` entries.
+        store: LightningStore used for span sequence and persistence.
+        host: Publicly reachable host used in resource endpoints. See `host` of `launcher_args` for more details.
+        litellm_config: Extra LiteLLM proxy config merged with `model_list`.
+        num_retries: Default LiteLLM retry count injected into `litellm_settings`.
+        num_workers: Number of workers to run in the server. Only applicable for "mp" launch mode. Ignored if launcher_args is provided.
+            When `num_workers > 1`, the server will be run using [gunicorn](https://gunicorn.org/).
+        launch_mode: Launch mode for the server. Defaults to "mp". Cannot be used together with launcher_args. Ignored if launcher_args is provided.
+            It's recommended to use `launch_mode="mp"` to launch the proxy, which will launch the server in a separate process.
+            `launch_mode="thread"` can also be used if used in caution. It will launch the server in a separate thread.
+            `launch_mode="asyncio"` launches the server in the current thread as an asyncio task.
+            It is NOT recommended because it often causes hanging requests. Only use it if you know what you are doing.
+        launcher_args: Arguments for the server launcher. If this is provided, host, port, and launch_mode will be ignored. Cannot be used together with port, host, and launch_mode.
+    """
+
+    def __init__(
+        self,
+        port: int | None = None,
+        model_list: List[ModelConfig] | None = None,
+        store: Optional[LightningStore] = None,
+        host: str | None = None,
+        litellm_config: Dict[str, Any] | None = None,
+        num_retries: int = 0,
+        num_workers: int = 1,
+        launch_mode: LaunchMode = "mp",
+        launcher_args: PythonServerLauncherArgs | None = None,
+        _add_return_token_ids: bool = True,
+    ):
+        self.store = store
+
+        if launcher_args is not None and (
+            port is not None or host is not None or launch_mode != "mp" or num_workers != 1
+        ):
+            raise ValueError("port, host, launch_mode, and num_workers cannot be set when launcher_args is provided.")
+
+        self.server_launcher_args = launcher_args or PythonServerLauncherArgs(
+            port=port,
+            host=host,
+            launch_mode=launch_mode,
+            n_workers=num_workers,
+            # NOTE: This /health endpoint can be slow sometimes because it actually probes the backend LLM service.
+            healthcheck_url="/health",
+            startup_timeout=60.0,
+        )
+
+        if self.server_launcher_args.healthcheck_url is None:
+            logger.warning("healthcheck_url is not set. LLM Proxy will not be checked for healthiness after starting.")
+
+        self.model_list = model_list or []
+        self.litellm_config = litellm_config or {}
+
+        # Ensure num_retries is present inside the litellm_settings block.
+        self.litellm_config.setdefault("litellm_settings", {})
+        self.litellm_config["litellm_settings"].setdefault("num_retries", num_retries)
+        self.server_launcher = PythonServerLauncher(app, self.server_launcher_args, noop_context())
+
+        self._config_file = None
+
+        self._add_return_token_ids = _add_return_token_ids
+
+    def get_store(self) -> Optional[LightningStore]:
+        """Get the store used by the proxy.
+
+        Returns:
+            The store used by the proxy.
+        """
+        return self.store
+
+    def set_store(self, store: LightningStore) -> None:
+        """Set the store for the proxy.
+
+        Args:
+            store: The store to use for the proxy.
+        """
+        self.store = store
+
+    def update_model_list(self, model_list: List[ModelConfig]) -> None:
+        """Replace the in-memory model list.
+
+        Args:
+            model_list: New list of model entries.
+        """
+        self.model_list = model_list
+        logger.info(f"Updating LLMProxy model list to: {model_list}")
+        # Do nothing if the server is not running.
+
+    def initialize(self):
+        """Initialize global middleware and LiteLLM callbacks.
+
+        Installs:
+
+        * A FastAPI middleware that rewrites /rollout/{rid}/attempt/{aid}/... paths,
+        injects rollout/attempt/sequence headers, and forwards downstream.
+        * LiteLLM callbacks for token ids and OpenTelemetry export.
+
+        The middleware can only be installed once because once the FastAPI app has started,
+        the middleware cannot be changed any more.
+
+        This function does not start any server. It only wires global hooks.
+        """
+        if self.store is None:
+            raise ValueError("Store is not set. Please set the store before initializing the LLMProxy.")
+
+        if _global_llm_proxy is not None:
+            logger.warning("A global LLMProxy is already set. Overwriting it with the new instance.")
+
+        # Set the global LLMProxy reference for middleware/exporter access.
+        set_active_llm_proxy(self)
+
+        # Install middleware if it's not already installed.
+        installed: bool = False
+        for mw in app.user_middleware:
+            if mw.cls is RolloutAttemptMiddleware:
+                # Check whether the middleware is installed.
+                # It could be installed by other LLM Proxy instances, but it doesn't matter.
+                logger.info("Found existing RolloutAttemptMiddleware installed. Will not install a new one.")
+                installed = True
+                break
+
+        if not installed:
+            # Fallback to adding a new middleware
+            logger.info("Adding a new middleware to the FastAPI app.")
+            app.add_middleware(RolloutAttemptMiddleware)
+
+        if not initialize_llm_callbacks(self._add_return_token_ids):
+            # If it's not the first time to initialize the callbacks, also
+            # reset LiteLLM's logging worker so its asyncio.Queue binds to the new loop.
+            _reset_litellm_logging_worker()
+
+    @asynccontextmanager
+    async def _serve_context(self) -> AsyncGenerator[None, None]:
+        """Context manager to serve the proxy server.
+
+        See [`start`][agentlightning.LLMProxy.start] and [`stop`][agentlightning.LLMProxy.stop] for more details.
+        """
+
+        if not self.store:
+            raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
+
+        # Initialize global middleware and callbacks.
+        self.initialize()
+
+        # Persist a temp worker config for LiteLLM and point the proxy at it.
+        self._config_file = tempfile.NamedTemporaryFile(suffix=".yaml", delete=False).name
+        with open(self._config_file, "w") as fp:
+            yaml.safe_dump(
+                {
+                    "model_list": self.model_list,
+                    **self.litellm_config,
+                },
+                fp,
+            )
+
+        save_worker_config(config=self._config_file)
+
+        # NOTE: When running the _serve_context in current process, you might encounter the following problems:
+        # Problem 1: in litellm worker, <Queue at 0x70f1d028cd90 maxsize=50000> is bound to a different event loop
+        # Problem 2: Proxy has conflicted opentelemetry setup with the main process.
+
+        # Ready
+        logger.info("LLMProxy preparation is done. Will start the server.")
+        yield
+
+        # Clean up
+
+        logger.info("LLMProxy server is cleaning up.")
+
+        # Remove worker config to avoid stale references.
+        if self._config_file and os.path.exists(self._config_file):
+            os.unlink(self._config_file)
+
+        logger.info("LLMProxy server finishes.")
+
+    async def start(self):
+        """Start the proxy server thread and initialize global wiring.
+
+        Side effects:
+
+        * Sets the module-level global store for middleware/exporter access.
+        * Calls `initialize()` once to register middleware and callbacks.
+        * Writes a temporary YAML config consumed by LiteLLM worker.
+        * Launches uvicorn in a daemon thread and waits for readiness.
+        """
+        # Refresh the serve context
+        self.server_launcher.serve_context = self._serve_context()
+
+        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"]:
+            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"]:
+            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"]:
+            raise RuntimeError("The store is not async-safe. Please use another store.")
+
+        logger.info(
+            f"Starting LLMProxy server in {self.server_launcher.args.launch_mode} mode with store capabilities: {store_capabilities}"
+        )
+
+        await self.server_launcher.start()
+
+    async def stop(self):
+        """Stop the proxy server and clean up temporary artifacts.
+
+        This is a best-effort graceful shutdown with a bounded join timeout.
+        """
+        if not self.is_running():
+            logger.warning("LLMProxy is not running. Nothing to stop.")
+            return
+
+        await self.server_launcher.stop()
+
+    async def restart(self, *, _port: int | None = None) -> None:
+        """Restart the proxy if running, else start it.
+
+        Convenience wrapper calling `stop()` followed by `start()`.
+        """
+        logger.info("Restarting LLMProxy server...")
+        if self.is_running():
+            await self.stop()
+        if _port is not None:
+            self.server_launcher_args.port = _port
+        await self.start()
+
+    def is_running(self) -> bool:
+        """Return whether the uvicorn server is active.
+
+        Returns:
+            bool: True if server was started and did not signal exit.
+        """
+        return self.server_launcher.is_running()
+
+    def as_resource(
+        self,
+        rollout_id: str | None = None,
+        attempt_id: str | None = None,
+        model: str | None = None,
+        sampling_parameters: Dict[str, Any] | None = None,
+    ) -> LLM:
+        """Create an `LLM` resource pointing at this proxy with rollout context.
+
+        The returned endpoint is:
+            `http://{host}:{port}/rollout/{rollout_id}/attempt/{attempt_id}`
+
+        Args:
+            rollout_id: Rollout identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
+            attempt_id: Attempt identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
+            model: Logical model name to use. If omitted and exactly one model
+                is configured or all models have the same name, that model is used.
+            sampling_parameters: Optional default sampling parameters.
+
+        Returns:
+            LLM: Configured resource ready for OpenAI-compatible calls.
+
+        Raises:
+            ValueError: If `model` is omitted and zero or multiple models are configured.
+        """
+        if model is None:
+            if len(self.model_list) == 1:
+                model = self.model_list[0]["model_name"]
+            elif len(self.model_list) == 0:
+                raise ValueError("No models found in model_list. Please specify the model.")
+            else:
+                first_model_name = self.model_list[0]["model_name"]
+                if all(model_config["model_name"] == first_model_name for model_config in self.model_list):
+                    model = first_model_name
+                else:
+                    raise ValueError(
+                        f"Multiple models found in model_list: {self.model_list}. Please specify the model."
+                    )
+
+        if rollout_id is None and attempt_id is None:
+            return ProxyLLM(
+                endpoint=self.server_launcher.access_endpoint,
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        elif rollout_id is not None and attempt_id is not None:
+            return LLM(
+                endpoint=f"{self.server_launcher.access_endpoint}/rollout/{rollout_id}/attempt/{attempt_id}",
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        else:
+            raise ValueError("Either rollout_id and attempt_id must be provided, or neither.")
+
+
+_global_llm_proxy: Optional[LLMProxy] = None
+_callbacks_before_litellm_start: Optional[List[Any]] = None
+
+
+def get_active_llm_proxy() -> LLMProxy:
+    """Get the current global LLMProxy instance.
+
+    Returns:
+        Optional[LLMProxy]: The current LLMProxy if set, else None.
+    """
+    if _global_llm_proxy is None:
+        raise ValueError("Global LLMProxy is not set. Please call llm_proxy.start() first.")
+    return _global_llm_proxy
+
+
+def set_active_llm_proxy(proxy: LLMProxy) -> None:
+    """Set the current global LLMProxy instance.
+
+    Args:
+        proxy: The LLMProxy instance to set as global.
+    """
+    global _global_llm_proxy
+    _global_llm_proxy = proxy
+
+
+def initialize_llm_callbacks(_add_return_token_ids: bool = True) -> bool:
+    """Restore `litellm.callbacks` to a state that is just initialized by agent-lightning.
+
+    When litellm is restarted multiple times in the same process, more and more callbacks
+    will be appended to `litellm.callbacks`, which may exceed the MAX_CALLBACKS limit.
+    This function remembers the initial state of `litellm.callbacks` and always restore to that state.
+
+    Args:
+        _add_return_token_ids: Whether to add the return token ids callback. Internal use only.
+        Ideally the callback should automatically be enabled when the backend supports it.
+
+    Returns:
+        Whether the callbacks are initialized for the first time.
+    """
+    global _callbacks_before_litellm_start
+
+    if _callbacks_before_litellm_start is None:
+        litellm.callbacks.extend(  # type: ignore
+            [
+                AddReturnTokenIds(),
+                LightningOpenTelemetry(),
+            ]
+            if _add_return_token_ids
+            else [
+                LightningOpenTelemetry(),
+            ]
+        )
+        _callbacks_before_litellm_start = [*litellm.callbacks]  # type: ignore
+        return True
+
+    _reset_litellm_logging_callback_manager()
+
+    # Check if tracer provider is malformed due to global tracer clear in tests.
+    if not _check_tracer_provider():
+        logger.warning(
+            "Global tracer provider might have been cleared outside. Re-initializing OpenTelemetry callback."
+        )
+        _callbacks_before_litellm_start = [
+            cb for cb in _callbacks_before_litellm_start if not isinstance(cb, LightningOpenTelemetry)
+        ] + [LightningOpenTelemetry()]
+    else:
+        logger.debug("Global tracer provider is valid. Reusing existing OpenTelemetry callback.")
+
+    litellm.callbacks.clear()  # type: ignore
+    litellm.callbacks.extend(_callbacks_before_litellm_start)  # type: ignore
+    return False
+
+
+def _check_tracer_provider() -> bool:
+    """Check if the global tracer provider is properly initialized.
+
+    We don't guarantee the tracer provider is our tracer provider.
+
+    Returns:
+        bool: True if the tracer provider is valid, else False.
+    """
+    if (
+        hasattr(trace_api, "_TRACER_PROVIDER")
+        and trace_api._TRACER_PROVIDER is not None  # pyright: ignore[reportPrivateUsage]
+    ):
+        return True
+    return False
@@ -1,7 +1,45 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 import logging
+import os
+import platform
+
+__all__ = ["configure_logger"]


 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
+    """Create or reset a namespaced logger with a consistent console format.
+
+    This helper clears any previously attached handlers before binding a single
+    `StreamHandler` that writes to standard output. The resulting logger does
+    not propagate to the root logger, preventing duplicate log emission when
+    applications compose multiple logging configurations.
+
+    Args:
+        level: Logging level applied both to the logger and the installed
+            handler. Defaults to `logging.INFO`.
+        name: Dotted path for the logger instance. Defaults to
+            `"agentlightning"`.
+
+    Returns:
+        Configured logger instance ready for immediate use.
+
+    Examples:
+        ```python
+        from agentlightning import configure_logger
+
+        logger = configure_logger(level=logging.INFO)
+        logger.info("agent-lightning is ready!")
+        ```
+    """
+
+    # Ensure UTF-8 encoding on Windows consoles
+    # Note: This change does not fully represent support for execution under the windown 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

@@ -1,66 +1,7 @@
-import asyncio
-import inspect
+# Copyright (c) Microsoft. All rights reserved.
+
 import warnings
-from typing import TypedDict, Optional

-from agentops.sdk.decorators import operation
+from .emitter.reward import *  # noqa: F401,F403

-
-class RewardSpanData(TypedDict):
-    type: "reward"
-    value: Optional[float]
-
-
-def reward(fn: callable) -> callable:
-    """
-    A decorator to wrap a function that computes rewards.
-    It will automatically handle the input and output of the function.
-    """
-
-    def wrap_result(result: Optional[float]) -> RewardSpanData:
-        """
-        Wrap the result of the function in a dict.
-        """
-        if result is None:
-            return {"type": "reward", "value": None}
-        if not isinstance(result, (float, int)):
-            warnings.warn(f"Reward is ignored because it is not a number: {result}")
-            return {"type": "reward", "value": None}
-        return {"type": "reward", "value": float(result)}
-
-    # Check if the function is async
-    is_async = asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn)
-
-    if is_async:
-
-        async def wrapper_async(*args, **kwargs):
-            result: Optional[float] = None
-
-            @operation
-            async def agentops_reward_operation() -> RewardSpanData:
-                # The reward function we are interested in tracing
-                # It takes zero inputs and return a formatted dict
-                nonlocal result
-                result = await fn(*args, **kwargs)
-                return wrap_result(result)
-
-            await agentops_reward_operation()
-            return result
-
-        return wrapper_async
-
-    else:
-
-        def wrapper(*args, **kwargs):
-            result: Optional[float] = None
-
-            @operation
-            def agentops_reward_operation() -> RewardSpanData:
-                nonlocal result
-                result = fn(*args, **kwargs)
-                return wrap_result(result)
-
-            agentops_reward_operation()
-            return result
-
-        return wrapper
+warnings.warn("agentlightning.reward is deprecated. Please use agentlightning.emitter instead.")
@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .agent import LitAgentRunner
+from .base import Runner
+from .legacy import LegacyAgentRunner
+
+__all__ = [
+    "Runner",
+    "LegacyAgentRunner",
+    "LitAgentRunner",
+]
@@ -0,0 +1,534 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent runner implementation for executing agent rollouts.
+
+This module provides the concrete implementation of the runner interface,
+handling the execution of agent rollouts with support for tracing, hooks,
+and distributed worker coordination.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, TypeVar, cast
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.litagent import LitAgent
+from agentlightning.reward import emit_reward, find_final_reward
+from agentlightning.store.base import LightningStore
+from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.tracer.base import Tracer
+from agentlightning.types import (
+    AttemptedRollout,
+    Hook,
+    NamedResources,
+    Rollout,
+    RolloutMode,
+    RolloutRawResult,
+    Span,
+)
+
+if TYPE_CHECKING:
+    from agentlightning.execution.events import ExecutionEvent
+
+from .base import Runner
+
+T_task = TypeVar("T_task")
+
+logger = logging.getLogger(__name__)
+
+
+class LitAgentRunner(Runner[T_task]):
+    """Execute [`LitAgent`][agentlightning.LitAgent] tasks with tracing support.
+
+    This runner manages the complete lifecycle of agent rollout execution,
+    including task polling, resource management, tracing, and hooks. It supports
+    both continuous iteration over tasks from the store and single-step execution.
+
+    Attributes:
+        worker_id: Identifier for the active worker process, if any.
+    """
+
+    def __init__(self, tracer: Tracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0) -> None:
+        """Initialize the agent runner.
+
+        Args:
+            tracer: [`Tracer`][agentlightning.Tracer] used for rollout spans.
+            max_rollouts: Optional cap on iterations processed by
+                [`iter`][agentlightning.LitAgentRunner.iter].
+            poll_interval: Seconds to wait between store polls when no work is available.
+        """
+        super().__init__()
+        self._tracer = tracer
+        self._max_rollouts = max_rollouts
+        self._poll_interval = poll_interval
+
+        # Set later
+        self._agent: Optional[LitAgent[T_task]] = None
+        self._hooks: Sequence[Hook] = []
+        self._store: Optional[LightningStore] = None
+        self.worker_id: Optional[int] = None
+
+    def init(self, agent: LitAgent[T_task], *, hooks: Optional[Sequence[Hook]] = None, **kwargs: Any) -> None:
+        """Initialize the runner with the agent.
+
+        This sets up the agent-runner relationship, registers hooks, and
+        initializes the tracer.
+
+        Args:
+            agent: [`LitAgent`][agentlightning.LitAgent] instance executed by the runner.
+            hooks: Optional sequence of [`Hook`][agentlightning.Hook]
+                callbacks invoked around tracing and rollout boundaries.
+            **kwargs: Additional initialization arguments (currently unused).
+        """
+        self._agent = agent
+        self._agent.set_runner(self)
+        self._hooks = [*hooks] if hooks is not None else []
+
+        self._tracer.init()
+
+    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        """Initialize the runner for each worker with worker_id and store.
+
+        This method is called once per worker in a distributed setup to provide
+        the worker with its ID and store connection.
+
+        Args:
+            worker_id: Unique identifier for this worker process.
+            store: [`LightningStore`][agentlightning.LightningStore]
+                used for task coordination and persistence.
+            **kwargs: Additional worker-specific initialization arguments (currently unused).
+        """
+        self._store = store
+        self.worker_id = worker_id
+
+        self._tracer.init_worker(worker_id)
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Teardown the runner and clean up all resources.
+
+        This method resets all internal state including the agent, store,
+        hooks, and worker ID, and calls the tracer's teardown method.
+
+        Args:
+            *args: Additional teardown arguments (currently unused).
+            **kwargs: Additional teardown keyword arguments (currently unused).
+        """
+        self._agent = None
+        self._store = None
+        self.worker_id = None
+        self._hooks = []
+
+        self._tracer.teardown()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Teardown the runner for a specific worker.
+
+        This method cleans up worker-specific resources and resets the worker ID.
+
+        Args:
+            worker_id: Unique identifier of the worker being torn down.
+            *args: Additional teardown arguments (currently unused).
+            **kwargs: Additional teardown keyword arguments (currently unused).
+        """
+        self.worker_id = None
+
+        self._tracer.teardown_worker(worker_id)
+
+    @property
+    def tracer(self) -> Tracer:
+        """Get the tracer instance.
+
+        Returns:
+            The Tracer instance used by this runner.
+        """
+        return self._tracer
+
+    def get_agent(self) -> LitAgent[T_task]:
+        """Get the agent instance.
+
+        Returns:
+            The LitAgent instance managed by this runner.
+
+        Raises:
+            ValueError: If the agent has not been initialized via [`init`][agentlightning.LitAgentRunner.init].
+        """
+        if self._agent is None:
+            raise ValueError("Agent not initialized. Call init() first.")
+        return self._agent
+
+    def get_store(self) -> LightningStore:
+        """Get the store instance.
+
+        Returns:
+            The LightningStore instance for this worker.
+
+        Raises:
+            ValueError: If the store has not been initialized via [`init_worker`][agentlightning.LitAgentRunner.init_worker].
+        """
+        if self._store is None:
+            raise ValueError("Store not initialized. Call init_worker() first.")
+        return self._store
+
+    def get_worker_id(self) -> str:
+        """Get the formatted worker ID string.
+
+        Returns:
+            A formatted string like "Worker-0" if initialized, or "Worker-Unknown"
+            if the worker ID has not been set.
+        """
+        return f"Worker-{self.worker_id}" if self.worker_id is not None else "Worker-Unknown"
+
+    def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
+        """Generate a standardized log prefix for the current worker.
+
+        This creates a consistent prefix format for log messages to identify
+        which worker and rollout the message is associated with.
+
+        Args:
+            rollout_id: Optional rollout ID to include in the prefix.
+
+        Returns:
+            A formatted log prefix string like "[Worker 0 | Rollout xyz]",
+            "[Worker 0]", "[Rollout xyz]", or "[Default Worker]".
+        """
+        if self.worker_id is not None:
+            if rollout_id:
+                return f"[Worker {self.worker_id} | Rollout {rollout_id}]"
+            else:
+                return f"[Worker {self.worker_id}]"
+        if rollout_id:
+            return f"[Rollout {rollout_id}]"
+        return "[Default Worker]"
+
+    async def _trigger_hooks(
+        self,
+        hook_type: Literal["on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end"],
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Trigger all registered hooks of a specific type.
+
+        This method calls the specified hook method on all registered hooks,
+        catching and logging any exceptions that occur during hook execution
+        to prevent them from disrupting the main execution flow.
+
+        Args:
+            hook_type: The type of hook to trigger. Valid values are:
+                "on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end".
+            *args: Positional arguments to pass to the hook methods.
+            **kwargs: Keyword arguments to pass to the hook methods.
+        """
+        for hook in self._hooks:
+            try:
+                await getattr(hook, hook_type)(*args, **kwargs)
+            except Exception:
+                logger.exception(f"{self._log_prefix()} Exception during {hook_type} hook {hook}.")
+
+    async def _post_process_rollout_result(
+        self, rollout: AttemptedRollout, raw_result: RolloutRawResult
+    ) -> List[ReadableSpan] | List[Span]:
+        """Standardizes the agent's return value and report what's needed to report to the store.
+
+        Args:
+            rollout: The rollout object for the current task.
+            raw_result: The output from the agent's rollout method.
+
+        Returns:
+            The spans that are assumed to be added to the store.
+            This only serves as an estimation for logging purposes. For precise tracking, use the store directly.
+        """
+        store = self.get_store()
+
+        trace_spans: list[ReadableSpan] | list[Span] = []
+
+        # Case 0: result is None
+        if raw_result is None:
+            trace_spans = self._tracer.get_last_trace()
+
+        # Case 1: result is a float (final reward)
+        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)
+            await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
+            trace_spans.append(reward_span)
+
+        if isinstance(raw_result, list):
+            # For rollout methods that return a list, we assume that the returned spans
+            # are the complete span set from the whole rollout
+            trace_spans = raw_result
+
+            # Case 2: result is a list of ReadableSpan (OpenTelemetry spans)
+            if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result):
+
+                if not isinstance(
+                    self._tracer, AgentOpsTracer
+                ):  # TODO: this should be replaced with general OpenTelemetry tracer in next version
+                    for span in raw_result:
+                        await store.add_otel_span(
+                            rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
+                        )
+                else:
+                    logger.warning(
+                        f"{self._log_prefix(rollout.rollout_id)} Tracer is already an OpenTelemetry tracer. "
+                        "The traces should have already been added to the store. "
+                        "No need to return anything from rollout."
+                    )
+
+            # Case 3: result is a list of Span (agentlightning spans)
+            elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result):
+                # Add the spans directly to the store
+                for span in raw_result:
+                    await store.add_span(cast(Span, span))
+                trace_spans = raw_result
+
+            # Left over cases for list
+            elif len(raw_result) == 0:
+                logger.warning(
+                    f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. "
+                    "Please check your rollout implementation."
+                )
+                trace_spans = raw_result
+
+            else:
+                types = [type(t).__name__ for t in raw_result][:10]
+                raise ValueError(
+                    f"Invalid raw result type. It's expected to be a list of ReadableSpan or Span, "
+                    f"but got: {', '.join(types)}..."
+                )
+
+        return trace_spans
+
+    async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
+        """Sleep until the next poll interval, with optional event-based interruption.
+
+        If an event is provided, the method will check it periodically (every 0.1s)
+        and return early if the event is set.
+
+        Args:
+            event: Optional [`ExecutionEvent`][agentlightning.ExecutionEvent] object that can be used to interrupt the sleep.
+                If set during the sleep period, the method returns immediately.
+        """
+        if event is None:
+            await asyncio.sleep(self._poll_interval)
+            return
+        current_time = time.time()
+        next_time = current_time + self._poll_interval
+        while time.time() < next_time:
+            await asyncio.sleep(0.1)
+            if event.is_set():
+                return
+
+    async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: bool = False) -> str:
+        """Execute a single rollout implementation.
+
+        This is the core method that handles the execution of a single rollout,
+        including resource fetching, hook triggering, agent invocation, tracing,
+        and result processing.
+
+        Args:
+            next_rollout: The rollout to execute, containing input data, mode,
+                and resources information.
+            raise_on_exception: If True, exceptions during rollout execution will
+                be re-raised. If False, exceptions are logged but not propagated.
+        """
+        store = self.get_store()
+        agent = self.get_agent()
+
+        rollout_id = next_rollout.rollout_id
+
+        resources_id = next_rollout.resources_id
+        resources_update = None
+        if resources_id:
+            resources_update = await store.get_resources_by_id(resources_id)
+        else:
+            logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
+            resources_update = await store.get_latest_resources()
+        if not resources_update:
+            if raise_on_exception:
+                raise RuntimeError(f"{self._log_prefix(rollout_id)} Failed to fetch resources")
+            else:
+                logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
+                return rollout_id
+
+        trace_spans: List[ReadableSpan] | List[Span] = []
+        has_exception: bool = False
+
+        try:
+            await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)
+
+            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
+            ):
+                await self._trigger_hooks(
+                    hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
+                )
+
+                # NOTE: This is the most costly step in the whole function
+                # If the rollout method becomes unresponsive or timeouts, there is nothing we can do within the runner.
+                # We might need some mechanisms in execution strategy to restart the runner. But that's a future work.
+                if agent.is_async():
+                    rollout_method = (
+                        agent.training_rollout_async if next_rollout.mode == "train" else agent.validation_rollout_async
+                    )
+                    result = await rollout_method(
+                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
+                    )
+                else:
+                    rollout_method = (
+                        agent.training_rollout if next_rollout.mode == "train" else agent.validation_rollout
+                    )
+                    result = rollout_method(
+                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
+                    )
+
+                await self._trigger_hooks(
+                    hook_type="on_trace_end", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
+                )
+
+            # Possible exceptions in post_process will be caught in the overall exception handler
+            trace_spans = await self._post_process_rollout_result(next_rollout, result)
+            last_reward = find_final_reward(trace_spans)
+
+            end_time = time.time()
+            logger.info(
+                f"{self._log_prefix(rollout_id)} Completed in "
+                f"{end_time - start_time:.2f}s. Collected {len(trace_spans)} span(s). "
+                f"Final reward: {last_reward}"
+            )
+
+        except Exception:
+            logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
+            has_exception = True
+
+            if raise_on_exception:
+                raise
+        finally:
+            try:
+                await self._trigger_hooks(
+                    hook_type="on_rollout_end", agent=agent, runner=self, rollout=next_rollout, spans=trace_spans
+                )
+            except Exception:
+                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+
+            try:
+                if has_exception:
+                    # possibly timed out and cancelled?
+                    await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="failed")
+                else:
+                    await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="succeeded")
+            except Exception:
+                logger.exception(
+                    f"{self._log_prefix(rollout_id)} Exception during update_attempt. Giving up the update."
+                )
+
+        return rollout_id
+
+    async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
+        """Run the runner, continuously iterating over tasks in the store.
+
+        This method polls the store for new rollouts and executes them until:
+
+        - The event is set (if provided)
+        - The max_rollouts limit is reached (if configured)
+        - No more tasks are available
+
+        All exceptions during rollout execution are caught and logged but not
+        propagated, allowing the runner to continue processing subsequent tasks.
+
+        Args:
+            event: Optional ExecutionEvent object to signal the runner to stop. The runner
+                will check this event periodically and stop gracefully when set.
+        """
+        num_tasks_processed = 0
+        logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_rollouts or 'unlimited'}).")
+        store = self.get_store()
+
+        while not (event is not None and event.is_set()) and (
+            self._max_rollouts is None or num_tasks_processed < self._max_rollouts
+        ):
+            # Retrieve the next rollout
+            next_rollout: Optional[Rollout] = None
+            while not (event is not None and event.is_set()):
+                logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
+                next_rollout = await store.dequeue_rollout()
+                if next_rollout is None:
+                    logger.debug(f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds.")
+                    await self._sleep_until_next_poll(event)
+                else:
+                    break
+
+            if next_rollout is None:
+                return
+
+            try:
+                # Claim the rollout but updating the current worker id
+                await store.update_attempt(
+                    next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
+                )
+            except Exception:
+                # This exception could happen if the rollout is dequeued and the other end died for some reason
+                logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
+                continue
+
+            # Execute the step
+            await self._step_impl(next_rollout)
+
+            num_tasks_processed += 1
+            if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
+                logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}")
+
+        logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[NamedResources] = None,
+        mode: Optional[RolloutMode] = None,
+        event: Optional[ExecutionEvent] = None,
+    ) -> Rollout:
+        """Execute a single task directly, bypassing the task queue.
+
+        This method creates a new rollout for the given input and executes it
+        immediately. Unlike [`iter()`][agentlightning.LitAgentRunner.iter],
+        exceptions are propagated to the caller.
+
+        Args:
+            input: The task input to be processed by the agent.
+            resources: Optional named resources to be used for this specific task.
+                If provided, a new resources entry will be created in the store.
+                If not provided, the latest resources from the store will be used.
+            mode: Optional rollout mode ("train" or "validation"). If not provided,
+                the agent's default mode will be used.
+            event: Optional ExecutionEvent object to signal interruption (currently unused
+                but included for interface consistency).
+
+        Returns:
+            The completed rollout.
+
+        Raises:
+            Exception: Any exception that occurs during rollout execution will be
+                re-raised to the caller.
+        """
+        store = self.get_store()
+
+        if resources is not None:
+            resources_update = await store.add_resources(resources)
+            resources_id = resources_update.resources_id
+        else:
+            resources_id = None
+
+        attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
+        rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
+
+        completed_rollout = await store.get_rollout_by_id(rollout_id)
+        if completed_rollout is None:
+            raise RuntimeError(f"{self._log_prefix()} Failed to fetch completed rollout by id after step: {rollout_id}")
+        return completed_rollout
@@ -0,0 +1,182 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Abstract runner interface for executing agent tasks."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Generic, Iterator, Optional, Sequence, TypeVar
+
+from agentlightning.execution.events import ExecutionEvent
+from agentlightning.litagent import LitAgent
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Hook, NamedResources, ParallelWorkerBase, Rollout, RolloutMode
+
+if TYPE_CHECKING:
+    from agentlightning.execution.events import ExecutionEvent
+
+
+T_task = TypeVar("T_task")
+
+logger = logging.getLogger(__name__)
+
+
+class Runner(ParallelWorkerBase, Generic[T_task]):
+    """Abstract base class for long-running agent executors.
+
+    Runner implementations coordinate [`LitAgent`][agentlightning.LitAgent]
+    instances, acquire work from a [`LightningStore`][agentlightning.LightningStore],
+    and emit [`Rollout`][agentlightning.Rollout] objects. Subclasses decide how
+    to schedule work (polling, streaming, etc.) while this base class provides a
+    minimal lifecycle contract.
+    """
+
+    def init(self, agent: LitAgent[T_task], **kwargs: Any) -> None:
+        """Prepare the runner to execute tasks for `agent`.
+
+        This method is called only once during the setup for all workers, not for each worker.
+
+        Args:
+            agent: Agent instance providing task-specific logic.
+            **kwargs: Optional runner-specific configuration.
+
+        Raises:
+            NotImplementedError: Subclasses must supply the initialization
+                routine.
+        """
+        raise NotImplementedError()
+
+    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        """Configure worker-local state before processing tasks.
+
+        This method is called for **each** worker during the setup.
+
+        Args:
+            worker_id: Unique identifier for this worker process or thread.
+            store: Shared [`LightningStore`][agentlightning.LightningStore]
+                backing task coordination.
+            **kwargs: Optional worker-specific configuration.
+
+        Raises:
+            NotImplementedError: Subclasses must prepare per-worker resources.
+        """
+        raise NotImplementedError()
+
+    def run(self, *args: Any, **kwargs: Any) -> None:
+        """Deprecated synchronous entry point.
+
+        Use [`iter()`][agentlightning.Runner.iter] or [`step()`][agentlightning.Runner.step] instead.
+
+        Raises:
+            RuntimeError: Always raised to direct callers to
+                [iter()][agentlightning.Runner.iter] or
+                [step()][agentlightning.Runner.step].
+        """
+        raise RuntimeError("The behavior of run() of Runner is undefined. Use iter() or step() instead.")
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Release resources acquired during [`init()`][agentlightning.Runner.init].
+
+        Raises:
+            NotImplementedError: Subclasses must implement the shutdown routine.
+        """
+        raise NotImplementedError()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Release per-worker resources allocated by [`init_worker()`][agentlightning.Runner.init_worker].
+
+        Args:
+            worker_id: Identifier of the worker being torn down.
+
+        Raises:
+            NotImplementedError: Subclasses must implement the shutdown routine.
+        """
+        raise NotImplementedError()
+
+    @contextmanager
+    def run_context(
+        self,
+        *,
+        agent: LitAgent[T_task],
+        store: LightningStore,
+        hooks: Optional[Sequence[Hook]] = None,
+        worker_id: Optional[int] = None,
+    ) -> Iterator[Runner[T_task]]:
+        """Initialize and tear down a runner within a simple context manager.
+
+        The helper is primarily intended for debugging runner implementations
+        outside of a full [`Trainer`][agentlightning.Trainer] stack.
+
+        Args:
+            agent: Agent executed by this runner.
+            store: Backing [`LightningStore`][agentlightning.LightningStore].
+                If you don't have one, you can easily create one with
+                [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
+            hooks: Optional sequence of hooks recognised by the runner.
+                Not all runners support hooks.
+            worker_id: Override the worker identifier used during setup. Defaults
+                to `0`.
+        """
+        _initialized: bool = False
+        _worker_initialized: bool = False
+        try:
+            self.init(agent=agent, hooks=hooks)
+            _initialized = True
+            self.init_worker(worker_id=0, store=store)
+            _worker_initialized = True
+            yield self
+        finally:
+            try:
+                if _worker_initialized:
+                    self.teardown_worker(worker_id=worker_id if worker_id is not None else 0)
+            except Exception:
+                logger.error("Error during runner worker teardown", exc_info=True)
+
+            try:
+                if _initialized:
+                    self.teardown()
+            except Exception:
+                logger.error("Error during runner teardown", exc_info=True)
+
+    async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
+        """Run the runner, continuously iterating over tasks in the store.
+
+        This method runs in a loop, polling the store for new tasks and executing
+        them until interrupted by the event or when no more tasks are available.
+
+        Args:
+            event: Cooperative stop signal. When set, the runner should complete
+                the current unit of work and exit the loop.
+
+        Raises:
+            NotImplementedError: Subclasses provide the iteration behavior.
+        """
+        raise NotImplementedError()
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[NamedResources] = None,
+        mode: Optional[RolloutMode] = None,
+        event: Optional[ExecutionEvent] = None,
+    ) -> Rollout:
+        """Execute a single task with the given input.
+
+        This method provides fine-grained control for executing individual tasks
+        directly, bypassing the store's task queue.
+
+        Args:
+            input: Task payload consumed by the agent.
+            resources: Optional named resources scoped to this invocation.
+            mode: Optional rollout mode such as `"train"` or `"eval"`.
+            event: Cooperative stop signal for long-running tasks.
+
+        Returns:
+            Completed rollout produced by the agent.
+
+        Raises:
+            NotImplementedError: Subclasses provide the execution behavior.
+        """
+        raise NotImplementedError()
@@ -1,25 +1,29 @@
-import asyncio
+# Copyright (c) Microsoft. All rights reserved.
+
 import json
 import logging
-import os
 import time
-from contextlib import nullcontext
-from typing import List, Optional, Union, Dict, Any
-
-import agentops
+from typing import Any, Dict, List, Optional, cast

 from opentelemetry.sdk.trace import ReadableSpan
-from .client import AgentLightningClient
-from .litagent import LitAgent
-from .types import Rollout, Task, Triplet, RolloutRawResult
-from .types import ParallelWorkerBase
-from .tracer.base import BaseTracer
-from .tracer import TripletExporter
+
+from agentlightning.adapter import TracerTraceToTriplet
+from agentlightning.client import AgentLightningClient
+from agentlightning.litagent import LitAgent
+from agentlightning.litagent.litagent import is_v0_1_rollout_api
+from agentlightning.tracer.base import Tracer
+from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet
+
+from .base import Runner

 logger = logging.getLogger(__name__)

+__all__ = [
+    "LegacyAgentRunner",
+]

-class AgentRunner(ParallelWorkerBase):
+
+class LegacyAgentRunner(Runner[Any]):
    """Manages the agent's execution loop and integrates with AgentOps.

    This class orchestrates the interaction between the agent (`LitAgent`) and
@@ -37,10 +41,10 @@ class AgentRunner(ParallelWorkerBase):

    def __init__(
        self,
-        agent: LitAgent,
+        agent: LitAgent[Any],
        client: AgentLightningClient,
-        tracer: BaseTracer,
-        triplet_exporter: TripletExporter,
+        tracer: Tracer,
+        triplet_exporter: TracerTraceToTriplet,
        worker_id: Optional[int] = None,
        max_tasks: Optional[int] = None,
    ):
@@ -54,30 +58,43 @@ class AgentRunner(ParallelWorkerBase):
        self.worker_id = worker_id
        self.max_tasks = max_tasks

+    # These methods are overridden by Runner, getting them back to old behavior.
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        self.worker_id = worker_id
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
    def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
        """Generates a standardized log prefix for the current worker."""
        if self.worker_id is not None:
            if rollout_id:
-                return f"[Worker {self.worker_id} | Rollout {rollout_id}]"
+                return f"[Worker {self.worker_id} | RolloutLegacy {rollout_id}]"
            else:
                return f"[Worker {self.worker_id}]"
        if rollout_id:
-            return f"[Rollout {rollout_id}]"
+            return f"[RolloutLegacy {rollout_id}]"
        return "[Default Worker]"

    def _to_rollout_object(
        self,
-        result: RolloutRawResult,
+        result: RolloutRawResultLegacy,
        rollout_id: str,
-    ) -> Rollout:
-        """Standardizes the agent's return value into a Rollout object.
+    ) -> RolloutLegacy:
+        """Standardizes the agent's return value into a RolloutLegacy object.

        Args:
            result: The output from the agent's rollout method.
            rollout_id: The unique identifier for the current task.

        Returns:
-            A standardized `Rollout` object for reporting to the server.
+            A standardized `RolloutLegacy` object for reporting to the server.
        """
        trace: Any = None
        final_reward: Optional[float] = None
@@ -98,8 +115,8 @@ class AgentRunner(ParallelWorkerBase):
        # Case 4: result is a list of dict (trace JSON)
        if isinstance(result, list) and all(isinstance(t, dict) for t in result):
            trace = result
-        # Case 5: result is a Rollout object
-        if isinstance(result, Rollout):
+        # Case 5: result is a RolloutLegacy object
+        if isinstance(result, RolloutLegacy):
            final_reward = result.final_reward
            triplets = result.triplets
            trace = result.trace
@@ -111,15 +128,15 @@ class AgentRunner(ParallelWorkerBase):
                trace = [json.loads(readable_span.to_json()) for readable_span in spans]
                trace_spans = spans

-        # Always extract triplets from the trace using TripletExporter
+        # Always extract triplets from the trace using TracerTraceToTriplet
        if trace_spans:
-            triplets = self.triplet_exporter.export(trace_spans)
+            triplets = self.triplet_exporter(trace_spans)  # type: ignore

        # If the agent has triplets, use the last one for final reward if not set
        if triplets and triplets[-1].reward is not None and final_reward is None:
            final_reward = triplets[-1].reward

-        # Create the Rollout object with standardized fields
+        # Create the RolloutLegacy object with standardized fields
        result_dict: Dict[str, Any] = {
            "rollout_id": rollout_id,
        }
@@ -130,11 +147,11 @@ class AgentRunner(ParallelWorkerBase):
        if trace is not None:
            result_dict["trace"] = trace

-        if isinstance(result, Rollout):
+        if isinstance(result, RolloutLegacy):
            return result.model_copy(update=result_dict)
-        return Rollout(**result_dict)
+        return RolloutLegacy(**result_dict)

-    def run(self) -> bool:
+    def run(self) -> bool:  # type: ignore
        """Poll the task and rollout once synchronously."""
        self.agent.set_runner(self)  # Ensure the agent has a reference to this runner

@@ -155,7 +172,7 @@ class AgentRunner(ParallelWorkerBase):
            logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
            return False

-        rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout

        try:
            try:
@@ -163,12 +180,20 @@ class AgentRunner(ParallelWorkerBase):
            except Exception:
                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")

-            with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
+            with self.tracer._trace_context_sync(name=f"rollout_{rollout_id}"):  # pyright: ignore[reportPrivateUsage]
                start_time = time.time()
                rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
                # Pass the task input, not the whole task object
-                result = rollout_method(task.input, task.rollout_id, resources_update.resources)
-                rollout_obj = self._to_rollout_object(result, task.rollout_id)
+                if is_v0_1_rollout_api(rollout_method):
+                    result = cast(
+                        RolloutRawResultLegacy,
+                        rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
+                    )  # type: ignore
+                else:
+                    result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
                end_time = time.time()
                logger.info(
                    f"{self._log_prefix(rollout_id)} Completed in "
@@ -181,14 +206,14 @@ class AgentRunner(ParallelWorkerBase):
            logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
        finally:
            try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
            except Exception:
                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
            self.client.post_rollout(rollout_obj)

        return True

-    def iter(self) -> int:
+    def iter(self) -> int:  # type: ignore
        """Executes the synchronous polling and rollout loop."""
        num_tasks_processed = 0
        logger.info(f"{self._log_prefix()} Started sync rollouts (max: {self.max_tasks or 'unlimited'}).")
@@ -224,7 +249,7 @@ class AgentRunner(ParallelWorkerBase):
            logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
            return False

-        rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout

        try:
            try:
@@ -232,24 +257,34 @@ class AgentRunner(ParallelWorkerBase):
            except Exception:
                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")

-            with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
+            async with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
                start_time = time.time()
                rollout_method = (
                    self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async
                )
                # Pass the task input, not the whole task object
-                result = await rollout_method(task.input, task.rollout_id, resources_update.resources)
-                rollout_obj = self._to_rollout_object(result, task.rollout_id)
+                if is_v0_1_rollout_api(rollout_method):
+                    result = cast(
+                        RolloutRawResultLegacy,
+                        await rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
+                    )  # type: ignore
+                else:
+                    result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
                end_time = time.time()
                logger.info(
                    f"{self._log_prefix(rollout_id)} Completed in "
-                    f"{end_time - start_time:.2f}s. Reward: {rollout_obj.final_reward}"
+                    f"{end_time - start_time:.2f}s. Triplet length: "
+                    f"{len(rollout_obj.triplets) if rollout_obj.triplets is not None else 'N/A'}. "
+                    f"Reward: {rollout_obj.final_reward}"
                )
        except Exception:
            logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
        finally:
            try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
            except Exception:
                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
            await self.client.post_rollout_async(rollout_obj)
@@ -1,37 +1,54 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Legacy HTTP server compatible with the original Agent Lightning protocol.
+
+The implementation in this module predates the modern store-powered runtime and
+is kept for backwards compatibility with older deployments. New applications
+should migrate to the store architecture where possible.
+"""
+
+from __future__ import annotations
+
 import asyncio
 import logging
+import threading
 import time
 import uuid
-import threading
+import warnings
 from contextlib import asynccontextmanager
-from typing import Any, Dict, List, Optional, Literal
+from typing import Any, Dict, List, Literal, Optional

 import uvicorn
 from fastapi import FastAPI, HTTPException, Path
-from pydantic import Field

 from .types import (
-    Rollout,
+    GenericResponse,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutLegacy,
    Task,
    TaskIfAny,
-    NamedResources,
-    GenericResponse,
-    ResourcesUpdate,
 )

 logger = logging.getLogger(__name__)


 class ServerDataStore:
-    """
-    A centralized, thread-safe, async, in-memory data store for the server's state.
-    This holds the task queue, versioned resources, and completed rollouts.
+    """Async-safe container for in-memory server state.
+
+    The store tracks queued tasks, claimed tasks, uploaded rollouts, and the
+    currently published resources. All interactions are guarded by asyncio locks
+    so that the FastAPI handlers can safely run in parallel.
+
+    !!! warning "Deprecated"
+        [`ServerDataStore`][agentlightning.server.ServerDataStore] is part of
+        the legacy client/server stack. Use [`LightningStore`][agentlightning.LightningStore] instead.
    """

    def __init__(self):
        self._task_queue: asyncio.Queue[Task] = asyncio.Queue()
        self._processing_tasks: Dict[str, Task] = {}  # Currently processing tasks
-        self._completed_rollouts: Dict[str, Rollout] = {}
+        self._completed_rollouts: Dict[str, RolloutLegacy] = {}

        # Store for versioned resources
        self._resource_versions: Dict[str, NamedResources] = {}
@@ -48,8 +65,18 @@ class ServerDataStore:
        resources_id: str | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> str:
-        """
-        Adds a new task to the queue with specific metadata and returns its unique ID.
+        """Enqueue a new task and return the generated rollout identifier.
+
+        Args:
+            sample: Payload that describes the task input.
+            mode: Phase in which the sample should be executed (`"train"`, `"val"`, or
+                `"test"`).
+            resources_id: Identifier of a resource bundle that the executor should
+                load before running the task.
+            metadata: Optional metadata forwarded to the executor.
+
+        Returns:
+            Unique rollout identifier assigned to the task.
        """
        rollout_id = f"rollout-{uuid.uuid4()}"
        task = Task(
@@ -66,9 +93,11 @@ class ServerDataStore:
        return rollout_id

    async def get_next_task(self) -> Optional[Task]:
-        """
-        Retrieves the next task from the queue without blocking.
-        Returns None if the queue is empty.
+        """Retrieve the next task from the queue without blocking.
+
+        Returns:
+            Next [`Task`][agentlightning.Task] ready to execute, or ``None``
+            when the queue is empty.
        """
        try:
            async with self._results_lock:
@@ -89,8 +118,10 @@ class ServerDataStore:
            return None

    async def update_resources(self, update: ResourcesUpdate):
-        """
-        Safely stores a new version of named resources and sets it as the latest.
+        """Persist a new resource bundle and mark it as the latest version.
+
+        Args:
+            update: Resource payload received from a client.
        """
        # TODO: evict old resources if necessary.
        async with self._resources_lock:
@@ -99,54 +130,70 @@ class ServerDataStore:
            logger.info(f"Resources updated. New version '{update.resources_id}' is now latest.")

    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves a specific version of named resources by its ID.
+        """Retrieve a specific resource bundle by identifier.
+
+        Args:
+            resources_id: Identifier that was previously published to the store.
+
+        Returns:
+            Matching [`ResourcesUpdate`][agentlightning.ResourcesUpdate]
+            instance, or ``None`` when the identifier is unknown.
        """
        async with self._resources_lock:
            resources = self._resource_versions.get(resources_id)
            if resources:
-                return ResourcesUpdate(resources_id=resources_id, resources=resources)
+                return ResourcesUpdate(
+                    resources_id=resources_id,
+                    resources=resources,
+                    create_time=time.time(),
+                    update_time=time.time(),
+                    version=1,
+                )
            return None

    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves the latest version of named resources.
-        """
+        """Return the most recent resource bundle, if one exists."""
        if self._latest_resources_id:
            return await self.get_resources_by_id(self._latest_resources_id)
        return None

-    async def store_rollout(self, rollout: Rollout):
-        """
-        Safely stores a completed rollout from a client.
+    async def store_rollout(self, rollout: RolloutLegacy):
+        """Persist a completed rollout for later inspection.
+
+        Args:
+            rollout: Rollout returned by a client.
        """
        async with self._results_lock:
            self._processing_tasks.pop(rollout.rollout_id, None)
            self._completed_rollouts[rollout.rollout_id] = rollout
            logger.info(f"Rollout received and stored: {rollout.rollout_id}")

-    async def retrieve_rollout(self, rollout_id: str) -> Optional[Rollout]:
-        """
-        Safely retrieves a single rollout by its ID, removing it from the store.
+    async def retrieve_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
+        """Retrieve and remove a stored rollout by identifier.
+
+        Args:
+            rollout_id: Identifier of the rollout to fetch.
+
+        Returns:
+            Stored [`RolloutLegacy`][agentlightning.RolloutLegacy], or ``None``
+            when the identifier is unknown.
        """
        async with self._results_lock:
            return self._completed_rollouts.pop(rollout_id, None)

-    async def retrieve_completed_rollouts(self) -> List[Rollout]:
-        """
-        Retrieves all completed rollouts and clears the store.
-        """
+    async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
+        """Return all completed rollouts and clear the internal buffer."""
        async with self._results_lock:
            rollouts = list(self._completed_rollouts.values())
            self._completed_rollouts.clear()
            return rollouts

    def get_processing_tasks(self) -> Dict[str, Task]:
-        """Returns a copy of currently processing tasks for timeout checking."""
+        """Return a copy of currently processing tasks for timeout checking."""
        return self._processing_tasks.copy()

    async def requeue_task(self, task: Task):
-        """Requeues a task that has timed out and removes it from processing."""
+        """Requeue a task that timed out while being processed."""
        logger.warning(f"Requeuing task {task.rollout_id} after timeout (attempt {task.num_claims})")
        async with self._results_lock:
            # Remove from processing tasks
@@ -155,22 +202,30 @@ class ServerDataStore:


 class AgentLightningServer:
-    """
-    The main SDK class for developers to control the Agent Lightning Server.
+    """High-level controller for the legacy Agent Lightning FastAPI server.

-    This class manages the server lifecycle, task queueing, resources updates,
-    and retrieval of results, providing a simple interface for the optimization logic.
+    The controller orchestrates server start-up, task queueing, resource updates,
+    and retrieval of client rollouts. It is primarily used by existing systems that
+    still rely on the HTTP-based workflow.
+
+    !!! warning "Deprecated"
+        [`AgentLightningServer`][agentlightning.server.AgentLightningServer] is part of
+        the legacy client/server stack. Prefer the store-based runtime for new
+        integrations.
    """

    def __init__(self, host: str = "127.0.0.1", port: int = 8000, task_timeout_seconds: float = 300.0):
-        """
-        Initializes the server controller.
+        """Initialize the controller.

        Args:
-            host: The host to bind the server to.
-            port: The port to bind the server to.
-            task_timeout_seconds: Time in seconds after which a claimed task is considered stale and requeued.
+            host: Hostname or IP address to bind the HTTP server to.
+            port: TCP port exposed by the server.
+            task_timeout_seconds: Seconds before a claimed task is considered stale and
+                re-queued.
        """
+        warnings.warn(
+            "AgentLightningServer is deprecated. Please use LightningStoreServer instead.", DeprecationWarning
+        )
        self.host = host
        self.port = port
        self.endpoint = f"http://{host}:{port}"
@@ -191,9 +246,7 @@ class AgentLightningServer:
    # --- ADDED: Lifespan context manager ---
    @asynccontextmanager
    async def _lifespan(self, app: FastAPI):
-        """
-        Manages server startup and shutdown. This runs inside the server's event loop.
-        """
+        """Manage server start-up and shutdown within the event loop."""
        logger.info("Server is starting up...")
        self.loop = asyncio.get_running_loop()
        self._store = ServerDataStore()  # Initialize data store here
@@ -207,16 +260,14 @@ class AgentLightningServer:
        self.loop = None

    async def _check_and_requeue_stale_tasks(self):
-        """
-        Check for stale tasks and requeue them. Called reactively during get_next_task.
-        """
+        """Check for stale tasks and requeue them when they exceed the timeout."""
        current_time = time.time()
        # Ensure store is initialized before checking
        if not self._store:
            return
        processing_tasks = self._store.get_processing_tasks()

-        for rollout_id, task in processing_tasks.items():
+        for _, task in processing_tasks.items():
            if task.last_claim_time and current_time - task.last_claim_time > self._task_timeout_seconds:
                await self._store.requeue_task(task)
                logger.warning(
@@ -224,11 +275,11 @@ class AgentLightningServer:
                )

    def _setup_routes(self):
-        """Setup FastAPI routes."""
+        """Configure the FastAPI routes that make up the legacy HTTP API."""

        @self._app.get("/task", response_model=TaskIfAny)
-        async def next_task() -> TaskIfAny:
-            """Endpoint for clients to poll for the next available task."""
+        async def next_task() -> TaskIfAny:  # type: ignore
+            """Provide the next available task to a client."""
            await self._check_and_requeue_stale_tasks()

            if not self._store:
@@ -243,8 +294,8 @@ class AgentLightningServer:
                return TaskIfAny(is_available=False)

        @self._app.get("/resources/latest", response_model=ResourcesUpdate)
-        async def fetch_latest_resources() -> ResourcesUpdate:
-            """Endpoint for clients to poll for the latest available resources."""
+        async def fetch_latest_resources() -> ResourcesUpdate:  # type: ignore
+            """Return the most recent resource bundle published to the server."""
            if not self._store:
                raise HTTPException(status_code=503, detail="Server not fully initialized.")
            resources_update = await self._store.get_latest_resources()
@@ -254,10 +305,10 @@ class AgentLightningServer:
            return resources_update

        @self._app.get("/resources/{resource_id}", response_model=ResourcesUpdate)
-        async def fetch_resources_by_id(
+        async def fetch_resources_by_id(  # type: ignore
            resource_id: str = Path(..., description="The unique identifier for the resource version.")
        ) -> ResourcesUpdate:
-            """Endpoint for clients to fetch a specific version of resources."""
+            """Return a specific version of resources by identifier."""
            if not self._store:
                raise HTTPException(status_code=503, detail="Server not fully initialized.")
            resources_update = await self._store.get_resources_by_id(resource_id)
@@ -267,8 +318,8 @@ class AgentLightningServer:
            return resources_update

        @self._app.post("/rollout", response_model=GenericResponse)
-        async def post_rollout(payload: Rollout) -> GenericResponse:
-            """Endpoint for clients to report a completed rollout."""
+        async def post_rollout(payload: RolloutLegacy) -> GenericResponse:  # type: ignore
+            """Persist the rollout reported by a client."""
            if not self._store:
                raise HTTPException(status_code=503, detail="Server not fully initialized.")
            await self._store.store_rollout(payload)
@@ -278,13 +329,13 @@ class AgentLightningServer:
            )

    async def start(self):
-        """Starts the FastAPI server in the background."""
+        """Start the FastAPI server in the background."""
        logger.info(f"Starting server at {self.endpoint}")
        asyncio.create_task(self._uvicorn_server.serve())
        await asyncio.sleep(1)  # Allow time for server to start up.

    async def stop(self):
-        """Gracefully stops the running FastAPI server."""
+        """Stop the FastAPI server and wait for a graceful shutdown."""
        if self._uvicorn_server.started:
            logger.info("Stopping server...")
            self._uvicorn_server.should_exit = True
@@ -292,10 +343,7 @@ class AgentLightningServer:
            logger.info("Server stopped.")

    async def run_forever(self):
-        """
-        Runs the server indefinitely until stopped.
-        This is useful when async start and stop methods do not work.
-        """
+        """Run the server indefinitely until `stop()` is invoked."""
        await self._uvicorn_server.serve()

    async def queue_task(
@@ -305,35 +353,37 @@ class AgentLightningServer:
        resources_id: str | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> str:
-        """
-        Adds a task to the queue for a client to process.
-        """
+        """Add a task to the queue for a client to process."""
        if not self._store:
            raise RuntimeError("Store not initialized. The server may not be running.")
        return await self._store.add_task(sample, mode=mode, resources_id=resources_id, metadata=metadata)

    async def update_resources(self, resources: NamedResources) -> str:
-        """
-        Updates the resources, creating a new version and setting it as the latest.
-        """
+        """Publish a new resource bundle and return its generated identifier."""
        if not self._store:
            raise RuntimeError("Store not initialized. The server may not be running.")
        resources_id = f"res-{uuid.uuid4()}"
-        update = ResourcesUpdate(resources_id=resources_id, resources=resources)
+        update = ResourcesUpdate(
+            resources_id=resources_id, resources=resources, create_time=time.time(), update_time=time.time(), version=1
+        )
        await self._store.update_resources(update)
        return resources_id

-    async def get_completed_rollout(self, rollout_id: str) -> Optional[Rollout]:
-        """
-        Retrieves a specific completed rollout by its ID.
-        """
+    async def get_completed_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
+        """Retrieve a specific completed rollout by identifier."""
        if not self._store:
            raise RuntimeError("Store not initialized. The server may not be running.")
        return await self._store.retrieve_rollout(rollout_id)

-    async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
-        """
-        Polls for a completed rollout by its ID, waiting up to `timeout` seconds.
+    async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[RolloutLegacy]:
+        """Poll for a completed rollout until it becomes available or a timeout expires.
+
+        Args:
+            rollout_id: Identifier of the rollout to wait for.
+            timeout: Maximum number of seconds to wait. ``None`` waits indefinitely.
+
+        Returns:
+            Retrieved rollout, or ``None`` when the timeout is reached without success.
        """
        start_time = time.time()
        while True:
@@ -344,10 +394,8 @@ class AgentLightningServer:
                return None
            await asyncio.sleep(1)

-    async def retrieve_completed_rollouts(self) -> List[Rollout]:
-        """
-        Retrieves all available completed trajectories and clears the internal store.
-        """
+    async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
+        """Return every completed rollout and clear the internal buffer."""
        if not self._store:
            raise RuntimeError("Store not initialized. The server may not be running.")
        return await self._store.retrieve_completed_rollouts()
@@ -0,0 +1,17 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import LightningStore, LightningStoreCapabilities
+from .client_server import LightningStoreClient, LightningStoreServer
+from .database import SqlLightningStore
+from .memory import InMemoryLightningStore
+from .threading import LightningStoreThreaded
+
+__all__ = [
+    "LightningStore",
+    "LightningStoreCapabilities",
+    "LightningStoreClient",
+    "LightningStoreServer",
+    "InMemoryLightningStore",
+    "LightningStoreThreaded",
+    "SqlLightningStore",
+]
@@ -0,0 +1,546 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional, Sequence, Union, TypedDict
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+)
+
+
+def is_queuing(rollout: Rollout) -> bool:
+    return rollout.status == "queuing" or rollout.status == "requeuing"
+
+
+def is_running(rollout: Rollout) -> bool:
+    return rollout.status == "preparing" or rollout.status == "running"
+
+
+def is_finished(rollout: Rollout) -> bool:
+    return rollout.status == "failed" or rollout.status == "succeeded" or rollout.status == "cancelled"
+
+
+class _UnsetType:
+    """A sentinel type to indicate an unset value."""
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "UNSET"
+
+    def __reduce__(self):
+        return (_get_unset, ())
+
+
+def _get_unset() -> _UnsetType:
+    return UNSET
+
+
+UNSET = _UnsetType()
+Unset = _UnsetType  # Alias for convenience
+
+
+class LightningStoreCapabilities(TypedDict):
+    """Capability of a LightningStore implementation."""
+
+    thread_safe: bool
+    """Whether the store is thread-safe."""
+    async_safe: bool
+    """Whether the store is async-safe."""
+    zero_copy: bool
+    """Whether the store has only one copy across all threads/processes."""
+
+
+class LightningStore:
+    """Contract for the persistent control-plane that coordinates training rollouts.
+
+    A `LightningStore` mediates every interaction between algorithms and runners:
+
+    - **Rollout lifecycle:** accept new rollouts, queue them for execution, create attempts,
+      and drive the rollout status machine (`"queuing"` → `"preparing"` → `"running"` →
+      `{"succeeded","failed","cancelled"}` or `"requeuing"` when a retry is justified).
+    - **Attempt tracking:** record each execution attempt, including progress heartbeats,
+      retry sequencing, and terminal states such as `"timeout"` or `"unresponsive"`.
+    - **Span ingest:** capture structured telemetry emitted by runners (either as native
+      [`Span`][agentlightning.Span] objects or as `opentelemetry.sdk.trace.ReadableSpan`
+      instances) so that algorithms can reconstruct trajectories and rewards.
+    - **Resource versioning:** manage immutable snapshots of named resources
+      (prompt templates, model checkpoints, proxy endpoints, …) and expose a single
+      "latest" snapshot that runners can fetch just after claiming work.
+
+    Implementations must provide thread-safe/async-safe semantics: each coroutine should
+    appear atomic to callers even when multiple algorithms or runners call the API concurrently.
+    Unless stated otherwise, missing identifiers should result in a `ValueError`.
+    """
+
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        return LightningStoreCapabilities(
+            thread_safe=False,
+            async_safe=False,
+            zero_copy=False,
+        )
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        """Register a rollout and immediately create its first attempt.
+
+        !!! note
+            Use [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout] when the
+            caller only wants to submit work for later scheduling.
+
+        The rollout must be persisted with `status="preparing"` and an initial attempt
+        with `sequence_id == 1` so the caller can begin execution without visiting the
+        public queue. Implementations are expected to:
+
+        1. Generate a unique `rollout_id` and `attempt_id`.
+        2. Record `start_time` for both rollout and attempt based on the current clock.
+        3. Copy `config` and `metadata` so later mutations do not leak shared references.
+        4. Resolve `resources_id` to the latest resource snapshot when `None` is supplied.
+
+        Args:
+            input: Arbitrary task payload supplied by an algorithm.
+            mode: Optional semantic mode for downstream analytics (`"train"`, `"val"`, `"test"`).
+            resources_id: Concrete resource snapshot to execute against; defaults to the latest stored snapshot.
+            config: Rollout retry/timeout policy. Should default to a fresh [`RolloutConfig`][agentlightning.RolloutConfig].
+            metadata: Free-form metadata persisted verbatim with the rollout.
+
+        Returns:
+            The fully-populated [`AttemptedRollout`][agentlightning.AttemptedRollout] including
+            the just-created attempt.
+
+        Raises:
+            NotImplementedError: Subclasses must provide durable storage for the rollout.
+            ValueError: Implementations should raise when `resources_id` does not exist.
+        """
+        raise NotImplementedError()
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        """Persist a rollout in `queuing` state so runners can claim it later.
+
+        !!! note
+            Different from [`start_rollout()`][agentlightning.LightningStore.start_rollout],
+            this method is called when the caller only wants to submit work for later scheduling.
+
+        Implementations must generate a unique `rollout_id`, stamp `start_time` with
+        the current time, default `config` to a fresh [`RolloutConfig`][agentlightning.RolloutConfig],
+        and insert the rollout at the tail of the scheduling queue. No attempt is created yet.
+
+        Args:
+            input: Arbitrary task payload supplied by an algorithm.
+            mode: Optional semantic mode indicator (`"train"`, `"val"`, `"test"`).
+            resources_id: Resource snapshot used when a runner eventually executes the rollout.
+            config: Fine-grained retry/timeout parameters to persist with the rollout.
+            metadata: Free-form metadata stored verbatim with the rollout record.
+
+        Returns:
+            The stored [`Rollout`][agentlightning.Rollout] in `queuing` status.
+
+        Raises:
+            NotImplementedError: Subclasses must persist the rollout.
+            ValueError: Implementations should raise when `resources_id` does not exist.
+        """
+        raise NotImplementedError()
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        """Claim the oldest queued rollout and transition it to `preparing`.
+
+        This function do not block.
+
+        Retrieval must be FIFO across rollouts that remain in `queuing` or `requeuing`
+        state. When a rollout is claimed, implementations must:
+
+        * Transition its status to `"preparing"`.
+        * Create a new attempt with `status="preparing"` and `sequence_id` equal to
+          the number of attempts already registered for the rollout plus one.
+        * Return an [`AttemptedRollout`][agentlightning.AttemptedRollout] snapshot so the
+          runner knows both rollout metadata and the attempt identifier.
+
+        Returns:
+            The next attempt to execute, or `None` when no eligible rollouts are queued.
+
+        Raises:
+            NotImplementedError: Subclasses must implement queue retrieval.
+        """
+        raise NotImplementedError()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        """Create a manual retry attempt for an existing rollout.
+
+        This is typically invoked by runners that wish to retry outside of the
+        normal queue flow (for example in an online RL setup).
+        Implementations must validate that the rollout exists, allocate a fresh `attempt_id`,
+        increment the `sequence_id` monotonically, stamp the new attempt with `status="preparing"`,
+        and return an up-to-date [`AttemptedRollout`][agentlightning.AttemptedRollout].
+
+        Args:
+            rollout_id: Unique identifier of the rollout receiving a new attempt.
+
+        Returns:
+            The rollout paired with its newly-created attempt.
+
+        Raises:
+            NotImplementedError: Subclasses must implement attempt creation.
+            ValueError: Implementations must raise when `rollout_id` is unknown.
+        """
+        raise NotImplementedError()
+
+    async def add_span(self, span: Span) -> Span:
+        """Persist a pre-constructed span emitted during rollout execution.
+
+        The provided [`Span`][agentlightning.Span] must already contain the `rollout_id`,
+        `attempt_id`, and `sequence_id`. Implementations must:
+
+        * Verify that both rollout and attempt exist.
+        * Ensure span ordering remains strictly increasing per attempt (rejecting or keeping duplicates).
+        * Treat the span arrival as a heartbeat: update the attempt's `last_heartbeat_time`
+          and transition both attempt and rollout to `"running"` if they were still
+          `"preparing"` or `"requeuing"`.
+
+        Args:
+            span: Fully populated span to persist.
+
+        Returns:
+            The stored span record (implementations may return a copy).
+
+        Raises:
+            NotImplementedError: Subclasses must implement span persistence.
+            ValueError: Implementations must raise when the referenced rollout or attempt is missing.
+        """
+        raise NotImplementedError()
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        """Convert and persist an OpenTelemetry span for a particular attempt.
+
+        Implementations must transform the `readable_span` into a [`Span`][agentlightning.Span]
+        (typically via [`Span.from_opentelemetry()`][agentlightning.Span.from_opentelemetry]),
+        assign a strictly increasing `sequence_id` when one is not provided, and persist it
+        using the same semantics as [`add_span()`][agentlightning.LightningStore.add_span].
+
+        Args:
+            rollout_id: Identifier of the rollout that produced the span.
+            attempt_id: Attempt identifier the span belongs to.
+            readable_span: OpenTelemetry span in SDK form.
+            sequence_id: Optional explicit ordering hint. When omitted, call
+                [`get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id]
+                automatically.
+
+        Returns:
+            The stored span record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement span persistence.
+            ValueError: Implementations must raise when the rollout or attempt is unknown.
+        """
+        raise NotImplementedError()
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[Union[Rollout, AttemptedRollout]]:
+        """Retrieve rollouts filtered by status and/or explicit identifiers.
+
+        Args:
+            status: Optional whitelist of [`RolloutStatus`][agentlightning.RolloutStatus] values.
+            rollout_ids: Optional whitelist of rollout identifiers to include.
+
+        Returns:
+            A list of matching rollouts. Ordering is backend-defined but must be deterministic.
+
+        Raises:
+            NotImplementedError: Subclasses must implement the query.
+        """
+        raise NotImplementedError()
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        """Return every attempt ever created for `rollout_id` in ascending sequence order.
+
+        Args:
+            rollout_id: Identifier of the rollout being inspected.
+
+        Returns:
+            Attempts sorted by `sequence_id` (oldest first). Returns an empty list when none exist.
+
+        Raises:
+            NotImplementedError: Subclasses must implement the query.
+            ValueError: Implementations must raise when the rollout does not exist.
+        """
+        raise NotImplementedError()
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
+        """Fetch a rollout by identifier without mutating its state.
+
+        Args:
+            rollout_id: Identifier to retrieve.
+
+        Returns:
+            The rollout when found, otherwise `None`.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
+        """
+        raise NotImplementedError()
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """Fetch the attempt with the highest `sequence_id` for `rollout_id`.
+
+        Args:
+            rollout_id: Identifier to inspect.
+
+        Returns:
+            The most recent attempt or `None` when no attempts exist yet.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
+            ValueError: Implementations must raise when the rollout does not exist.
+        """
+        raise NotImplementedError()
+
+    async def query_resources(self) -> List[ResourcesUpdate]:
+        """List every stored resource snapshot in insertion order.
+
+        Returns:
+            A chronological list of [`ResourcesUpdate`][agentlightning.ResourcesUpdate] objects.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
+        """
+        raise NotImplementedError()
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        """Return a specific named resource snapshot by identifier.
+
+        Args:
+            resources_id: Identifier of the snapshot.
+
+        Returns:
+            The stored [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or `None` when missing.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
+        """
+        raise NotImplementedError()
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """Fetch the latest resource snapshot marked as the global default.
+
+        Returns:
+            The current latest [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or `None` when
+            no resources have been registered yet.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
+        """
+        raise NotImplementedError()
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        """Allocate the next strictly increasing sequence number used to order spans.
+
+        Implementations must retain counters so repeated calls return `1, 2, ...` without
+        gaps unless spans were explicitly inserted with a custom `sequence_id`. The
+        counter may be scoped per rollout or per attempt, but the sequence must be
+        strictly increasing for spans emitted by the specified attempt so traces remain
+        totally ordered.
+
+        See [Distributed Tracing][distributed-tracing] for detailed motivations.
+
+        Args:
+            rollout_id: Identifier of the rollout emitting spans.
+            attempt_id: Attempt identifier for the upcoming span.
+
+        Returns:
+            The next integer sequence identifier, unique within the attempt.
+
+        Raises:
+            NotImplementedError: Subclasses must provide the allocator.
+            ValueError: Implementations must raise when the rollout or attempt does not exist.
+        """
+        raise NotImplementedError()
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        """Block until the targeted rollouts reach a terminal status or the timeout expires.
+
+        Terminal statuses are `"succeeded"`, `"failed"`, and `"cancelled"`. When the timeout
+        elapses, implementations should return the subset of rollouts that are already terminal
+        and omit the rest.
+
+        !!! warning
+            It's dangerous and might be event-loop blocking to call this function
+            with a long timeout. It's a good idea to poll for the method to check
+            if new completed rollouts can coming. Be careful in implementing the sleep logic
+            to avoid busy-waiting.
+
+        Args:
+            rollout_ids: Identifiers of rollouts to watch.
+            timeout: Maximum time in seconds to wait. `None` waits indefinitely.
+
+        Returns:
+            Rollouts that finished before the deadline, in arbitrary order.
+
+        Raises:
+            NotImplementedError: Subclasses must implement waiting semantics.
+            ValueError: Implementations must raise when a rollout identifier is unknown.
+        """
+        raise NotImplementedError()
+
+    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+        """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.
+
+        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.
+
+        Returns:
+            An ordered list of spans (possibly empty).
+
+        Raises:
+            NotImplementedError: Subclasses must implement the query.
+            ValueError: Implementations must raise when the rollout or attempt is unknown.
+        """
+        raise NotImplementedError()
+
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        """Persist a new immutable snapshot of named resources and mark it as latest.
+
+        Implementations must assign a fresh `resources_id` and ensure subsequent calls to
+        [`get_latest_resources()`][agentlightning.LightningStore.get_latest_resources] return the
+        snapshot produced here.
+
+        Args:
+            resources: Mapping of resource names to their serialized payloads.
+
+        Returns:
+            The stored [`ResourcesUpdate`][agentlightning.ResourcesUpdate] including its generated id.
+
+        Raises:
+            NotImplementedError: Subclasses must implement resource persistence.
+        """
+        raise NotImplementedError()
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        """Overwrite or extend an existing resource snapshot and mark it as latest.
+
+        This API is typically used by algorithms that maintain mutable resources (e.g., model
+        checkpoints) under a stable identifier.
+
+        If `resources_id` does not exist, implementations should add it as a new snapshot.
+
+        Args:
+            resources_id: Identifier of the snapshot to replace.
+            resources: Updated mapping of resource names to payloads.
+
+        Returns:
+            The persisted [`ResourcesUpdate`][agentlightning.ResourcesUpdate].
+
+        Raises:
+            NotImplementedError: Subclasses must implement resource persistence.
+        """
+        raise NotImplementedError()
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        """Update rollout metadata and, when provided, drive status transitions.
+
+        Parameters default to the sentinel [`UNSET`][agentlightning.store.base.UNSET] to
+        distinguish omitted fields from explicit `None` assignments. Implementations must:
+
+        * Validate the rollout exists before mutating it.
+        * Replace each property when a concrete value (including `None`) is supplied.
+        * When the status switches into a terminal state, set `end_time` and signal any waiters.
+        * When the status re-enters a queueing state, ensure the rollout is enqueued exactly once.
+
+        Args:
+            rollout_id: Identifier of the rollout to update.
+            input: Replacement task payload; pass `None` to explicitly clear the input.
+            mode: Replacement rollout mode.
+            resources_id: Replacement resources snapshot reference.
+            status: Target rollout status.
+            config: Replacement retry/timeout configuration.
+            metadata: Replacement metadata dictionary.
+
+        Returns:
+            The updated rollout record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement mutation logic.
+            ValueError: Implementations must raise when the rollout is unknown or the update is invalid.
+        """
+        raise NotImplementedError()
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        """Update attempt bookkeeping such as status, worker ownership, and heartbeats.
+
+        When `attempt_id` is `"latest"` the update must target the attempt with the highest
+        `sequence_id`; otherwise it must target the specific attempt. Implementations should
+        propagate status changes to the rollout (for example via [`propagate_status()`][agentlightning.store.utils.propagate_status])
+        once the latest attempt transitions to a terminal state.
+
+        Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],
+        parameters also default to the sentinel [`UNSET`][agentlightning.store.base.UNSET].
+
+        Args:
+            rollout_id: Identifier of the rollout whose attempt will be updated.
+            attempt_id: Attempt identifier or `"latest"` as a convenience.
+            status: Replacement attempt status. Terminal statuses must set `end_time`.
+            worker_id: Identifier for the worker currently processing the attempt.
+            last_heartbeat_time: Wall-clock timestamp (seconds) of the latest heartbeat/span.
+            metadata: Replacement metadata dictionary.
+
+        Returns:
+            The updated attempt record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement mutation logic.
+            ValueError: Implementations must raise when the rollout or attempt is unknown.
+        """
+        raise NotImplementedError()
@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .sqlite import SqlLightningStore
+
+__all__ = [
+    "SqlLightningStore",
+]
@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .attempt import AttemptInDB, SpanSeqIdInDB
+from .base import (
+    AttemptStatusUpdateMessage,
+    SqlAlchemyBase,
+)
+from .resources import ResourcesUpdateInDB
+from .rollout import RolloutInDB
+from .span import SpanInDB
+
+__all__ = [
+    "SqlAlchemyBase",
+    "AttemptStatusUpdateMessage",
+    "RolloutInDB",
+    "AttemptInDB",
+    "ResourcesUpdateInDB",
+    "SpanSeqIdInDB",
+    "SpanInDB",
+]
@@ -0,0 +1,251 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import time
+import uuid
+from dataclasses import InitVar
+from typing import Any, Dict, List, Optional
+
+from sqlalchemy import JSON, Float, Integer, String, select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from sqlalchemy.orm import Mapped, mapped_column
+
+from agentlightning.types import Attempt
+
+from .base import AttemptStatusUpdateMessage, SqlAlchemyBase
+
+logger = logging.getLogger(__name__)
+
+
+def _generate_attempt_id() -> str:
+    """We don't need that long because attempts are limited to rollouts."""
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:8]
+    return "at-" + short_id
+
+
+class AttemptInDB(SqlAlchemyBase):
+    __tablename__ = "attempts"
+
+    rollout_id: Mapped[str] = mapped_column(String, nullable=False)
+    attempt_id: Mapped[str] = mapped_column(String, primary_key=True, default_factory=_generate_attempt_id)
+    sequence_id: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+    start_time: Mapped[float] = mapped_column(Float, default_factory=time.time, nullable=False)
+    end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True, default=None)
+    status: Mapped[str] = mapped_column(String, default="preparing", nullable=False)
+    worker_id: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
+    last_heartbeat_time: Mapped[Optional[float]] = mapped_column(Float, nullable=False, default_factory=time.time)
+    attempt_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=None)
+
+    # addition columns for processing
+    max_duration: Mapped[Optional[float]] = mapped_column(
+        Float, nullable=True, default=None
+    )  # maximum duration allowed for this attempt in seconds
+    max_heartbeat_interval: Mapped[Optional[float]] = mapped_column(
+        Float, nullable=True, default=None
+    )  # maximum allowed heartbeat interval in seconds
+
+    version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
+    __mapper_args__ = {
+        "version_id_col": version_id,
+    }
+
+    def is_unresponsive(self, current_time: float) -> bool:
+        """Check if the attempt is unresponsive based on the last heartbeat time and max_heartbeat_interval."""
+        if self.max_heartbeat_interval is None:
+            return False
+        if self.last_heartbeat_time is None:
+            return False
+        return (current_time - self.last_heartbeat_time) > self.max_heartbeat_interval
+
+    def is_timed_out(self, current_time: float) -> bool:
+        """Check if the attempt has timed out based on the start time and max_duration."""
+        if self.max_duration is None:
+            return False
+        return (current_time - self.start_time) > self.max_duration
+
+    def as_attempt(self) -> Attempt:
+        return Attempt(
+            **self.model_dump(
+                exclude={"max_duration", "max_heartbeat_interval", "version_id"},
+                mapper={"metadata": lambda obj: obj.attempt_metadata},  # type: ignore
+            )
+        )
+
+    def _validate_status_message(self, msg: Dict[str, Any]) -> None:
+        """This function validates the status update message from caller.
+        Raises ValueError if the message is invalid.
+        """
+        if "event" not in msg:
+            raise ValueError("Status update message must contain 'event' field.")
+        if "timestamp" not in msg:
+            msg["timestamp"] = time.time()
+        if msg["event"] not in [
+            "user_update",  # user update attempt status via dbstore.update_attempt()
+            "span_received",  # new span received
+            "single_step_timeout",  # single step timeout detected (from last span heartbeat)
+            "overall_timeout",  # overall timeout detected
+        ]:
+            raise ValueError(f"Unsupported event type: {msg['event']}")
+        if msg["event"] == "user_update" and "new_status" not in msg:
+            raise ValueError("User update event must contain 'new_status' field.")
+
+    def get_finished_statuses(self) -> List[str]:
+        """This function returns the list of statuses that are considered finished."""
+        return [
+            "succeeded",
+            "failed",
+            "timeout",
+        ]
+
+    def update_status(self, msg: Dict[str, Any]) -> Optional[AttemptStatusUpdateMessage]:
+        """This function updates the status of the attempt based on the event.
+        Args:
+            msg: A dictionary containing the status update message. It must contain an "event" field, and optionally a "new_status" field.
+                More details about the message format can be found in the `_validate_status_message`() method.
+            current_time: The current time to use for updating timestamps. If None, uses time.time().
+        Returns:
+            A dictionary containing the status update message: {"event": "attempt_status_updated", "old_status": old_status, "new_status": new_status}.
+                IF no meaningful status update is performed, returns None.
+        Raises:
+            ValueError: If the event is not recognized or the status transition is invalid.
+            NotImplementedError: If the event handling is not implemented for the current status.
+            RuntimeError: If the new status is not set after processing the event.
+        """
+        self._validate_status_message(msg)
+        event = msg["event"]
+        current_time = msg.get("timestamp", time.time())
+        old_status = self.status
+        new_status = msg.get("new_status", None)
+
+        # Step 1: Determine the new status based on the event and current status
+        if event == "user_update":
+            if not new_status:
+                raise ValueError("new_status must be provided for user_update event.")
+        elif event == "span_received":
+            self.last_heartbeat_time = current_time
+            if old_status in ["preparing", "unresponsive", "running"]:
+                new_status = "running"
+            elif old_status in self.get_finished_statuses():
+                logger.warning(
+                    f"Span received after attempt is already in status {self.status}. No status update performed."
+                )
+                return  # no further status update needed
+            else:
+                raise NotImplementedError(f"Event {event} is not implemented for status {old_status}.")
+        elif event == "single_step_timeout":
+            if old_status in [
+                "preparing",
+                "running",
+            ]:
+                new_status = "unresponsive"
+            else:
+                logger.warning(
+                    f"Single step timeout detected but attempt is in status {self.status}. No status update performed."
+                )
+                return  # no further status update needed
+        elif event == "overall_timeout":
+            if old_status not in self.get_finished_statuses():
+                new_status = "timeout"
+            else:
+                logger.warning(
+                    f"Overall timeout detected but attempt is in status {self.status}. No status update performed."
+                )
+                return  # no further status update needed
+        else:
+            raise NotImplementedError(f"Event {event} is not implemented for status update.")
+
+        # Step 2: Update the status
+        if not new_status:
+            raise RuntimeError(
+                f"new_status should not be {new_status} after processing event for {event} on status {old_status}."
+            )
+        if new_status == old_status:
+            return  # no status change
+        if new_status in self.get_finished_statuses():
+            # when attempt is finished, set end_time
+            self.end_time = current_time
+        self.status = new_status
+
+        # Step 3: Return the status update info for further processing
+        return AttemptStatusUpdateMessage(
+            attempt_id=self.attempt_id,
+            rollout_id=self.rollout_id,
+            timestamp=current_time,
+            old_status=old_status,
+            new_status=new_status,
+        )
+
+    @classmethod
+    async def get_latest_attempt_for_rollout(
+        cls: type[AttemptInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
+    ) -> Optional[Attempt]:
+        async with session_factory() as session:
+            async with session.begin():
+                result = await session.scalars(
+                    select(cls).where(cls.rollout_id == rollout_id).order_by(cls.sequence_id.desc()).limit(1)
+                )
+                attempt_obj = result.one_or_none()
+                if attempt_obj is None:
+                    return None
+                return attempt_obj.as_attempt()
+
+    @classmethod
+    async def get_attempts_for_rollout(
+        cls: type[AttemptInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
+    ) -> List[Attempt]:
+        async with session_factory() as session:
+            async with session.begin():
+                result = await session.scalars(
+                    select(cls).where(cls.rollout_id == rollout_id).order_by(cls.sequence_id.asc())
+                )
+                return [attempt.as_attempt() for attempt in result.all()]
+
+
+class SpanSeqIdInDB(SqlAlchemyBase):
+    __tablename__ = "span_sequence"
+
+    rollout_id: Mapped[str] = mapped_column(nullable=False, primary_key=True)
+
+    # FIXME InMemoryLightningStore let all attempts under the same rollout share the same span sequence for sorting
+    # attempt_id: Mapped[str] = mapped_column(nullable=False)
+    attempt_id: InitVar[str]  # not mapped column, just for type hinting
+
+    current_sequence: Mapped[int] = mapped_column(default=1, nullable=False)
+
+    # Versioning for optimistic concurrency control
+    version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
+    __mapper_args__ = {
+        "version_id_col": version_id,
+        # "primary_key": [rollout_id, attempt_id],
+        # "primary_key": [rollout_id],
+    }
+
+    @classmethod
+    async def get_next_sequence_id(
+        cls: type[SpanSeqIdInDB],
+        session_factory: async_sessionmaker[AsyncSession],
+        rollout_id: str,
+        attempt_id: str,
+        external_seq_id: Optional[int] = None,
+    ) -> int:
+        """Get the next sequence ID with retries to handle race conditions.
+        IF external_seq_id is provided and is greater than current_sequence, set current_sequence to external_seq_id.
+        """
+        async with session_factory() as session:
+            async with session.begin():
+                seq_obj = await session.get(cls, rollout_id)
+                # seq_obj = await session.get(cls, [rollout_id, attempt_id])
+                if seq_obj is None:
+                    raise ValueError(f"Rollout {rollout_id} not found")
+                else:
+                    current_seq = (
+                        external_seq_id
+                        if external_seq_id is not None and external_seq_id > seq_obj.current_sequence
+                        else seq_obj.current_sequence
+                    )
+                    seq_obj.current_sequence = current_seq + 1
+                    await session.flush()
+                    return current_seq
@@ -0,0 +1,186 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+from pydantic import BaseModel, Field, TypeAdapter, computed_field
+
+# from dataclasses import asdict
+from sqlalchemy import JSON, TypeDecorator
+from sqlalchemy.ext.asyncio import AsyncAttrs
+from sqlalchemy.orm import DeclarativeBase, MappedAsDataclass
+
+
+class SqlAlchemyBase(AsyncAttrs, MappedAsDataclass, DeclarativeBase):
+    pass
+
+    def model_dump(
+        self,
+        exclude: set[str] | None = None,
+        mapper: Dict[str, Callable[["SqlAlchemyBase"], Any]] | None = None,
+    ) -> Dict[str, Any]:
+        """Dump the SQLAlchemy model to a dictionary.
+        Args:
+            exclude: set[str]
+                The set of field names to exclude.
+            mapper: Dict[str, Callable[[SqlAlchemyBase], Any]]
+                A mapping from field names to functions that take the model instance and return the value to be used for that field.
+                If the key is "*", the function should return a dictionary of additional fields to be added to the output.
+        Returns:
+            Dict[str, Any]: The dumped model as a dictionary.
+        """
+        exclude = exclude or set()
+        mapper = mapper or {}
+        dic = {k: getattr(self, k) for k in self.__table__.columns.keys() if k not in exclude}
+        for k, func in mapper.items():
+            if k == "*":
+                dic.update(func(self))
+            else:
+                dic[k] = func(self)
+        return dic
+
+
+class PydanticInDB(TypeDecorator[BaseModel]):
+    """Custom SQLAlchemy type to store pydantic.BaseModel as JSON in the database.
+    Attributes:
+        target_type: type[BaseModel], the type of the pydantic model to be stored.
+    """
+
+    impl = JSON
+    target_type: type[BaseModel] | None = None
+
+    def process_bind_param(self, value: BaseModel | None, dialect: Any) -> Optional[str]:
+        if value is None:
+            return None
+        if self.target_type is not None:
+            return TypeAdapter(self.target_type).validate_python(value).model_dump_json()  # type: ignore
+        return json.dumps(value)
+
+    def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[BaseModel]:
+        if value is None:
+            return None
+        if self.target_type is not None:
+            return TypeAdapter(self.target_type).validate_json(value)  # type: ignore
+        dic = json.loads(value)
+        return dic  # type: ignore
+
+
+class PydanticListInDB(TypeDecorator[list[BaseModel]]):
+    """Custom SQLAlchemy type to store List[pydantic.BaseModel] as JSON in the database.
+    Attributes:
+        value_type: type[BaseModel], the type of the pydantic model to be stored in the list.
+    """
+
+    impl = JSON
+    value_type: type[BaseModel] | None = None
+
+    def process_bind_param(self, value: List[BaseModel] | None, dialect: Any) -> Optional[str]:
+        if value is None:
+            return None
+        if self.value_type is not None:
+            lst = [TypeAdapter(self.value_type).validate_python(v).model_dump() for v in value]
+            return json.dumps(lst)
+        raise ValueError("target_type must be set for PydanticListInDB")
+
+    def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[List[BaseModel]]:
+        if value is None:
+            return None
+        if self.value_type is not None:
+            dic = json.loads(value)
+            return [TypeAdapter(self.value_type).validate_python(v) for v in dic]  # type: ignore
+        raise ValueError("target_type must be set for PydanticListInDB")
+
+
+class NamedDictBase(TypeDecorator[Dict[str, Any]]):
+    """Custom SQLAlchemy type to store Dict[str, pydantic.BaseModel] as JSON in the database.
+    Attributes:
+        target_alias: type[Dict[str, BaseModel]], the alias type of the dict.
+        value_type: type[BaseModel], the type of the values in the dict.
+
+    For example, given NamedResources = Dict[str, ResourceUnion],
+    we can define NamedDictBase with target_alias=NamedResources and target_type=ResourceUnion.
+    """
+
+    impl = JSON
+    target_alias: type | None = None
+    value_type: type[BaseModel] | Any = None
+
+    def process_bind_param(self, value: Dict[str, Any] | None, dialect: Any) -> Optional[str]:
+        if value is None:
+            return None
+
+        # ignore target_alias for when dumping because Dict is not a pydantic model
+        if self.value_type is not None:
+            dic = {
+                k: TypeAdapter(self.value_type).validate_python(v).model_dump() if isinstance(v, BaseModel) else v
+                for k, v in value.items()
+            }
+            return json.dumps(dic)
+        dic = {k: v.model_dump() if isinstance(v, BaseModel) else v for k, v in value.items()}
+        return json.dumps(dic)
+
+    def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[Dict[str, Any]]:
+        if value is None:
+            return None
+        if self.target_alias is not None:
+            return TypeAdapter(self.target_alias).validate_json(value)  # type: ignore
+        if self.value_type is not None:
+            dic = json.loads(value)
+            return {k: TypeAdapter(self.value_type).validate_python(v) for k, v in dic.items()}  # type: ignore
+        return json.loads(value)
+
+
+class DatabaseRuntimeError(Exception):
+    """Raised when a runtime error occurs during database operations.
+    Particularly used when the execution of a query fails.
+    """
+
+    pass
+
+
+class RaceConditionError(Exception):
+    """Raised when a race condition is detected during database operations."""
+
+    pass
+
+
+class NoRolloutToDequeueError(Exception):
+    """Raised when there is no rollout available to dequeue."""
+
+    pass
+
+
+class AttemptStatusUpdateMessage(BaseModel):
+    attempt_id: str
+    rollout_id: str
+    timestamp: float = Field(default_factory=time.time)
+    old_status: Optional[str] = None
+    new_status: str
+
+    @computed_field
+    @property
+    def event(self) -> str:
+        return "attempt_status_update"
+
+    @computed_field
+    @property
+    def is_failed(self) -> bool:
+        return self.new_status in ["failed", "timeout", "unresponsive"]
+
+    @computed_field
+    @property
+    def is_succeeded(self) -> bool:
+        return self.new_status == "succeeded"
+
+    @computed_field
+    @property
+    def is_finished(self) -> bool:
+        return self.is_failed or self.is_succeeded
+
+    @computed_field
+    @property
+    def is_running(self) -> bool:
+        return self.new_status in ["running", "preparing"]
@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import hashlib
+import time
+import uuid
+from typing import Optional
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from sqlalchemy.orm import Mapped, mapped_column
+
+from agentlightning.types import NamedResources, ResourcesUpdate
+
+from .base import NamedDictBase, SqlAlchemyBase
+
+
+def _generate_resources_id() -> str:
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+    return "rs-" + short_id
+
+
+class NamedResourcesInDB(NamedDictBase):
+    """Custom SQLAlchemy type to store NamedResources as JSON in the database."""
+
+    target_alias = NamedResources
+
+
+class ResourcesUpdateInDB(SqlAlchemyBase):
+    __tablename__ = "resources"
+    resources: Mapped[NamedResources] = mapped_column(
+        NamedResourcesInDB, nullable=False
+    )  # JSON serialized, convert to NamedResources when needed
+    resources_id: Mapped[str] = mapped_column(primary_key=True, default_factory=_generate_resources_id)
+    create_time: Mapped[float] = mapped_column(nullable=False, default_factory=time.time)
+    update_time: Mapped[float] = mapped_column(nullable=False, default_factory=time.time, onupdate=time.time)
+    version: Mapped[int] = mapped_column(nullable=False, default=1)
+
+    __mapper_args__ = {
+        "version_id_col": version,
+    }
+
+    @classmethod
+    async def get_resources_by_id(
+        cls, session_factory: async_sessionmaker[AsyncSession], resources_id: str
+    ) -> Optional[ResourcesUpdate]:
+        async with session_factory() as session:
+            async with session.begin():
+                obj = await session.get(cls, resources_id)
+                if obj is None:
+                    return None
+                return obj.as_resources_update()
+
+    def as_resources_update(self) -> ResourcesUpdate:
+        return ResourcesUpdate(**self.model_dump())
@@ -0,0 +1,201 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import time
+import uuid
+from typing import Any, Dict, List, Optional, cast
+
+from sqlalchemy import JSON, Float, Integer, String, and_, select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from sqlalchemy.orm import Mapped, mapped_column
+
+from agentlightning.types import AttemptedRollout, Rollout, RolloutConfig, RolloutStatus
+
+from ...base import is_finished, is_queuing
+from .attempt import AttemptInDB
+from .base import AttemptStatusUpdateMessage, PydanticInDB, SqlAlchemyBase
+
+logger = logging.getLogger(__name__)
+
+
+def _generate_rollout_id() -> str:
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+    return "ro-" + short_id
+
+
+class RolloutConfigInDB(PydanticInDB):
+    """Custom SQLAlchemy type to store RolloutConfig as JSON in the database."""
+
+    target_type = RolloutConfig
+
+
+class RolloutInDB(SqlAlchemyBase):
+    __tablename__ = "rollouts"
+
+    input: Mapped[Any] = mapped_column(JSON, nullable=False)
+    rollout_id: Mapped[str] = mapped_column(String, primary_key=True, default_factory=_generate_rollout_id)
+    start_time: Mapped[float] = mapped_column(Float, default_factory=time.time, nullable=False)
+    end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True, default=None)
+    mode: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
+    resources_id: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
+    status: Mapped[RolloutStatus] = mapped_column(String, default="queuing", nullable=False)
+    config: Mapped[RolloutConfig] = mapped_column(
+        RolloutConfigInDB, nullable=False, default_factory=RolloutConfig
+    )  # JSON serialized, convert to RolloutConfig when needed
+    rollout_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(
+        JSON, nullable=True, default=None
+    )  # JSON serialized, convert to Dict when needed
+
+    # Attempt-related helper methods can be added here if needed
+    num_attempts: Mapped[int] = mapped_column(
+        Integer, default=0, nullable=False
+    )  # number of attempts made for this rollout
+    enqueue_time: Mapped[Optional[float]] = mapped_column(
+        Float, nullable=True, default_factory=time.time
+    )  # time when the rollout was enqueued (for FIFO scheduling)
+    latest_attempt_id: Mapped[Optional[str]] = mapped_column(
+        String, nullable=True, default=None
+    )  # the attempt_id of the latest attempt
+
+    # use optimistic concurrency control
+    version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
+    __mapper_args__ = {
+        "version_id_col": version_id,
+    }
+
+    def __post_init__(self):
+        if self.status not in ["queuing", "running", "succeeded", "failed", "requeuing"]:
+            raise ValueError(f"Invalid rollout status: {self.status}")
+
+    def as_rollout(self) -> Rollout:
+        return Rollout(
+            **self.model_dump(
+                exclude={"rollout_metadata", "num_attempts", "enqueue_time", "latest_attempt_id", "version_id"},
+                mapper={
+                    "metadata": lambda obj: obj.rollout_metadata,  # type: ignore
+                    "config": lambda obj: obj.config if obj.config is not None else RolloutConfig(),  # type: ignore
+                },
+            )
+        )
+
+    def _validate_status_message(self, msg: Dict[str, str]) -> None:
+        """Validate the status update message.
+        Raises:
+            ValueError: If the message is invalid.
+        """
+        if "event" not in msg:
+            raise ValueError("Status update message must contain 'event' field.")
+        event = msg["event"]
+        if event not in [
+            "attempt_status_update",  # from attempt status update
+            "user_update",  # from user-initiated update
+        ]:
+            raise ValueError(f"Invalid event type in status update message: {event}")
+        if event == "user_update":
+            if "new_status" not in msg:
+                raise ValueError("Status update message for event 'user_update' must contain 'new_status' field.")
+        if event == "attempt_status_update":
+            # leverage AttemptStatusUpdateMessage for validation
+            pass
+
+    async def update_status(self, msg: Dict[str, Any] | AttemptStatusUpdateMessage) -> None:
+        """Update the rollout status based on the provided message.
+        Args:
+            msg (Dict[str, str]): The status update message. Refer to `_validate_status_message` for the expected format.
+            current_time (Optional[float]): The current time to set end_time or enqueue_time if needed.
+        """
+        if isinstance(msg, dict):
+            self._validate_status_message(msg)
+            event = msg["event"]
+            current_time = msg.get("timestamp", time.time())
+        else:
+            event = msg.event
+            current_time = msg.timestamp
+
+        old_status = self.status
+        new_status = self.status  # initialize new_status with old_status
+
+        # Step 1: Determine the new status based on the event
+        if event == "user_update":
+            assert isinstance(msg, dict)
+            new_status = msg["new_status"]
+        elif event == "attempt_status_update":
+            msg = AttemptStatusUpdateMessage(**msg) if isinstance(msg, dict) else msg
+            if msg.attempt_id == self.latest_attempt_id:
+                new_status = msg.new_status  # directly take the latest attempt status
+                if msg.is_succeeded:
+                    new_status = "succeeded"
+                elif msg.is_failed:
+                    # no other attempts running, decide whether to requeue or fail
+                    config = self.config
+                    if config.max_attempts > self.num_attempts and msg.new_status in config.retry_condition:
+                        new_status = "requeuing"
+                    else:
+                        new_status = "failed"
+                # elif msg.is_running and old_status in ["failed", "requeuing"]:
+                # new_status = "running"
+            else:
+                # ignore attempts from old attempts
+                new_status = old_status
+
+        # Step 2: Update the status if it has changed and handle follow-up actions
+        if new_status is None:
+            raise RuntimeError(
+                f"New status of `{old_status}` and `{self.latest_attempt_id}` could not be determined from the message {msg}."
+            )
+        if new_status == old_status:
+            return
+        self.status = cast(RolloutStatus, new_status)
+
+        if is_finished(self):  # type: ignore
+            self.end_time = current_time
+        if is_queuing(self):  # type: ignore
+            self.enqueue_time = current_time
+            # When requeuing, we do not reset latest_attempt_id or num_attempts,
+            # as they should persist across requeues.
+
+    @classmethod
+    async def get_rollout_by_id(
+        cls: type[RolloutInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
+    ) -> Optional[Rollout | AttemptedRollout]:
+        """Query a specific rollout from the database."""
+        async with session_factory() as session:
+            async with session.begin():
+                rollout_obj = await session.get(cls, rollout_id)
+                if rollout_obj is None:
+                    return None
+                if rollout_obj.latest_attempt_id is not None:
+                    attempt_obj = await session.get(AttemptInDB, rollout_obj.latest_attempt_id)
+                    if attempt_obj is not None:
+                        return AttemptedRollout(
+                            **rollout_obj.as_rollout().model_dump(), attempt=attempt_obj.as_attempt()
+                        )
+                return rollout_obj.as_rollout()
+
+    @classmethod
+    async def query_rollouts(
+        cls: type[RolloutInDB],
+        session_factory: async_sessionmaker[AsyncSession],
+        *,
+        statuses: Optional[List[str]] = None,
+        ids: Optional[List[str]] = None,
+    ) -> List[RolloutInDB]:
+        """
+        Query rollouts from the database with optional filters.
+        """
+        async with session_factory() as session:
+            async with session.begin():
+                conditions: list[Any] = []
+                if statuses is not None:
+                    conditions.append(cls.status.in_(statuses))
+                if ids is not None:
+                    conditions.append(cls.rollout_id.in_(ids))
+                query = select(cls)
+                if conditions:
+                    query = query.where(and_(*conditions))
+                result = await session.scalars(query)
+                rollout_objs = result.all()
+                return list(rollout_objs)
@@ -0,0 +1,101 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from sqlalchemy import JSON, Float, Integer, String
+from sqlalchemy.orm import Mapped, mapped_column
+
+logger = logging.getLogger(__name__)
+
+from agentlightning.types.tracer import (
+    Attributes,
+    AttributeValue,
+    Event,
+    Link,
+    OtelResource,
+    Span,
+    SpanContext,
+    TraceStatus,
+)
+
+from .base import NamedDictBase, PydanticInDB, PydanticListInDB, SqlAlchemyBase
+
+
+class TraceStatusInDB(PydanticInDB):
+    target_type = TraceStatus
+
+
+class AttributesInDB(NamedDictBase):
+    target_alias = None  # type: ignore
+    value_type = AttributeValue
+
+
+class EventListInDB(PydanticListInDB):
+    value_type = Event
+
+
+class LinkListInDB(PydanticListInDB):
+    value_type = Link
+
+
+class SpanContextInDB(PydanticInDB):
+    target_type = SpanContext
+
+
+class OtelResourceInDB(PydanticInDB):
+    target_type = OtelResource
+
+
+class SpanInDB(SqlAlchemyBase):
+    __tablename__ = "spans"
+
+    rollout_id: Mapped[str] = mapped_column(String, nullable=False)  # The rollout which this span belongs to.
+    attempt_id: Mapped[str] = mapped_column(String, nullable=False)  # The attempt which this span belongs to.
+    sequence_id: Mapped[int] = mapped_column(
+        Integer, nullable=False
+    )  # The ID to make spans ordered within a single attempt.
+
+    # Current ID (in hex, formatted via trace_api.format_*)
+    trace_id: Mapped[str] = mapped_column(
+        String, nullable=False
+    )  # one rollout can have traces coming from multiple places
+
+    # FIXME: span_id may be not unique across different attempts/rollouts, use (rollout_id, attempt_id, sequence_id) as the primary key instead
+    span_id: Mapped[str] = mapped_column(
+        String, nullable=False
+    )  # The span ID of the span. This ID comes from the OpenTelemetry span ID generator.
+    parent_id: Mapped[Optional[str]] = mapped_column(String, nullable=True)  # The parent span ID of the span.
+
+    # Core ReadableSpan fields
+    name: Mapped[str] = mapped_column(String, nullable=False)
+    status: Mapped[TraceStatus] = mapped_column(TraceStatusInDB, nullable=False)
+    attributes: Mapped[Attributes] = mapped_column(AttributesInDB, nullable=False)
+    events: Mapped[List[Event]] = mapped_column(EventListInDB, nullable=False)
+    links: Mapped[List[Link]] = mapped_column(LinkListInDB, nullable=False)
+
+    # Timestamps
+    start_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
+    end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
+
+    # Other parsable fields
+    context: Mapped[Optional[SpanContext]] = mapped_column(SpanContextInDB, nullable=True)
+    parent: Mapped[Optional[SpanContext]] = mapped_column(SpanContextInDB, nullable=True)
+    resource: Mapped[OtelResource] = mapped_column(OtelResourceInDB, nullable=False)
+
+    # extra fields can be added here as needed
+    extra: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=None)
+
+    __mapper_args__ = {
+        "primary_key": [rollout_id, attempt_id, sequence_id],
+    }
+
+    def as_span(self) -> Span:
+        return Span(
+            **self.model_dump(
+                exclude={"extra"},
+                mapper={"*": lambda obj: obj.extra or {}},  # type: ignore
+            )
+        )
@@ -0,0 +1,316 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This file contains a configurable async retry decorator based on exception type."""
+
+from __future__ import annotations
+
+import functools
+import importlib
+import logging
+import random
+from dataclasses import asdict, dataclass
+from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Optional, Type, TypeVar
+
+from tenacity import AsyncRetrying, RetryCallState, retry_if_exception
+
+# ----------------------------------------------------------------------
+# Logging setup
+# ----------------------------------------------------------------------
+logger = logging.getLogger("async_retry")
+logging.basicConfig(level=logging.INFO)
+
+# ----------------------------------------------------------------------
+# Type alias for async callable
+# ----------------------------------------------------------------------
+F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
+
+
+# ----------------------------------------------------------------------
+# Dataclass definition for retry configuration
+# ----------------------------------------------------------------------
+@dataclass
+class RetryStrategy:
+    """Configuration schema for retry behavior of a specific exception type.
+    The wait time before $n$-th retry is calculated as ($n$ starts from 1):
+        wait_time = wait_seconds * (backoff ** (n - 1)) * (1 + jitter * U(-1, 1))
+    where U(-1, 1) is a uniform random variable between -1 and 1.
+    Attributes:
+        max_attempts: Maximum number of attempts before giving up. Default is 1 (no retry). None means infinite retries.
+        max_retry_delay: Optional maximum delay between retries in seconds. Default is None (no limit).
+        wait_seconds: Base wait time in seconds before the first retry. Default is 0.0.
+        max_wait_seconds: Maximum wait time in seconds between retries. Default is None (no limit).
+        backoff: Exponential backoff multiplier. Default is 1.0 (no backoff).
+        jitter: Fractional (relative) jitter to apply to wait time. Default is 0.0 (no jitter).
+        log: Whether to log each retry attempt. Default is False.
+    """
+
+    max_attempts: Optional[int] = 1
+    max_retry_delay: Optional[float] = None
+    wait_seconds: float = 0.0
+    max_wait_seconds: Optional[float] = None
+    backoff: float = 1.0
+    jitter: float = 0.0
+    log: bool = False
+
+    def asdict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    def __post_init__(self):
+        if self.max_attempts is not None and self.max_attempts < 1:
+            raise ValueError("max_attempts must be at least 1 or None for infinite retries")
+        if self.wait_seconds < 0.0:
+            raise ValueError("wait_seconds must be non-negative")
+        if self.backoff < 1.0:
+            raise ValueError("backoff must be at least 1.0")
+        if not (0.0 <= self.jitter <= 1.0):
+            raise ValueError("jitter must be between 0.0 and 1.0")
+
+    def _get_wait_time(self, attempt_number: int) -> float:
+        """Calculate the wait time before the given attempt number."""
+        base_wait = self.wait_seconds * (self.backoff ** (attempt_number - 1))
+        if self.jitter > 0:
+            delta = base_wait * self.jitter
+            wait_time = random.uniform(base_wait - delta, base_wait + delta)
+        else:
+            wait_time = base_wait
+        wait_time = max(wait_time, 0.0)
+        if self.max_wait_seconds is not None:
+            wait_time = min(wait_time, self.max_wait_seconds)
+        return wait_time
+
+    def wait_func(self, retry_state: RetryCallState) -> float:
+        """Tenacity wait function based on the given strategy."""
+        return self._get_wait_time(retry_state.attempt_number)
+
+    def stop_func(self, retry_state: RetryCallState) -> bool:
+        """Tenacity stop function based on the given strategy."""
+        if self.max_attempts is not None:
+            if retry_state.attempt_number >= self.max_attempts:
+                return True
+        if self.max_retry_delay is not None:
+            time_since_start = retry_state.seconds_since_start
+            if time_since_start is None:
+                logger.warning("Cannot determine time since start for retry stop condition.")
+                return False
+            if time_since_start >= self.max_retry_delay:
+                return True
+        return False
+
+    async def before_sleep(self, retry_state: RetryCallState):
+        """Tenacity before_sleep callback to log retry attempts."""
+        if self.log:
+            exc = retry_state.outcome.exception() if retry_state.outcome else None
+            next_wait = self.wait_func(retry_state)
+            logger.warning(
+                f"[Retry] {exc.__class__.__name__}: attempt={retry_state.attempt_number}, "
+                f"next_wait={next_wait:.2f}s, message={exc}"
+            )
+
+
+# ----------------------------------------------------------------------
+# Exception Registry — shared, reusable, and extensible
+# ----------------------------------------------------------------------
+class ExceptionRegistry:
+    """
+    Global registry for mapping string keys to Exception classes.
+    Supports dynamic registration and fallback to importlib.
+    """
+
+    _registry: Dict[str, Type[BaseException]] = {}
+
+    @classmethod
+    def register(cls, name: str, exc_type: Type[BaseException] | None = None) -> None:
+        """Register an exception type under a given name."""
+        if name in cls._registry:
+            logger.warning(f"Overwriting existing exception registration for name '{name}'.")
+        if exc_type is None:
+            # Try to dynamically import the exception class
+            try:
+                module_name, class_name = name.rsplit(".", 1)
+                module = importlib.import_module(module_name)
+                exc_type = getattr(module, class_name)
+                if exc_type is None:
+                    raise TypeError(f"{name} is not an Exception type.")
+            except (ImportError, AttributeError, ValueError, TypeError) as e:
+                raise ValueError(f"Cannot resolve exception type for name '{name}': {e}")
+        cls._registry[name] = exc_type
+
+    @classmethod
+    def all_registered(cls) -> Dict[str, Type[BaseException]]:
+        """Return the current registry mapping."""
+        return dict(cls._registry)
+
+    @classmethod
+    def clear(cls):
+        """Clear all registered exception mappings."""
+        cls._registry.clear()
+
+
+# ----------------------------------------------------------------------
+# Async Retry Decorator
+# ----------------------------------------------------------------------
+class AsyncTypeBasedRetry:
+    """
+    A configurable async retry decorator based on exception type.
+
+    - Takes configuration as a Dict[str, RetryStrategy].
+    - Provides `from_json()` for quick loading.
+    - Uses a global ExceptionRegistry to resolve exception names.
+    """
+
+    def __init__(self, strategies: Dict[str, RetryStrategy], default_strategy: RetryStrategy | None = None):
+        self.exception_map = self._build_exception_map(strategies)
+        self.default_strategy = default_strategy or RetryStrategy()
+
+    # ------------------------------------------------------------------
+    # Build exception map
+    # ------------------------------------------------------------------
+    def _build_exception_map(self, strategies: Dict[str, RetryStrategy]) -> Dict[Type[BaseException], RetryStrategy]:
+        mapping: Dict[Type[BaseException], RetryStrategy] = {}
+        all_registered = ExceptionRegistry.all_registered()
+        for name, strat in strategies.items():
+            if name in all_registered:
+                exc_type = all_registered[name]
+            else:
+                raise ValueError(f"Exception type '{name}' is not registered in ExceptionRegistry.")
+            mapping[exc_type] = strat
+        return mapping
+
+    # ------------------------------------------------------------------
+    # Retry core logic
+    # ------------------------------------------------------------------
+    def get_exception(self, retry_state: RetryCallState) -> Optional[BaseException]:
+        """Get the exception from the given retry state, if any."""
+        return retry_state.outcome.exception() if retry_state.outcome else None
+
+    def get_strategy(self, retry_state: RetryCallState) -> Optional[RetryStrategy]:
+        """Get the RetryStrategy for the exception in the given retry state.
+        IF no matching exception type is found, return the default strategy.
+        IF no exception is found, return None.
+        """
+        exc = self.get_exception(retry_state)
+        if exc is None:
+            return None
+        for exc_type, strat in self.exception_map.items():
+            if isinstance(exc, exc_type):
+                return strat
+        return self.default_strategy
+
+    def should_retry(self, exc: BaseException) -> bool:
+        return any(isinstance(exc, t) for t in self.exception_map.keys())
+
+    def wait_func(self, retry_state: RetryCallState) -> float:
+        strat = self.get_strategy(retry_state)
+        if strat is None:
+            return 0.0
+        return strat.wait_func(retry_state)
+
+    def stop_func(self, retry_state: RetryCallState) -> bool:
+        strat = self.get_strategy(retry_state)
+        if strat is None:
+            return False
+        return strat.stop_func(retry_state)
+
+    async def before_sleep(self, retry_state: RetryCallState):
+        strat = self.get_strategy(retry_state)
+        if strat is None:
+            return
+        await strat.before_sleep(retry_state)
+
+    # ------------------------------------------------------------------
+    # Decorator entry point
+    # ------------------------------------------------------------------
+    def __call__(self, func: F) -> F:
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):  # type: ignore
+            async for attempt in AsyncRetrying(
+                retry=retry_if_exception(lambda e: self.should_retry(e)),
+                wait=self.wait_func,
+                stop=self.stop_func,
+                before_sleep=self.before_sleep,
+                reraise=True,
+            ):
+                with attempt:
+                    return await func(*args, **kwargs)
+
+        return wrapper  # type: ignore
+
+
+# ----------------------------------------------------------------------
+# A configurable async retrier for any code block
+# ----------------------------------------------------------------------
+
+
+class AsyncRetryBlock:
+    """
+    Async retry helper for a single exception type and strategy.
+
+    Usage:
+        async with AsyncRetryBlock(strategy):
+            await some_async_function()
+    """
+
+    def __init__(self, strategy: RetryStrategy, **retry_kwargs):  # type: ignore
+        self.strategy = strategy
+        self._retryer = AsyncRetrying(
+            wait=self._wait_func,
+            stop=self._stop_func,
+            before_sleep=self._before_sleep,
+            **retry_kwargs,  # type: ignore
+        )
+
+    async def run(self, coro: Callable[..., Awaitable[Any]]) -> Any:
+        """Run the given coroutine with retries according to the strategy.
+        For example:
+            async def my_coro():
+                ...
+            retry_block = AsyncRetryBlock(strategy)
+            result = await retry_block.run(my_coro)
+        """
+        async for attempt in self._retryer:
+            with attempt:
+                return await coro()
+
+    # ------------------------------------------------------------------
+    # Core: async iterator interface
+    # ------------------------------------------------------------------
+    def __aiter__(self) -> AsyncIterator[Any]:
+        """Return an async iterator that yields retry attempts.
+        Usage:
+            async for attempt in retry_block:
+                with attempt:
+                    await some_async_function()
+        """
+        return self._retryer.__aiter__()
+
+    # ------------------------------------------------------------------
+    # Context manager entry
+    # ------------------------------------------------------------------
+    async def __aenter__(self):
+        self._aiter = self._retryer.__aiter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):  # type: ignore
+        # Consume the retry iterator
+        try:
+            # If exception occurred, let the retryer handle it
+            async for attempt in self._aiter:
+                with attempt:
+                    if exc_val:
+                        raise exc_val
+        except Exception:
+            # Allow exception to propagate if retries exhausted
+            pass
+        return False
+
+    # ------------------------------------------------------------------
+    # Strategy function
+    # ------------------------------------------------------------------
+    def _wait_func(self, retry_state: RetryCallState) -> float:
+        return self.strategy.wait_func(retry_state)
+
+    def _stop_func(self, retry_state: RetryCallState) -> bool:
+        return self.strategy.stop_func(retry_state)
+
+    async def _before_sleep(self, retry_state: RetryCallState):
+        await self.strategy.before_sleep(retry_state)
@@ -0,0 +1,685 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Literal, Optional, Sequence, Union
+
+from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.triggers.interval import IntervalTrigger
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel
+from sqlalchemy import and_, or_, select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.orm.exc import StaleDataError
+from tenacity import RetryError
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+)
+
+from ..base import UNSET, LightningStore, Unset, is_finished
+from .orm import (
+    AttemptInDB,
+    ResourcesUpdateInDB,
+    RolloutInDB,
+    SpanInDB,
+    SpanSeqIdInDB,
+    SqlAlchemyBase,
+)
+from .retry_helper import AsyncRetryBlock, AsyncTypeBasedRetry, ExceptionRegistry, RetryStrategy
+
+logger = logging.getLogger(__name__)
+
+# TODO add periodic cleanup of old rollouts/attempts/spans
+
+ExceptionRegistry.register("sqlalchemy.orm.exc.StaleDataError")
+ExceptionRegistry.register("sqlalchemy.exc.OperationalError")
+
+db_retry = AsyncTypeBasedRetry(
+    {
+        "sqlalchemy.exc.OperationalError": RetryStrategy(
+            max_attempts=5, wait_seconds=1, backoff=1.5, jitter=0.3, log=True
+        ),
+        "sqlalchemy.orm.exc.StaleDataError": RetryStrategy(
+            max_attempts=100, wait_seconds=1e-3, backoff=1.0, jitter=0.1, log=True
+        ),
+    }
+)
+
+
+class _WaitForRolloutsCompleted(Exception):
+    """Internal exception to signal that not all rollouts have completed yet."""
+
+    pass
+
+
+class BackgroundTaskConfig(BaseModel):
+    name: str  # unique name for the task
+    method: str  # method name to call, currently only supports methods of SqlLightningStore
+    interval: Dict[Literal["seconds", "minutes", "hours"], float]  # interval for the task
+    is_async: bool = True  # whether the task method is async, default to True
+
+
+class SqlLightningStore(LightningStore):
+    """
+    A LightningStore implementation that uses a database backend to store and manage rollouts and attempts.
+    The database backend is expected to support asynchronous operations.
+    The store uses SQLAlchemy ORM models to interact with the database
+    Args:
+        database_url (string):
+            The database URL for connecting to the database.
+            If None, will read from the 'DATABASE_URL' environment variable.
+        retry_for_waiting (RetryStrategy):
+            Retry strategy for polling when waiting for rollouts to complete.
+            If None, a default strategy will be used.
+        wait_for_nonexistent_rollout (Bool):
+            If True, when waiting for rollouts, will wait for all specified rollouts to complete, including non-existing ones.
+            If False, will ignore non-existing rollouts as completed. (Default: False)
+        background_tasks_cfg (list[Dict[str, Any]]):
+            The configuration for in-process periodic tasks, following the definition of `BackgroundTaskConfig`.
+            IF not provided (None as default), the dbstore will incorporate a default set of periodic tasks as follows:
+            [
+                BackgroundTaskConfig(name="check_attempt_timeout", method="check_attempt_timeout", interval={"seconds": 10.0}),
+            ]
+            To disable all periodic tasks, provide an empty list `[]`.
+    Note:
+        Explicitly use async `start()` and `stop()` methods to manage the database connection lifecycle.
+    """
+
+    def __init__(
+        self,
+        database_url: Optional[str] = None,
+        *,
+        retry_for_waiting: Optional[dict[str, Any] | RetryStrategy] = None,
+        wait_for_nonexistent_rollout: bool = False,
+        background_tasks_cfg: list[Dict[str, Any]] | None = None,
+    ) -> None:
+        super().__init__()
+        if database_url is None:
+            database_url = os.getenv("DATABASE_URL", None)
+        if database_url is None:
+            raise ValueError(
+                "A database URL must be provided either via the 'database_url' parameter or the 'DATABASE_URL' environment variable."
+            )
+
+        self._engine = create_async_engine(database_url, echo=False)
+        self._async_session = async_sessionmaker(self._engine, expire_on_commit=False)
+
+        self._latest_resources_id = None
+
+        # special handling for retry strategy
+        retry_for_waiting = retry_for_waiting or RetryStrategy(
+            max_attempts=10,  # set a limit for retries if timeout is specified, otherwise will change to None later
+            max_retry_delay=None,  # set later
+            wait_seconds=10.0,  # poll every 10 seconds
+            max_wait_seconds=60.0,  # at most wait 60 seconds between retries
+            backoff=1.0,
+            jitter=0.0,
+            log=True,
+        )
+        self.retry_for_waiting = (
+            retry_for_waiting if isinstance(retry_for_waiting, RetryStrategy) else RetryStrategy(**retry_for_waiting)
+        )
+        self.wait_for_nonexistent_rollout = wait_for_nonexistent_rollout
+
+        # setup in-process periodic tasks
+        if background_tasks_cfg is None:
+            self.background_tasks_cfg = [
+                BackgroundTaskConfig(
+                    name="check_attempt_timeout", method="check_attempt_timeout", interval={"seconds": 10.0}
+                ),
+            ]
+        else:
+            self.background_tasks_cfg = [BackgroundTaskConfig(**cfg) for cfg in background_tasks_cfg]
+        self._background_scheduler = BackgroundScheduler()
+
+    async def start(self):
+        async with self._engine.begin() as conn:
+            await conn.run_sync(SqlAlchemyBase.metadata.create_all)
+        for task_cfg in self.background_tasks_cfg:
+            self.add_background_task(task_cfg, to_scheduler_only=True)
+        self._background_scheduler.start()  # type: ignore
+
+    async def stop(self):
+        await self._engine.dispose()
+        self._background_scheduler.shutdown()  # type: ignore
+
+    def add_background_task(
+        self, task_cfg: Dict[str, Any] | BackgroundTaskConfig, to_scheduler_only: bool = False
+    ) -> None:
+        """Add a new periodic background task to the scheduler.
+        Args:
+            task_cfg (Dict[str, Any] | BackgroundTaskConfig): The configuration for the background task.
+            to_scheduler_only (bool): If True, only add the task to the scheduler without updating the configuration list.
+        Raises:
+            ValueError: If the task method is not defined in SqlLightningStore.
+        """
+        config = task_cfg if isinstance(task_cfg, BackgroundTaskConfig) else BackgroundTaskConfig(**task_cfg)
+        if not to_scheduler_only:
+            # check existing tasks
+            for existing in self.background_tasks_cfg:
+                if existing.name == config.name:
+                    logger.warning(
+                        f"Background task {config.name} is already scheduled, will update its configuration."
+                    )
+            self.background_tasks_cfg.append(config)
+        delta_t = timedelta(**config.interval)
+        if not hasattr(self, config.method):
+            raise ValueError(f"Periodic task method {config.method} is not defined in SqlLightningStore.")
+        if config.is_async:
+            func = lambda: asyncio.run(getattr(self, config.method)())
+        else:
+            func = lambda: getattr(self, config.method)()
+
+        self._background_scheduler.add_job(  # type: ignore
+            func=func,
+            trigger=IntervalTrigger(**config.interval),  # type: ignore
+            name=f"SqlLightningStore.{config.name}",
+            replace_existing=True,
+            next_run_time=datetime.now() + delta_t,  # schedule the first run after the interval
+        )
+
+    # ------------------------------------------------------
+    # Public methods defined in LightningStore
+    # ------------------------------------------------------
+
+    @db_retry
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        async with self._async_session() as session:
+            async with session.begin():
+                rollout_obj = RolloutInDB(
+                    input=input,
+                    mode=mode,
+                    resources_id=resources_id or self._latest_resources_id,
+                    status="queuing",
+                    config=config or RolloutConfig(),
+                    rollout_metadata=metadata,
+                )
+                session.add(rollout_obj)
+                attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
+                await session.flush()  # ensure the object is written to the DB
+                return attempted_rollout
+
+    @db_retry
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        async with self._async_session() as session:
+            async with session.begin():
+                rollout_obj = RolloutInDB(
+                    input=input,
+                    mode=mode,
+                    resources_id=resources_id or self._latest_resources_id,
+                    status="queuing",
+                    config=config or RolloutConfig(),
+                    rollout_metadata=metadata,
+                )
+                session.add(rollout_obj)
+                await session.flush()  # ensure the object is written to the DB
+                return rollout_obj.as_rollout()
+
+    @db_retry
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        return await self._fifo_dequeue_rollout()
+
+    @db_retry
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        async with self._async_session() as session:
+            async with session.begin():
+                rollout_obj = await session.get(RolloutInDB, rollout_id)
+                if rollout_obj is None:
+                    raise ValueError(f"Rollout {rollout_id} not found")
+                attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
+                await session.flush()  # ensure the object is written to the DB
+                return attempted_rollout
+
+    @db_retry
+    async def add_span(self, span: Span) -> Span:
+        seq_id = await SpanSeqIdInDB.get_next_sequence_id(self._async_session, span.rollout_id, span.attempt_id)
+        return await self._add_span(span.model_dump(), seq_id=seq_id)
+
+    @db_retry
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        sequence_id = await SpanSeqIdInDB.get_next_sequence_id(self._async_session, rollout_id, attempt_id, sequence_id)
+        span = Span.from_opentelemetry(
+            src=readable_span,
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+        )
+        return await self._add_span(span.model_dump(), seq_id=sequence_id)
+
+    @db_retry
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[Rollout]:
+        rollouts = await RolloutInDB.query_rollouts(self._async_session, statuses=status, ids=rollout_ids)  # type: ignore
+        attempt_ids = [r.latest_attempt_id for r in rollouts if r.latest_attempt_id is not None]
+        async with self._async_session() as session:
+            async with session.begin():
+                scalars = await session.scalars(select(AttemptInDB).where(AttemptInDB.attempt_id.in_(attempt_ids)))
+                attempts = scalars.all()
+                attempt_map = {a.attempt_id: a.as_attempt() for a in attempts}
+                return [
+                    (
+                        AttemptedRollout(**r.as_rollout().model_dump(), attempt=attempt_map[r.latest_attempt_id])
+                        if r.latest_attempt_id in attempt_map
+                        else r.as_rollout()
+                    )
+                    for r in rollouts
+                ]  # type: ignore
+
+    @db_retry
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        return await AttemptInDB.get_attempts_for_rollout(self._async_session, rollout_id)  # type: ignore
+
+    @db_retry
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
+        return await RolloutInDB.get_rollout_by_id(self._async_session, rollout_id)
+
+    @db_retry
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        return await AttemptInDB.get_latest_attempt_for_rollout(self._async_session, rollout_id)
+
+    @db_retry
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        return await ResourcesUpdateInDB.get_resources_by_id(self._async_session, resources_id)
+
+    @db_retry
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        if self._latest_resources_id is None:
+            return None
+        return await ResourcesUpdateInDB.get_resources_by_id(self._async_session, self._latest_resources_id)
+
+    @db_retry
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        return await SpanSeqIdInDB.get_next_sequence_id(self._async_session, rollout_id, attempt_id)
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        # implementation the timeout via tenacity retry mechanism, by a `with` context
+        strategy = RetryStrategy(**self.retry_for_waiting.asdict())
+        if timeout is not None:
+            strategy.max_retry_delay = timeout
+            if strategy.max_attempts is not None:
+                strategy.wait_seconds = min(strategy.wait_seconds, timeout / (strategy.max_attempts + 1))
+        else:
+            strategy.max_attempts = None  # infinite retries
+
+        non_completed_ids, non_existing_ids = set(rollout_ids), set(rollout_ids)
+        completed_rollouts: Dict[str, Rollout] = {}
+        if len(non_completed_ids) < len(rollout_ids):
+            logger.warning("Duplicate rollout_ids found in wait_for_rollouts input. Duplicates will be ignored.")
+
+        try:
+            async for attempt in AsyncRetryBlock(
+                strategy,
+                reraise=True,
+            ):
+                with attempt:
+                    async with self._async_session() as session:
+                        async with session.begin():
+                            result = await session.scalars(
+                                select(RolloutInDB).where(RolloutInDB.rollout_id.in_(non_completed_ids))
+                            )
+                            rollouts = [r.as_rollout() for r in result.all()]
+                            for r in rollouts:
+                                if r.rollout_id in non_existing_ids:
+                                    non_existing_ids.discard(r.rollout_id)  # found existing rollout
+                                if is_finished(r):
+                                    completed_rollouts[r.rollout_id] = r
+                                    non_completed_ids.discard(r.rollout_id)
+                            # check termination conditions
+                            if self.wait_for_nonexistent_rollout:
+                                if len(non_completed_ids) == 0:
+                                    return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
+                                raise _WaitForRolloutsCompleted(
+                                    f"WaitForRolloutsCompleted: requested={len(rollout_ids)}, completed={len(completed_rollouts)}, non_existing={len(non_existing_ids)}"
+                                )
+                            else:
+                                if len(non_completed_ids) == len(non_existing_ids):
+                                    logger.warning(f"All remaining rollouts are non-existing: {non_existing_ids}.")
+                                    return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
+                                raise _WaitForRolloutsCompleted(
+                                    f"WaitForRolloutsCompleted: requested={len(rollout_ids)}, completed={len(completed_rollouts)}, non_existing={len(non_existing_ids)}"
+                                )
+
+        except (RetryError, _WaitForRolloutsCompleted):
+            return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
+        except Exception as e:
+            logger.error(f"Error while waiting for rollouts: {e}")
+            raise e
+
+        # Ensure a return value in case no rollouts are completed
+        return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
+
+    @db_retry
+    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+        async with self._async_session() as session:
+            async with session.begin():
+                conditions: List[Any] = [SpanInDB.rollout_id == rollout_id]
+                if attempt_id is not None:
+                    if attempt_id == "latest":
+                        rollout_obj = await session.get(RolloutInDB, rollout_id)
+                        if rollout_obj is None:
+                            logger.warning(f"Rollout {rollout_id} does not exist. Cannot query latest attempt spans.")
+                            return []
+                        attempt_id = rollout_obj.latest_attempt_id
+                    conditions.append(SpanInDB.attempt_id == attempt_id)
+                query = select(SpanInDB).where(and_(*conditions)).order_by(SpanInDB.sequence_id.asc())
+                result = await session.scalars(query)
+                span_objs = result.all()
+                return [obj.as_span() for obj in span_objs]
+
+    @db_retry
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        async with self._async_session() as session:
+            async with session.begin():
+                current_time = time.time()
+                resource_obj = ResourcesUpdateInDB(
+                    resources=resources,
+                    create_time=current_time,
+                    update_time=current_time,
+                )
+                session.add(resource_obj)
+                await session.flush()  # ensure the object is written to the DB
+                self._latest_resources_id = resource_obj.resources_id
+                return resource_obj.as_resources_update()
+
+    @db_retry
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        async with self._async_session() as session:
+            async with session.begin():
+                obj = await session.get(ResourcesUpdateInDB, resources_id)
+                if obj is None:
+                    # raise ValueError(f"Failed to update resources {resources_id}. It may not exist.")
+                    # FIXME InMemoryLightningStore will create the resources if not exist, but the base method require to raise error
+                    # HACK here stick to the behavior of InMemoryLightningStore for compatibility
+                    current_time = time.time()
+                    obj = ResourcesUpdateInDB(
+                        resources_id=resources_id,
+                        resources=resources,
+                        create_time=current_time,
+                        update_time=current_time,
+                    )
+                    session.add(obj)
+                else:
+                    obj.resources = resources
+                await session.flush()
+                self._latest_resources_id = resources_id
+                return obj.as_resources_update()
+
+    @db_retry
+    async def query_resources(self) -> List[ResourcesUpdate]:
+        async with self._async_session() as session:
+            async with session.begin():
+                result = await session.scalars(
+                    select(ResourcesUpdateInDB).order_by(ResourcesUpdateInDB.create_time.asc())
+                )
+                resource_objs = result.all()
+                return [obj.as_resources_update() for obj in resource_objs]
+
+    @db_retry
+    async def update_rollout(
+        self,
+        rollout_id: str | None,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        if rollout_id is None:
+            raise ValueError("rollout_id must be provided for updating a rollout.")
+
+        async with self._async_session() as session:
+            async with session.begin():
+                rollout_obj = await session.get(RolloutInDB, rollout_id)
+                if rollout_obj is None:
+                    raise ValueError(f"Rollout {rollout_id} not found")
+                # udpate fields
+                if not isinstance(input, Unset):
+                    rollout_obj.input = input
+                if not isinstance(mode, Unset):
+                    rollout_obj.mode = mode
+                if not isinstance(resources_id, Unset):
+                    rollout_obj.resources_id = resources_id
+                if not isinstance(status, Unset):
+                    await rollout_obj.update_status(dict(event="user_update", new_status=status))
+                if not isinstance(config, Unset):
+                    rollout_obj.config = config
+                if not isinstance(metadata, Unset):
+                    rollout_obj.rollout_metadata = metadata
+                await session.flush()  # ensure the object is written to the DB
+                return rollout_obj.as_rollout()
+
+    @db_retry
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        async with self._async_session() as session:
+            async with session.begin():
+                rollout_obj = await session.get(RolloutInDB, rollout_id)
+                if rollout_obj is None:
+                    raise ValueError(f"Rollout {rollout_id} not found")
+                if attempt_id == "latest":
+                    if rollout_obj.latest_attempt_id is None:
+                        raise ValueError(f"Rollout {rollout_id} has no attempts. Cannot update latest attempt.")
+                    attempt_id = rollout_obj.latest_attempt_id
+                if attempt_id != rollout_obj.latest_attempt_id:
+                    logger.warning(
+                        f"Updating attempt {attempt_id} which is not the latest attempt for rollout {rollout_id}. Latest is {rollout_obj.latest_attempt_id}."
+                    )
+                attempt_obj = await session.get(AttemptInDB, attempt_id)
+                if attempt_obj is None:
+                    raise ValueError(f"No attempts found")
+                if attempt_obj.rollout_id != rollout_id:
+                    raise ValueError(f"Attempt {attempt_id} does not belong to rollout {rollout_id}.")
+                # update fields
+                if not isinstance(status, Unset):
+                    msg = attempt_obj.update_status(dict(event="user_update", new_status=status))
+                    if msg is not None:
+                        await rollout_obj.update_status(msg)
+                if not isinstance(worker_id, Unset):
+                    attempt_obj.worker_id = worker_id
+                if not isinstance(last_heartbeat_time, Unset):
+                    attempt_obj.last_heartbeat_time = last_heartbeat_time
+                if not isinstance(metadata, Unset):
+                    attempt_obj.attempt_metadata = metadata
+                await session.flush()  # ensure the object is written to the DB
+                return attempt_obj.as_attempt()
+
+    # ------------------------------------------------------
+    # periodic background tasks can be added here
+    # ------------------------------------------------------
+
+    async def check_attempt_timeout(self):
+        """Periodically check for attempts that have timed out and update their status accordingly."""
+        # use update with where condition to find and update timed-out attempts
+        current_time = time.time()
+
+        timed_out_results = await self._attempt_timeout_check(current_time)
+
+        # TODO run the tasks with a wrapper with asyncio semaphore to limit concurrency and handle exceptions
+        tasks = [self._process_timed_out_attempt(attempt, current_time) for attempt in timed_out_results]
+        await asyncio.gather(*tasks)
+
+    async def _process_timed_out_attempt(self, attempt_ref: AttemptInDB, current_time: float) -> None:
+        async with self._async_session() as session:
+            async with session.begin():
+                # Step 1: Update attempt status
+                attempt_obj = await session.get(
+                    AttemptInDB, attempt_ref.attempt_id
+                )  # refresh the object in the new session
+                if attempt_obj is None:
+                    raise ValueError(f"Attempt {attempt_ref.attempt_id} not found during timeout processing")
+                if attempt_obj.version_id != attempt_ref.version_id:
+                    # version mismatch, skip processing to avoid race conditions
+                    raise StaleDataError(f"Attempt {attempt_ref.attempt_id} version mismatch during timeout processing")
+                msg = {}
+                if attempt_obj.is_timed_out(current_time):
+                    msg = dict(event="overall_timeout", timestamp=current_time)
+                elif attempt_obj.is_unresponsive(current_time):
+                    msg = dict(event="single_step_timeout", timestamp=current_time)
+                else:
+                    raise ValueError(f"Attempt {attempt_ref.attempt_id} is not timed out during timeout processing")
+                msg2rollout = attempt_obj.update_status(msg)
+                if msg2rollout is None:
+                    return  # no further update needed
+
+                # Step 2: Update rollouts
+                rollout_obj = await session.get(RolloutInDB, attempt_obj.rollout_id)
+                if rollout_obj is None:
+                    raise ValueError(f"Rollout {attempt_obj.rollout_id} not found during timeout processing")
+                await rollout_obj.update_status(msg2rollout)
+
+    # ------------------------------------------------------
+    # internal helper methods can be added here
+    # ------------------------------------------------------
+
+    async def _add_span(self, span: Dict[str, Any], seq_id: Optional[int] = None) -> Span:
+        """Add a new span to the database."""
+        if seq_id is not None:
+            span["sequence_id"] = seq_id
+        extra_dic: Dict[str, Any] = {}
+        for k in list(span.keys()):
+            if k not in SpanInDB.__table__.columns.keys():
+                extra_dic[k] = span.pop(k)
+        span["extra"] = extra_dic if extra_dic else None
+
+        async with self._async_session() as session:
+            async with session.begin():
+                # create SpanInDB object
+                span_obj = SpanInDB(**span)
+                session.add(span_obj)
+                # update attempt's last_heartbeat_time and status
+                attempt_obj = await session.get(AttemptInDB, span["attempt_id"])
+                if attempt_obj is None:
+                    raise ValueError(f"Attempt {span['attempt_id']} not found")
+                # ensure the attempt and rollout are in running status
+                msg = attempt_obj.update_status(dict(event="span_received"))
+                if msg is not None:
+                    rollout_obj = await session.get(RolloutInDB, attempt_obj.rollout_id)
+                    if rollout_obj is None:
+                        raise ValueError(f"Rollout {attempt_obj.rollout_id} not found")
+                    await rollout_obj.update_status(msg)
+                await session.flush()  # ensure the object is written to the DB
+                return span_obj.as_span()
+
+    async def _fifo_dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        """Dequeue the next rollout in FIFO order (the one with the earliest enqueue_time).
+        Returns the RolloutInDB object if found, else None.
+        Note: This method does not update the status of the rollout. The caller should handle that.
+        """
+        async with self._async_session() as session:
+            async with session.begin():
+                # use the update...returning to atomically select the next rollout and claim it by updating its status to 'preparing'
+                result = await session.scalars(
+                    select(RolloutInDB)
+                    .where(RolloutInDB.status.in_(["queuing", "requeuing"]), RolloutInDB.enqueue_time.isnot(None))
+                    .order_by(RolloutInDB.enqueue_time.asc())
+                    .limit(1)
+                )
+                rollout_obj = result.one_or_none()
+                if rollout_obj is None:
+                    return None  # no rollout available
+                # update the status of the rollout to 'preparing' via Compare-and-Swap to avoid race
+                attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
+                await session.flush()  # ensure the object is written to the DB
+                return attempted_rollout
+
+    async def _start_attempt_for_rollout(self, session: AsyncSession, rollout_obj: RolloutInDB) -> AttemptedRollout:
+        """Create a new attempt for the given rollout and update the rollout's fields."""
+        # create a new attempt for this rollout
+        rollout_config = rollout_obj.config
+        attempt_obj = AttemptInDB(
+            rollout_id=rollout_obj.rollout_id,
+            sequence_id=rollout_obj.num_attempts + 1,
+            status="preparing",
+            max_duration=rollout_config.timeout_seconds,
+            max_heartbeat_interval=rollout_config.unresponsive_seconds,
+        )
+        session.add(attempt_obj)
+        # pre-update the rollout_obj fields for CAS
+        rollout_obj.status = attempt_obj.status  # type: ignore pre-update the status in the object for CAS
+        rollout_obj.enqueue_time = None  # pre-update the enqueue_time in the object for CAS
+        rollout_obj.num_attempts += 1  # pre-update the num_attempts in the object for CAS
+        rollout_obj.latest_attempt_id = attempt_obj.attempt_id  # pre-update the latest_attempt_id in the object for CAS
+
+        # create a sequence id tracker for each attempt
+        # FIXME currently InMemoryLightningStore let all attempts under the same rollout share the same span sequence for sorting
+        # create a sequence id tracker for this rollout, only if not exists
+        existing = await session.get(SpanSeqIdInDB, rollout_obj.rollout_id)
+        if existing is None:
+            seq_obj = SpanSeqIdInDB(
+                rollout_id=rollout_obj.rollout_id,
+                attempt_id=attempt_obj.attempt_id,
+            )
+            session.add(seq_obj)
+
+        return AttemptedRollout(**rollout_obj.as_rollout().model_dump(), attempt=attempt_obj.as_attempt())
+
+    async def _attempt_timeout_check(self, now: float) -> Sequence[AttemptInDB]:
+        """Scan the table for attempts that have timed out based on the given mode, and return them for further processing.
+        Returns:
+            list[AttemptInDB]:
+                A list of AttemptInDB objects that timed out.
+        """
+        async with self._async_session() as session:
+            async with session.begin():
+                scalars = await session.scalars(
+                    select(AttemptInDB).where(
+                        and_(
+                            AttemptInDB.status.in_(["preparing", "running"]),
+                            or_(
+                                and_(
+                                    AttemptInDB.max_duration.isnot(None),
+                                    (now - AttemptInDB.start_time) > AttemptInDB.max_duration,
+                                ),
+                                and_(
+                                    AttemptInDB.max_heartbeat_interval.isnot(None),
+                                    (now - AttemptInDB.last_heartbeat_time) > AttemptInDB.max_heartbeat_interval,
+                                ),
+                            ),
+                        )
+                    )
+                )
+                return scalars.all()
@@ -0,0 +1,182 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict, List, Literal, Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+)
+
+from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
+
+
+class LightningStoreThreaded(LightningStore):
+    """Facade that delegates all store operations to a underlying store instance.
+
+    The operations are guaranteed to be thread-safe.
+    Make sure the threaded stores are instantiated before initializing the threads.
+    """
+
+    def __init__(self, store: LightningStore) -> None:
+        super().__init__()  # watchdog relies on the underlying store
+        self.store = store
+        self._lock = threading.Lock()
+
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        capabilities = self.store.capabilities()
+        return {
+            **capabilities,
+            "async_safe": True,
+            "thread_safe": True,
+        }
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_rollout(input, mode, resources_id, config, metadata)
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        with self._lock:
+            return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        with self._lock:
+            return await self.store.dequeue_rollout()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_attempt(rollout_id)
+
+    async def query_rollouts(
+        self,
+        *,
+        status: Optional[Sequence[RolloutStatus]] = None,
+        rollout_ids: Optional[Sequence[str]] = None,
+    ) -> List[Rollout]:
+        with self._lock:
+            return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        with self._lock:
+            return await self.store.query_attempts(rollout_id)
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        with self._lock:
+            return await self.store.get_rollout_by_id(rollout_id)
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        with self._lock:
+            return await self.store.get_latest_attempt(rollout_id)
+
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        with self._lock:
+            return await self.store.add_resources(resources)
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        with self._lock:
+            return await self.store.update_resources(resources_id, resources)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.get_resources_by_id(resources_id)
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.get_latest_resources()
+
+    async def add_span(self, span: Span) -> Span:
+        with self._lock:
+            return await self.store.add_span(span)
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        with self._lock:
+            return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        # This method does not change the state of the store, and it's not thread-safe.
+        return await self.store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        with self._lock:
+            return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        with self._lock:
+            return await self.store.query_spans(rollout_id, attempt_id)
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        with self._lock:
+            return await self.store.update_rollout(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                status=status,
+                config=config,
+                metadata=metadata,
+            )
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        with self._lock:
+            return await self.store.update_attempt(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                status=status,
+                worker_id=worker_id,
+                last_heartbeat_time=last_heartbeat_time,
+                metadata=metadata,
+            )
@@ -0,0 +1,127 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import time
+from typing import Awaitable, Callable, List, cast
+
+from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
+
+UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[Rollout]]
+UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]
+
+
+async def propagate_status(
+    update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
+    attempt: Attempt,
+    config: RolloutConfig,
+) -> Rollout:
+    """
+    Propagate the status of an attempt to the rollout.
+
+    The rollout should be made sure in a state to be outdated.
+    Requeue the rollout if it should be retried.
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
+    """
+    # Propagate the status directly to the rollout
+    if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
+        return await update_rollout_status(
+            attempt.rollout_id,
+            attempt.status,
+        )
+
+    if attempt.status == "failed" or attempt.status == "timeout" or attempt.status == "unresponsive":
+        # Check if this status should trigger a retry
+        if attempt.status in config.retry_condition:
+            # If we haven't exceeded max attempts, retry
+            if attempt.sequence_id < config.max_attempts:
+                return await update_rollout_status(
+                    attempt.rollout_id,
+                    "requeuing",
+                )
+
+        # If we can't retry or shouldn't retry, mark as failed
+        return await update_rollout_status(
+            attempt.rollout_id,
+            "failed",
+        )
+
+    raise ValueError(f"Invalid attempt status: {attempt.status}")
+
+
+async def healthcheck(
+    rollouts: List[AttemptedRollout],
+    update_rollout_status: UpdateRolloutStatus,
+    update_attempt_status: UpdateAttemptStatus,
+) -> None:
+    """
+    Perform health check on all running rollouts in the store.
+
+    This method should be called periodically to:
+
+    1. Update rollout status to failed to succeeded when the attempt is done
+    2. Check for unresponsive attempts (no heartbeat or spans for a while)
+    3. Check for timed-out rollouts (running too long since start_time)
+    4. Update attempt/rollout status accordingly
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
+
+    Args:
+        store: The LightningStore instance to check rollouts from
+    """
+    current_time = time.time()
+
+    for rollout in rollouts:
+        config = rollout.config  # policy for retry and timeout
+
+        # Get the latest attempt for this rollout
+        latest_attempt = rollout.attempt
+        if not latest_attempt:
+            continue
+
+        # Check if the attempt has already failed or succeeded
+        if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
+            await propagate_status(update_rollout_status, latest_attempt, config)
+            continue
+
+        # Check for timeout condition (based on attempt start_time, instead of rollout start_time)
+        if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "timeout",
+            )
+            continue
+
+        # Check for unresponsive condition (based on last heartbeat)
+        if latest_attempt.last_heartbeat_time:
+            if latest_attempt.status == "preparing":
+                # If still preparing, mark it as running
+                latest_attempt = await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "running",
+                )
+
+            # Haven't received heartbeat for a while
+            if (
+                config.unresponsive_seconds is not None
+                and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
+            ):
+                await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "unresponsive",
+                )
+                continue
+
+        # Check if there's no last heartbeat (no spans) at all
+        if (
+            latest_attempt.last_heartbeat_time is None
+            and config.unresponsive_seconds is not None
+            and current_time - latest_attempt.start_time > config.unresponsive_seconds
+        ):
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "unresponsive",
+            )
@@ -1,3 +1,7 @@
-from .base import BaseTracer
+# Copyright (c) Microsoft. All rights reserved.
+
 from .agentops import AgentOpsTracer
-from .triplet import TripletExporter
+from .base import Tracer
+from .otel import OtelTracer
+
+__all__ = ["AgentOpsTracer", "Tracer", "OtelTracer"]
@@ -1,20 +1,26 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 from __future__ import annotations

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

-import agentops.sdk.core
 import agentops
+import agentops.sdk.core
 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.trace.status import StatusCode

-from agentlightning.instrumentation.agentops import AgentOpsServerManager
 from agentlightning.instrumentation import instrument_all, uninstrument_all
-from .base import BaseTracer
+from agentlightning.store.base import LightningStore

+from .base import Tracer

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


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

    This tracer provides functionality to capture execution details using the
@@ -50,46 +56,11 @@ class AgentOpsTracer(BaseTracer):
        self.instrument_managed = instrument_managed
        self.daemon = daemon

-        self._agentops_server_manager = AgentOpsServerManager(self.daemon)
-        self._agentops_server_port_val: Optional[int] = None
-
        if not self.agentops_managed:
            logger.warning("agentops_managed=False. You are responsible for AgentOps setup.")
        if not self.instrument_managed:
            logger.warning("instrument_managed=False. You are responsible for all instrumentation.")

-    def __getstate__(self):
-        state = self.__dict__.copy()
-        state["_agentops_server_manager"] = None  # Exclude the unpicklable server manager
-        # _agentops_server_port_val (int) is inherently picklable and will be included.
-        logger.debug(f"Getting state for pickling Trainer (PID {os.getpid()}). _agentops_server_manager excluded.")
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        # In child process, self._agentops_server_manager will be None.
-        logger.debug(f"Setting state for unpickled Trainer (PID {os.getpid()}). _agentops_server_manager is None.")
-
-    def init(self, *args, **kwargs):
-        if self.agentops_managed and self._agentops_server_manager:
-            self._agentops_server_manager.start()
-            self._agentops_server_port_val = self._agentops_server_manager.get_port()
-            if self._agentops_server_port_val is None:
-                if (
-                    self._agentops_server_manager.server_process is not None
-                    and self._agentops_server_manager.server_process.is_alive()
-                ):
-                    raise RuntimeError("AgentOps server started but port is None. Check server manager logic.")
-                elif (
-                    self._agentops_server_port_val is None and self._agentops_server_manager.server_process is None
-                ):  # Server failed to start
-                    raise RuntimeError("AgentOps server manager indicates server is not running and port is None.")
-
-    def teardown(self):
-        if self.agentops_managed:
-            self._agentops_server_manager.stop()
-            logger.info("AgentOps server stopped.")
-
    def instrument(self, worker_id: int):
        instrument_all()

@@ -105,24 +76,9 @@ class AgentOpsTracer(BaseTracer):
            logger.info(f"[Worker {worker_id}] Instrumentation applied.")

        if self.agentops_managed:
-            if self._agentops_server_port_val:  # Use the stored, picklable port value
-                base_url = f"http://localhost:{self._agentops_server_port_val}"
-                env_vars_to_set = {
-                    "AGENTOPS_API_KEY": "dummy",
-                    "AGENTOPS_API_ENDPOINT": base_url,
-                    "AGENTOPS_APP_URL": f"{base_url}/notavailable",
-                    "AGENTOPS_EXPORTER_ENDPOINT": f"{base_url}/traces",
-                }
-                for key, value in env_vars_to_set.items():
-                    os.environ[key] = value
-                    logger.info(f"[Worker {worker_id}] Env var set: {key}={value}")
-            else:
-                logger.warning(
-                    f"[Worker {worker_id}] AgentOps managed, but local server port is not available. Client may not connect as expected."
-                )
-
+            os.environ.setdefault("AGENTOPS_API_KEY", "dummy")
            if not agentops.get_client().initialized:
-                agentops.init()
+                agentops.init(auto_start_session=False)  # type: ignore
                logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
            else:
                logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized.")
@@ -132,11 +88,13 @@ class AgentOpsTracer(BaseTracer):
        try:
            # new versions
            instance = agentops.sdk.core.tracer
-            instance.provider.add_span_processor(self._lightning_span_processor)
+            # 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()
-            instance._provider.add_span_processor(self._lightning_span_processor)
+            instance = TracingCore.get_instance()  # type: ignore
+            instance._provider.add_span_processor(self._lightning_span_processor)  # type: ignore

    def teardown_worker(self, worker_id: int) -> None:
        super().teardown_worker(worker_id)
@@ -145,22 +103,69 @@ class AgentOpsTracer(BaseTracer):
            self.uninstrument(worker_id)
            logger.info(f"[Worker {worker_id}] Instrumentation removed.")

-    @contextmanager
-    def trace_context(self, name: Optional[str] = None):
+    @asynccontextmanager
+    async def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> AsyncGenerator[LightningSpanProcessor, None]:
        """
        Starts a new tracing context. This should be used as a context manager.

        Args:
            name: Optional name for the tracing context.
+            store: Optional store to add the spans to.
+            rollout_id: Optional rollout ID to add the spans to.
+            attempt_id: Optional attempt ID to add the spans to.

        Yields:
-            The LightningSpanProcessor instance to collect spans.
+            The [`LightningSpanProcessor`][agentlightning.tracer.agentops.LightningSpanProcessor] 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
+
+    @contextmanager
+    def _trace_context_sync(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[LightningSpanProcessor]:
+        """Implementation of `trace_context` for synchronous execution."""
        if not self._lightning_span_processor:
            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")

-        with self._lightning_span_processor:
-            yield self._lightning_span_processor
+        kwargs: dict[str, Any] = {}
+        if name is not None:
+            kwargs["trace_name"] = name
+        elif rollout_id is not None:
+            kwargs["trace_name"] = rollout_id
+        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")
+        except Exception as e:
+            status = StatusCode.ERROR  # type: ignore
+            logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
+        finally:
+            agentops.end_trace(trace, end_state=status)  # type: ignore

    def get_last_trace(self) -> List[ReadableSpan]:
        """
@@ -173,7 +178,7 @@ class AgentOpsTracer(BaseTracer):
            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
        return self._lightning_span_processor.spans()

-    def get_langchain_callback_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
+    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
        """
        Get the Langchain callback handler for integrating with Langchain.

@@ -197,18 +202,89 @@ class AgentOpsTracer(BaseTracer):
            )
        return LangchainCallbackHandler(api_key=api_key, tags=tags)

+    get_langchain_callback_handler = get_langchain_handler  # alias
+

 class LightningSpanProcessor(SpanProcessor):
+    """Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
+    to a [`LightningStore`][agentlightning.LightningStore].
+    """

-    _spans: List[ReadableSpan] = []
+    def __init__(self):
+        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 = 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, exc_val, exc_tb):
-        pass
+    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]:
        """
@@ -220,6 +296,22 @@ class LightningSpanProcessor(SpanProcessor):
        """
        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.
@@ -231,10 +323,16 @@ class LightningSpanProcessor(SpanProcessor):
        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)
-
-    def shutdown(self) -> None:
-        pass
-
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
-        return True
@@ -1,11 +1,22 @@
-from contextlib import contextmanager
-from typing import Iterator, List, Optional, Callable, Any, Awaitable
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional

 from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.store.base import LightningStore
 from agentlightning.types import ParallelWorkerBase

+if TYPE_CHECKING:
+    from langchain_core.callbacks.base import BaseCallbackHandler  # type: ignore

-class BaseTracer(ParallelWorkerBase):
+logger = logging.getLogger(__name__)
+
+
+class Tracer(ParallelWorkerBase):
    """
    An abstract base class for tracers.

@@ -14,7 +25,7 @@ class BaseTracer(ParallelWorkerBase):
    designed to be backend-agnostic, allowing for different implementations
    (e.g., for AgentOps, OpenTelemetry, Docker, etc.).

-    The primary interaction pattern is through the `trace_context`
+    The primary interaction pattern is through the [`trace_context`][agentlightning.Tracer.trace_context]
    context manager, which ensures that traces are properly started and captured,
    even in the case of exceptions.

@@ -24,9 +35,9 @@ class BaseTracer(ParallelWorkerBase):
    tracer = YourTracerImplementation()

    try:
-        with tracer.trace_context(name="my_traced_task"):
+        async with tracer.trace_context(name="my_traced_task"):
            # ... code to be traced ...
-            run_my_agent_logic()
+            await run_my_agent_logic()
    except Exception as e:
        print(f"An error occurred: {e}")

@@ -35,26 +46,48 @@ class BaseTracer(ParallelWorkerBase):

    # Process the trace data
    if trace_tree:
-        rl_triplets = TripletExporter().export(spans)
+        rl_triplets = TracerTraceToTriplet().adapt(spans)
        # ... do something with the triplets
    ```
    """

-    @contextmanager
-    def trace_context(self, name: Optional[str] = None) -> Iterator[Any]:
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> AsyncContextManager[Any]:
        """
        Starts a new tracing context. This should be used as a context manager.

        The implementation should handle the setup and teardown of the tracing
        for the enclosed code block. It must ensure that any spans generated
        within the `with` block are collected and made available via
-        `get_last_trace`.
+        [`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.
+            rollout_id: The rollout ID to add the spans to.
+            attempt_id: The attempt ID to add the spans to.
        """
        raise NotImplementedError()

+    def _trace_context_sync(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> ContextManager[Any]:
+        """Internal API for CI backward compatibility."""
+        raise NotImplementedError()
+
    def get_last_trace(self) -> List[ReadableSpan]:
        """
        Retrieves the raw list of captured spans from the most recent trace.
@@ -64,10 +97,12 @@ class BaseTracer(ParallelWorkerBase):
        """
        raise NotImplementedError()

-    def trace_run(self, func: Callable, *args, **kwargs) -> Any:
+    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
        """
        A convenience wrapper to trace the execution of a single synchronous function.

+        Deprecated in favor of customizing Runners.
+
        Args:
            func: The synchronous function to execute and trace.
            *args: Positional arguments to pass to the function.
@@ -76,13 +111,15 @@ class BaseTracer(ParallelWorkerBase):
        Returns:
            The return value of the function.
        """
-        with self.trace_context(name=func.__name__):
+        with self._trace_context_sync(name=func.__name__):
            return func(*args, **kwargs)

-    async def trace_run_async(self, func: Callable[..., Awaitable], *args, **kwargs) -> Any:
+    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
        """
        A convenience wrapper to trace the execution of a single asynchronous function.

+        Deprecated in favor of customizing Runners.
+
        Args:
            func: The asynchronous function to execute and trace.
            *args: Positional arguments to pass to the function.
@@ -91,5 +128,13 @@ class BaseTracer(ParallelWorkerBase):
        Returns:
            The return value of the function.
        """
-        with self.trace_context(name=func.__name__):
+        async with self.trace_context(name=func.__name__):
            return await func(*args, **kwargs)
+
+    def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:  # type: ignore
+        """Get a handler to install in langchain agent callback.
+
+        Agents are expected to use this handler in their agents to enable tracing.
+        """
+        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
+        return None
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Yuqing	dc2c5e5d8b	Merge branch 'main' into dev/database-store	2025-11-12 12:02:35 +08:00
Yuge Zhang	5ae7933d41	Use unified server launcher for LiteLLM Proxy (#292 )	2025-11-12 02:45:29 +08:00
yuqing	2eb207ec79	fix lint warning and update uv.lock	2025-11-11 16:53:01 +08:00
Yuge Zhang	2ab977ed18	Dashboard - build into Python package (#291 )	2025-11-11 16:29:02 +08:00
Yuge Zhang	1eae9a34f0	Fix dashboard pipeline (#289 )	2025-11-11 00:27:02 +08:00
Yuge Zhang	4e7748b059	Dashboard - tests and infrastructure (#288 )	2025-11-10 22:30:00 +08:00
Yuge Zhang	582f67cade	Python Server Launcher (#286 )	2025-11-10 16:08:05 +08:00
부창규	9e23ba6b50	Rename the function properly in Spider (#285 )	2025-11-10 14:57:06 +08:00
yuqing	de8df23805	Merge branch 'dev/database-store' of https://github.com/microsoft/agent-lightning into dev/database-store	2025-11-07 12:05:59 +08:00
yuqing	5fc8ee0bf0	fix lint errors	2025-11-07 11:58:35 +08:00
yuqing	fe2d05218f	To break the big transaction inside timeout healthy checking	2025-11-07 11:05:12 +08:00
yuqing	c1b7827e5f	fix lint error	2025-11-06 15:54:19 +08:00
yuqing	3162ed6fb6	fix lint issue	2025-11-06 15:39:23 +08:00
Yuqing	5983eec570	Update agentlightning/store/database/retry_helper.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>	2025-11-06 15:23:59 +08:00
Yuqing	c21e065d24	Update agentlightning/store/database/orm/rollout.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>	2025-11-06 15:23:29 +08:00
yuqing	b6312db2a5	fix pre-commit warnings	2025-11-06 15:22:54 +08:00
yuqing	6131c1f7b5	rename to SqlLightningStore operations with SQLAlchemy. - Refactored existing tests to utilize the new SqlLightningStore, ensuring compatibility with the new implementation. - Adjusted timeout behavior in tests to align with the new polling intervals.	2025-11-06 15:07:30 +08:00
yuqing	294a1cc1f6	Merge remote-tracking branch 'origin/main' into dev/database-store	2025-11-06 14:41:14 +08:00
yuqing	93814a9a9c	Enhance rollout and attempt handling with type unification and new resource management features - Updated type hints to include `Union` for rollouts and attempted rollouts across multiple files. - Improved `query_rollouts` and `get_rollout_by_id` methods to return `AttemptedRollout` where applicable. - Added `query_resources` method for resource management in `DatabaseLightningStore`. - Refactored `as_attempt` and `as_rollout` methods to utilize `model_dump` for better serialization. (Serialization cost is reduced) - Updated tests to validate status transitions and ensure correct behavior with new rollout logic.	2025-11-06 11:09:24 +08:00
Yuge Zhang	3f8a3ac0f1	Preserve interface for SQL store testing (#279 )	2025-11-06 00:05:40 +08:00
yuqing	70838ee86c	Merge remote-tracking branch 'origin/main' into dev/database-store	2025-11-05 21:51:28 +08:00
Yuge Zhang	e0b55ab057	Fix preparing status transition on rollout when creating attempts (#278 )	2025-11-05 17:45:57 +08:00
yuqing	5950276103	fix typo	2025-11-05 15:51:24 +08:00
Yuge Zhang	421f2773c7	RESTful API improvements (#275 )	2025-11-05 14:12:58 +08:00
yuqing	1857e39d5e	reuse test_memory.py	2025-11-05 12:05:43 +08:00
Yuge Zhang	f717f9982f	Fix: Port conflict in tracer tests (#271 )	2025-11-05 11:10:05 +08:00
Shenghua Chen	44dbfde0b4	fix room_selector example which always run the first task (#270 )	2025-11-05 10:47:07 +08:00
yuqing	f7fe24a4b9	all tests passed with some FIXME	2025-11-04 22:49:31 +08:00
yuqing	a477bec3e1	only corner cases for status propagation	2025-11-04 20:33:29 +08:00
Yuge Zhang	713511902d	Add Tinker × Agent-lightning tuning articles to docs (#269 )	2025-11-04 15:43:24 +08:00
Ni Hao	80531c9c28	fix openai_agent version for compatibility issue. (#265 ) * fix openai_agent version for compatibility issue. * gen uv.lock --------- Co-authored-by: Hao Ni (CSI Interfusion Co Ltd) <v-nhao@microsoft.com>	2025-11-04 14:54:44 +08:00
Yuge Zhang	37daf2104f	Adding VERL replacement for Tinker (#264 ) * Adding VERL replacement for Tinker * Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>	2025-11-04 12:09:36 +08:00
yuqing	cbd6498e73	update error messages	2025-11-04 10:50:10 +08:00
yuqing	a0b68333b8	support periodic background tasks for attempt timeout checking	2025-11-04 10:14:33 +08:00
Yuge Zhang	9afdd4570c	docs: add deepwiki badge to readme (#263 )	2025-11-03 08:09:06 +00:00
yuqing	96a0d58fce	configurable retry added	2025-11-03 15:55:40 +08:00
yuqing	b8940feffb	interface implemented, but still have some issues for the periodical operations	2025-11-03 11:35:49 +08:00
Vishal V	55fbe66fe7	docs: fix typos in train-first-agent.md (#260 ) Co-authored-by: Vishal <VishalV@ibm.com>	2025-11-02 16:01:46 +08:00
Yuge Zhang	3794c97c1e	Store RESTful API updates (#259 )	2025-11-02 16:01:32 +08:00
Yuge Zhang	848623766d	Azure OpenAI Finetuning example (#256 )	2025-11-01 23:44:30 +08:00
Yuge Zhang	4cd09ec900	Add contributor and maintainer guides (#239 )	2025-11-01 13:08:59 +08:00
yuqing	47b0a4e493	initial version with all dbstore interface implemented, but not all tests passed	2025-10-31 23:35:30 +08:00
Zhiyuan He	3ed5e1e5b5	Fix training metrics before and after processing (#145 )	2025-10-31 23:09:10 +08:00
Ni Hao	3f372ff7b3	Replace AgentOps mock server with bypassable client (#202 ) --------- Co-authored-by: Hao Ni (CSI Interfusion Co Ltd) <v-nhao@microsoft.com>	2025-10-31 23:05:39 +08:00
Yuge Zhang	c453c41fd2	Fix tests-full failure to checkout PR branch (#253 )	2025-10-31 11:29:17 +00:00
Yuge Zhang	5c8ac61af6	Bump version to 0.2.2 (#248 )	2025-10-31 18:14:43 +08:00
Zhiyuan He	a02e1b91d9	Add support for verl 0.6.0 (#246 )	2025-10-31 15:30:38 +08:00
Yuge Zhang	496e793f0b	Tinker Integration (#245 )	2025-10-30 12:44:17 +08:00
Yuge Zhang	80d306ff54	[Release] v0.2.1 (#243 )	2025-10-30 08:31:45 +08:00
Yuge Zhang	5f67bfe137	Normalize Store FastAPI (#241 )	2025-10-29 21:43:37 +08:00
John Eismeier	f8c45b6ca8	propose fix a couple of typos and avoid emacs backup files (#237 ) Signed-off-by: John E <jeis4wpi@outlook.com>	2025-10-29 09:42:58 +08:00
ddsfda99	01955aead7	Fix store port conflict handling (issue #221 ) (#227 )	2025-10-28 14:58:41 +08:00
Yuge Zhang	0a9e3d75f2	Fix CI GPU trigger (#234 )	2025-10-28 14:55:18 +08:00
Ni Hao	a3b2db18fa	Make the number of tasks on the server and client consistent. (#187 )	2025-10-28 14:27:46 +08:00
Yuge Zhang	268bd77ce6	Refine notes and triggering conditions (#233 )	2025-10-28 12:30:37 +08:00
Yuge Zhang	ab6ea3c131	Augment backport implementation (#230 )	2025-10-28 03:57:16 +00:00
Yuge Zhang	d16538da96	Internal API update in preparation for Tinker integration (#226 )	2025-10-28 11:26:49 +08:00
Yuge Zhang	8ce40a0410	Add backport support (#229 )	2025-10-28 11:21:36 +08:00
Yuge Zhang	a9c7dbef22	Track Actions Status in Issue Comment Responder (#228 )	2025-10-28 00:01:31 +08:00
Yuge Zhang	e69d24f4a8	Fix response ID security (#217 )	2025-10-27 21:07:10 +08:00
Yuge Zhang	955a0cc9a3	Revert #224 (#225 )	2025-10-27 20:38:12 +08:00
Yuge Zhang	3966db6d2a	Post comment on workflow run ready (#224 )	2025-10-27 19:46:41 +08:00
Yuge Zhang	584600d72e	Fix property list too long in issue responder (#223 )	2025-10-27 11:20:34 +00:00
Yuge Zhang	955524658d	Update CI triggering mechanism (#222 )	2025-10-27 19:10:16 +08:00
Yuge Zhang	4d5e133a06	Add vLLM blog link to resources (#215 )	2025-10-26 09:57:07 +08:00
Yuge Zhang	df2a159b00	Add tutorial for launching workers on separate machines (#213 )	2025-10-25 14:16:27 +08:00
Yuge Zhang	675fc86727	Issue comment responder (#214 )	2025-10-25 13:00:02 +08:00
Yuge Zhang	c16b3a21b6	Add dependency groups from Tinker and CrewAI (#212 )	2025-10-25 12:27:37 +08:00
Yuge Zhang	aab976558b	Add Trainer port option for client-server strategies (#198 )	2025-10-25 01:41:41 +08:00
Yuge Zhang	91c85aef7e	Serialize docs deployment workflow (#205 )	2025-10-25 01:29:12 +08:00
scott-vsi	a1c36b55a0	Update verl.md (#210 ) included a link to the VERL Framework	2025-10-25 01:07:50 +08:00
Yuge Zhang	6700878f64	Switch tracer models to ConfigDict for Pydantic v2 compliance (#211 )	2025-10-25 01:07:33 +08:00
Yuge Zhang	56fa8d6881	Fix LiteLLM dual init issue (#206 )	2025-10-25 00:58:37 +08:00
Yuge Zhang	b0f28423b2	Fix model name selection in LLMProxy (#197 )	2025-10-24 01:38:49 +08:00
Yuge Zhang	e28fb8cb6b	Fix LiteLLM logging worker reset on proxy restart (#174 )	2025-10-24 00:56:16 +08:00
Yuge Zhang	fae0fba3d7	Fix trigger on label (#204 )	2025-10-23 23:05:43 +08:00
Yuge Zhang	0decbabfbe	Remove extra blank lines after Examples headers (#201 )	2025-10-23 15:06:49 +08:00
Yuge Zhang	af7a6aa2cc	Bump version to 0.2.1 (#199 )	2025-10-23 15:06:06 +08:00
Yuge Zhang	ae4e992771	Use pull_request_target trigger to allow fork-origin PRs to check on privileged workflows (#200 )	2025-10-23 13:46:44 +08:00
Yuge Zhang	8abe85ad91	Using Group Subscription for CI (#195 )	2025-10-22 21:39:58 +08:00
Yuge Zhang	22454adedb	Fix release pipeline (#194 ) Deploy Documentation / deploy (push) Has been cancelled Details PyPI Release / check-version (push) Has been cancelled Details PyPI Release / publish-pypi (push) Has been cancelled Details	2025-10-22 14:02:36 +08:00
Yuge Zhang	483c518d74	Adjust CI status placement and example catalog details (#193 )	2025-10-22 13:11:40 +08:00
Yuge Zhang	948506f3b6	Badge aggregation on workflow dispatch (#192 )	2025-10-22 11:29:50 +08:00
Yuge Zhang	34437dd6f5	Trigger the workflows on certain labels and reopen event (#190 )	2025-10-22 11:29:29 +08:00
Yuge Zhang	951fa685b5	Aggregate badge statuses (#191 )	2025-10-22 10:55:34 +08:00
Yuge Zhang	2b12e29f32	Clarify platform and runtime requirements (#188 )	2025-10-22 10:03:43 +08:00
Ni Hao	0e04363f4c	force to output logs on windows. (#176 )	2025-10-21 17:37:38 -07:00
Yuge Zhang	e91187b491	Update installation instructions (#179 )	2025-10-20 19:44:13 +08:00
Yuge Zhang	c4b829dbe7	Update README and documentation README (#183 )	2025-10-20 19:36:00 +08:00
Yuge Zhang	5c274703fe	Split examples.yml into multiple workflow definitions (#181 )	2025-10-20 14:59:20 +08:00
Yuge Zhang	55284f8394	Ensure store server serializes access per thread (#175 )	2025-10-20 14:56:19 +08:00
Yuge Zhang	895bffc5b6	Split Examples test into mulitple jobs (#180 )	2025-10-20 14:36:19 +08:00
Yuge Zhang	8e06fe6902	Migrate to use uv as dependency manager (#170 )	2025-10-20 01:42:48 +08:00
Yuge Zhang	7d8dccd2b0	Refresh Python Package Docstrings and Minor Documentation Refinement (#173 )	2025-10-19 23:08:10 +08:00
Yuge Zhang	89a887d835	Documentation update: Serving LLM, Unsloth SFT, Parallelize (#169 )	2025-10-18 20:48:04 +08:00
Yuge Zhang	d31090e9ee	Pin LangChain version to less than 1.0 (#168 )	2025-10-18 12:48:06 +08:00
Yuge Zhang	c6298a96fd	Documentation update: Train SQL Agent, Traces, Debugging (#167 )	2025-10-17 18:24:27 +00:00
Yuge Zhang	fcb2a0811e	Upgrade Calc-X Agent Example and Misc Bug Fixes (#166 )	2025-10-17 18:58:52 +08:00
Yuge Zhang	bdf6a8f223	Make `tracer.trace_context` async (#165 )	2025-10-17 02:03:27 +08:00
Yuge Zhang	8c673c241e	Streamline in-memory span eviction thresholds (#161 )	2025-10-17 00:51:22 +08:00
Yuge Zhang	b7d2d6d6cb	Upgrade SQL Agent Example to v0.2 (#164 )	2025-10-16 17:49:45 +08:00
Yuge Zhang	46a08d7272	Documentation update: Write agents and Understanding Store (#163 )	2025-10-16 15:01:39 +08:00
Yuge Zhang	cca9e9d62f	fix: include optional fields in rollout requests (#162 )	2025-10-16 14:53:40 +08:00
Yuge Zhang	8aeb0ec1ba	Preserve timeout status when spans arrive (#160 )	2025-10-16 14:12:27 +08:00
Yuge Zhang	65ba916743	Support unmanaged execution store configuration (#159 ) * Support unmanaged execution store configuration * Refine managed store cleanup and extend strategy tests	2025-10-16 12:04:46 +08:00
Yuge Zhang	d35a33dc14	[BREAKING] Strip "Base" from base classes (#158 )	2025-10-16 09:52:12 +08:00
Nanako	418691e5a2	Add Search-R1 Example and Per-Source Test Statistics (#147 ) * update Search_R1 Example * update per-source statistics comments	2025-10-16 02:46:59 +08:00
Yuge Zhang	8b33ddc028	Quickstart tutorials update and documentation structure update (#157 )	2025-10-15 18:01:09 +08:00
Yuge Zhang	bdc0b7e2a8	[BREAKING] Update API names and imports (#155 )	2025-10-15 13:37:44 +08:00
Yuge Zhang	2adaddbf7c	Pin unsloth to 2025.10.1 (#154 )	2025-10-15 01:13:45 +08:00
Yuge Zhang	86becfdbff	Add Built-in APO algorithm and Associated Examples (#153 )	2025-10-14 16:17:01 +00:00
Yuge Zhang	994384cb9b	Fix MessagesAdapter (cont.) (#152 )	2025-10-14 15:40:42 +08:00
Yuge Zhang	7f8395b941	Fix Messages Adapter (#150 )	2025-10-14 09:01:56 +08:00
Yuge Zhang	495d4eba38	Refine algo decorator fallbacks (#146 )	2025-10-13 23:46:47 +08:00
Yuge Zhang	8dd4c5a1a3	Local SFT example and bug fixes (#149 )	2025-10-13 23:10:00 +08:00
Yuge Zhang	b1ae0b75c4	Ensure Trainer.dev validates FastAlgorithm usage (#148 )	2025-10-13 22:02:52 +08:00
Yuge Zhang	9f8ec4950f	Refactor emitter module and extract common utilities (#143 )	2025-10-12 00:37:44 +08:00
Yuge Zhang	c08da2ae37	Use Azure OpenAI and MSR W&B (#142 )	2025-10-12 00:01:49 +08:00
Yuge Zhang	f032ffa319	Refactor LitAgent decorators and add prompt_rollout support (#141 )	2025-10-11 13:56:08 +08:00
Yuge Zhang	a4cf2fd5fd	Update docs palette and add version warning (#139 )	2025-10-10 17:59:01 +00:00
Yuge Zhang	8a4ecbacf6	Improve debugging experience and add new APO example (#138 )	2025-10-11 01:42:54 +08:00
Yuge Zhang	dd337d456e	Add AgentFlow community project links (#137 )	2025-10-11 01:28:19 +08:00
Yuge Zhang	a0626bdea9	Bird's eye view (v0.2) (#136 )	2025-10-10 17:09:02 +08:00
Yuge Zhang	2489d068ba	Add env overrides for client/server execution strategy (#135 )	2025-10-10 00:43:46 +08:00
Yuge Zhang	f4814949cb	Supports collecting trace data from LLMProxy (#133 )	2025-10-10 00:16:03 +08:00
Yuge Zhang	a0791e8b13	Fix flaky pipeline (#131 )	2025-10-09 14:54:33 +08:00
Yuge Zhang	26d1df698d	Show pytest durations in workflows (#130 )	2025-10-09 01:54:06 +08:00
Yuge Zhang	7bf418ea67	Use v0.2 interfaces in Trainer (#129 )	2025-10-09 01:53:12 +08:00
Yuge Zhang	d735fb27c4	Handle recovering attempts in memory store (#124 )	2025-10-09 01:45:50 +08:00
lhx	347638f218	fix: add BaseMessage import to sql_agent.py (#126 ) * fix: add BaseMessage import to sql_agent.py * trigger --------- Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>	2025-10-09 01:39:41 +08:00
Yuge Zhang	a42839b7fb	Move test files (#128 )	2025-10-08 14:36:54 +08:00
Yuge Zhang	e11036cf7b	Fix HTTP tracer tests (#127 )	2025-10-08 13:30:36 +08:00
Yuge Zhang	685eea70a6	Implement AgentRunnerV2 (#125 )	2025-10-03 01:10:09 +08:00
Yuge Zhang	a1a4fe39c6	Add missing future import annotations (#123 )	2025-10-02 01:48:49 +08:00
Yuge Zhang	a9d0c9237d	Add LLM proxy (#122 )	2025-10-02 01:26:04 +08:00
Yuge Zhang	1513b52a05	Add full pytest workflow on GPU (#121 )	2025-10-01 08:06:40 +00:00
xJkie	4ec1029577	Rollout with rollout_manager with REMAX advantage estimator (#115 )	2025-10-01 13:10:45 +08:00
Yuge Zhang	63c133051d	Add Execution Strategies: Client-Server and Shared Memory implementations (#120 )	2025-10-01 13:09:45 +08:00
Yuge Zhang	2316a8451e	Add LightningStore interface and implementation (#118 )	2025-09-30 09:29:14 +00:00
Yuge Zhang	138ad0e487	Add DeepWerewolf to community projects (#112 )	2025-09-27 13:59:56 +08:00
Yuge Zhang	504ef2c627	Add visualization to adapter (#113 ) * Add visualization to adapter * Update agentlightning/adapter/triplet.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>	2025-09-27 12:21:40 +08:00
Yuge Zhang	a63197355c	Initialize trace adapter (#103 )	2025-09-20 12:57:12 +00:00
Yuge Zhang	3eb725fade	Update pre-commit configurations and bump to 0.2.0 (#102 )	2025-09-20 07:20:14 +00:00
Yuge Zhang	66bcfeba11	Fix pyright issues (#101 )	2025-09-20 14:51:56 +08:00
Yuge Zhang	a9208ab700	Sort imports (#100 )	2025-09-20 11:30:34 +08:00
Yuge Zhang	ddc8997b8c	Embed algorithm into trainer (#99 )	2025-09-19 18:42:58 +08:00
Yuge Zhang	0a92600a4c	Ensure blank line after copyright header (#97 ) * Ensure copyright header has separating blank line * fix check script * trigger pre-commits	2025-09-19 05:32:24 +00:00
Yuge Zhang	ba10c845e1	Bug fixes and RL improvements for Agent-framework support (#96 )	2025-09-18 21:22:09 +08:00
Yuge Zhang	7ad967daf7	Skip unexpected rollout ids (#85 )	2025-09-18 20:43:44 +08:00
Yuge Zhang	f6db2dc8ab	Migrate to 1ES pipeline (#95 )	2025-09-18 20:26:16 +08:00
Yuge Zhang	5724f63cfc	Bump agentops, pin verl (#90 )	2025-09-13 10:14:53 +08:00
Luna Qiu	bd6c62dd7c	Update RAI transparency documentation (#84 ) * Update RAI transparency documentation * Update RAI_README.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>	2025-09-06 13:52:00 +08:00