.

2025-11-05 09:27:59 +08:00 · 2025-11-04 10:45:14 +08:00 · 2025-11-04 01:27:48 +08:00 · 2025-11-04 01:23:21 +08:00 · 2025-11-04 00:47:24 +08:00 · 2025-11-03 16:42:51 +00:00
238 changed files with 36327 additions and 3845 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 });
@@ -8,6 +8,10 @@ on:
      - 'v*'
  workflow_dispatch:

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

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

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

      - name: Get current version
        id: get_version
@@ -44,16 +43,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,26 @@ 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'
-
-      - name: Install build dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e .[dev]
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: uv sync --frozen --no-default-groups --group dev

      - 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')"
@@ -8,61 +8,81 @@ on:

  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:
-        setup: [stable, latest]
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
      fail-fast: false
    steps:
      - name: Check GPU status
        run: nvidia-smi
      - uses: actions/checkout@v4
-      - name: Create a virtual environment
-        run: python3 -m venv .venv
-      - name: Install dependencies (${{ matrix.setup }})
-        run: |
-          . .venv/bin/activate
-          ./scripts/setup_${{ matrix.setup }}_gpu.sh
+        with:
+          ref: ${{ github.event.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: |
-          . .venv/bin/activate
-          which python
-          which pip
-          which uvx
-          pip list | tee requirements-freeze.txt
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
      - name: Upload dependencies artifact
        uses: actions/upload-artifact@v4
        with:
-          name: dependencies-${{ matrix.setup }}
+          name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
          path: requirements-freeze.txt
          compression-level: 0

      - name: Launch LiteLLM Proxy
        run: |
-          set -ex
-          . .venv/bin/activate
-          litellm --config scripts/litellm_ci.yaml --port 12306 &
-          sleep 10  # Wait for the proxy to be up
+          ./scripts/litellm_run.sh
        env:
-          AZURE_API_BASE: ${{ secrets.AZURE_API_BASE }}
-          AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}
-      - name: Verify LiteLLM Proxy
-        run: |
-          set -ex
-          . .venv/bin/activate
-          python scripts/litellm_sanity_check.py
-        env:
-          OPENAI_BASE_URL: http://localhost:12306/
-          OPENAI_API_KEY: dummy
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}

      - name: Run tests
        run: |
-          set -ex
-          . .venv/bin/activate
-          pytest -v --durations=0 tests
+          uv run pytest -v --durations=0 tests
        env:
          PYTEST_ADDOPTS: "--color=yes"
          OPENAI_BASE_URL: http://localhost:12306/
@@ -5,9 +5,9 @@ permissions:

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

  schedule:
@@ -16,70 +16,71 @@ on:

 jobs:

-  lint-fast:
-    name: Lint - Fast
+  lint:
+    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'
      - name: Run pre-commit
        uses: pre-commit/action@v3.0.1
      - name: Check Python headers
-        run: |
-          python scripts/check_python_headers.py
+        run: uv run --locked --no-sync scripts/check_python_headers.py
      - name: Run Black
-        run: black --check .
+        run: uv run --locked --no-sync black --check .
      - name: Run isort
-        run: isort --check-only .
-      - name: Run pyright
-        run: pyright -p pyrightconfig.fast.json
-
-  lint-slow:
-    name: Lint - Slow
-    runs-on: ubuntu-latest
-    timeout-minutes: 10
-    steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.12'
-      - name: Install dependencies
-        run: |
-          ./scripts/setup_type_checking.sh
-      - name: Run Black
-        run: black --check .
-      - name: Run isort
-        run: isort --check-only .
-      - name: Run pyright
-        run: pyright -p pyrightconfig.json
+        run: 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'

  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:
@@ -92,35 +93,47 @@ 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
      - name: Run tests
        run: |
-          pytest -v --durations=0 tests
+          uv run pytest -v --durations=0 tests
        env:
          PYTEST_ADDOPTS: "--color=yes"
@@ -189,6 +189,9 @@ cython_debug/
 #  you could uncomment the following to ignore the enitre vscode folder
 .vscode/

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

@@ -8,6 +8,8 @@ repos:
        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
@@ -0,0 +1 @@
+3.12
@@ -1,11 +1,11 @@
-<div style="text-align:center; margin-bottom:20px;">
-  <img src="docs/assets/readme-banner.png" alt="Agent-lightning-banner" style="max-width:600px"/>
-</div>
+<p align="center">
+  <img src="docs/assets/readme-banner.svg" alt="Agent-lightning-banner" style="width:600px"/>
+</p>

 # Agent Lightning⚡

-[![CPU Test](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml)
-[![GPU Test](https://github.com/microsoft/agent-lightning/actions/workflows/examples.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples.yml)
+[![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)
 [![Discord](https://img.shields.io/badge/Discord-Join-5865F2?logo=discord&logoColor=white)](https://discord.gg/RYk7CdvDR7)
@@ -17,14 +17,29 @@ Join our [Discord community](https://discord.gg/RYk7CdvDR7) to connect with othe
 ## ⚡ Core Features

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

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

-## ⚡ Resources
+<p align="center">
+  <img src="docs/assets/readme-diff.svg" alt="Agent-Lightning Core Quickstart" style="width:100%"/>
+</p>

+## ⚡ Installation
+
+```bash
+pip install agentlightning
+```
+
+Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.
+
+To start using Agent-lightning, check out our [documentation](https://microsoft.github.io/agent-lightning/) and [examples](./examples).
+
+## ⚡ Articles
+
+- 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.
@@ -35,114 +50,27 @@ Join our [Discord community](https://discord.gg/RYk7CdvDR7) to connect with othe
 - [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
 - [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.

-## ⚡ Installation
-
-First, let's get your environment set up. We'll be using `/path/to/agentlightning` to refer to the directory containing this README file.
-
-### 1. Set Up Your Environment
-
-We strongly recommend creating a new virtual environment to avoid conflicts with other packages. You can use either `conda` or `venv`. **Python 3.10 or later** is recommended.
-
-### 2. Install Core Training Dependencies (Optional)
-
-If you are running RL with Agent-Lightning, the next step is to install the essential packages: `PyTorch`, `FlashAttention`, `vLLM` and `VERL`. The following versions and installation order have been tested and are confirmed to work.
-
-```bash
-pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
-pip install flash-attn --no-build-isolation
-pip install vllm==0.9.2
-pip install verl==0.5.0
-```
-
-See `scripts/setup_stable_gpu.sh` for a full installation script.
-
-### 3. Install Agent Lightning
-
-Now, you're ready to install Agent Lightning itself.
-
-```bash
-pip install agentlightning
-```
-
-### 4. Install Agent Frameworks (Optional)
-
-If you plan to use other agent frameworks, you can install them with the following commands. If you don't need these, feel free to skip this step.
-We recommend doing this as the final step to avoid dependency versions being overwritten by mistake.
-
-```bash
-# AutoGen (Recommended to install first)
-pip install "autogen-agentchat" "autogen-ext[openai]"
-
-# LiteLLM
-pip install "litellm[proxy]"
-
-# MCP
-pip install mcp
-
-# UV
-pip install uv
-
-# OpenAI Agents
-pip install openai-agents
-
-# LangChain
-pip install langgraph "langchain[openai]" langchain-community langchain-text-splitters
-
-# SQL-related dependencies
-pip install sqlparse nltk
-```
-
-Don't worry if dependency conflicts arise during this step. Follow the installation order above and the conflicts generally do not matter.
-
-## ⚡ Examples
-
-For more detailed examples, please see the `examples` folder:
-
-1. [calc_x](examples/calc_x): An agent built with AutoGen with calculator tool use, trained on Calc-X dataset with Reinforcement Learning.
-2. [spider](examples/spider): A write-check-rewrite looped agent with LangGraph with SQL execution; selectively optimize write and rewrite on Spider dataset with Reinforcement Learning.
-3. [apo](examples/apo): An example to customize an optimization algorithm: Automatic Prompt Optimization.
-
-## ⚡ Important Caveats
-
-1.  **AgentOps Integration**: Agent Lightning uses [AgentOps](https://github.com/AgentOps-AI/agentops) for agent tracking by default. If you're already using AgentOps in your own code, you'll need to disable our managed AgentOps client by modifying the `tracer` parameter of trainer.
-2.  **Debugging Traces**: If you encounter issues with tracing, you can visualize the trace tree using `tracer.last_trace().visualize("tree_graph")`. Please note that this API is experimental and may change in future releases.
-3.  **Launching the Server and Agents**: Currently, the training server and agent clients must be launched in separate processes. You can open two terminal windows or run one of them in the background. The launching order generally doesn't matter.
-4.  **Environment Variables**: The environment variables and working directory at the time of `ray init` are important. If you run into "file not found" errors, try restarting Ray from your current working directory.
-5.  **Handling Timeouts**: The training server may hang if samples fail or time out on the agent side. To prevent this, we recommend setting limits on the prompt and response lengths, as this is the most common cause of failures.
-6.  **VERL Failures**: Save checkpoints frequently, as VERL with vLLM may sometimes experience out-of-memory issues. If you encounter a VERL failure, you can resume training from the last checkpoint.
-
 ## ⚡ 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

@@ -1,22 +0,0 @@
-import asyncio
-
-
-async def a():
-    print("a")
-    b()
-    print("finish")
-
-
-def b():
-    print("b")
-    loop = asyncio.get_running_loop()
-    fut = asyncio.run_coroutine_threadsafe(c(), loop)
-    fut.result(timeout=5.0)
-
-
-async def c():
-    print("c")
-    await asyncio.sleep(0.1)
-
-
-asyncio.run(a())
@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.

-__version__ = "0.2.0"
+__version__ = "0.2.1"

 from .adapter import *
 from .algorithm import *
@@ -1,11 +1,12 @@
 # Copyright (c) Microsoft. All rights reserved.

-from .base import Adapter, TraceAdapter
+from .base import Adapter, OtelTraceAdapter, TraceAdapter
 from .messages import TraceToMessages
 from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase

 __all__ = [
    "TraceAdapter",
+    "OtelTraceAdapter",
    "Adapter",
    "TraceToTripletBase",
    "TracerTraceToTriplet",
@@ -13,18 +13,20 @@ T_to = TypeVar("T_to")
 class Adapter(Generic[T_from, T_to]):
    """Base class for synchronous adapters that convert data from one format to another.

-    This class defines a simple protocol for transformation:
+    The class defines a minimal protocol so that adapters can be treated like callables while
+    still allowing subclasses to supply the concrete transformation logic.

-    - The `__call__` method makes adapters callable, so they can be used like functions.
-    - Subclasses must implement the `adapt` method to define the actual conversion logic.
+    !!! note
+        Subclasses must override [`adapt()`][agentlightning.Adapter.adapt] to provide
+        the actual conversion.

-    Type parameters:
+    Type Variables:

-    - T_from: The source data type (input).
-    - T_to: The target data type (output).
+        T_from: Source data type supplied to the adapter.

-    Example:
+        T_to: Target data type produced by the adapter.

+    Examples:
        >>> class IntToStrAdapter(Adapter[int, str]):
        ...     def adapt(self, source: int) -> str:
        ...         return str(source)
@@ -37,8 +39,9 @@ class Adapter(Generic[T_from, T_to]):
    def __call__(self, source: T_from, /) -> T_to:
        """Convert the data to the target format.

-        This method delegates to `adapt` and allows the adapter
-        to be invoked as a function.
+        This method delegates to [`adapt()`][agentlightning.Adapter.adapt] so that an
+        instance of [`Adapter`][agentlightning.Adapter] can be used like a standard
+        function.

        Args:
            source: Input data in the source format.
@@ -51,8 +54,8 @@ class Adapter(Generic[T_from, T_to]):
    def adapt(self, source: T_from, /) -> T_to:
        """Convert the data to the target format.

-        Subclasses should override this method with the concrete
-        transformation logic.
+        Subclasses must override this method with the concrete transformation logic. The base
+        implementation raises `NotImplementedError` to make the requirement explicit.

        Args:
            source: Input data in the source format.
@@ -66,17 +69,12 @@ class Adapter(Generic[T_from, T_to]):
 class OtelTraceAdapter(Adapter[List[ReadableSpan], T_to], Generic[T_to]):
    """Base class for adapters that convert OpenTelemetry trace spans into other formats.

-    This class specializes `Adapter` for working with OpenTelemetry `ReadableSpan`
-    objects. It expects a list of spans as input and produces a custom target format
-    (e.g., reinforcement learning training data, SFT datasets, logs, metrics).
+    This specialization of [`Adapter`][agentlightning.Adapter] expects a list of
+    `opentelemetry.sdk.trace.ReadableSpan` instances and produces any target format, such as
+    reinforcement learning trajectories, structured logs, or analytics-ready payloads.

-    Subclasses should override `adapt` to define the desired conversion.
-
-    Type parameters:
-        T_to: The target data type that spans should be converted into.
-
-    Example:
-        >>> class TraceToDictAdapter(TraceAdapter[dict]):
+    Examples:
+        >>> class TraceToDictAdapter(OtelTraceAdapter[dict]):
        ...     def adapt(self, spans: List[ReadableSpan]) -> dict:
        ...         return {"count": len(spans)}
        ...
@@ -89,7 +87,8 @@ class OtelTraceAdapter(Adapter[List[ReadableSpan], T_to], Generic[T_to]):
 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` for working with trace spans. It expects a list of
-    Agent-lightning spans as input and produces a custom target format
-    (e.g., reinforcement learning training data, SFT datasets, logs, metrics).
+    This class specializes [`Adapter`][agentlightning.Adapter] for working with
+    [`Span`][agentlightning.Span] instances emitted by Agent Lightning instrumentation.
+    Subclasses receive entire trace slices and return a format suited for the downstream consumer,
+    for example reinforcement learning training data or observability metrics.
    """
@@ -1,28 +1,48 @@
 # Copyright (c) Microsoft. All rights reserved.

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

-from openai.types.chat import (
-    ChatCompletionAssistantMessageParam,
-    ChatCompletionFunctionToolParam,
-    ChatCompletionMessageFunctionToolCallParam,
-    ChatCompletionMessageParam,
-)
 from pydantic import TypeAdapter

 from agentlightning.types import Span

 from .base import TraceAdapter

+if TYPE_CHECKING:
+    from openai.types.chat import (
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+

 class OpenAIMessages(TypedDict):
+    """OpenAI-style chat messages with optional tool definitions.
+
+    Attributes:
+        messages: Ordered chat messages that describe the conversation.
+        tools: Tool specifications available to the assistant, if any.
+    """
+
    messages: List[ChatCompletionMessageParam]
    tools: Optional[List[ChatCompletionFunctionToolParam]]


 class _RawSpanInfo(TypedDict):
+    """Intermediate representation parsed from a span.
+
+    Attributes:
+        prompt: Prompt messages reconstructed from span attributes.
+        completion: Assistant completions following tool invocations.
+        request: Request payload recorded in the trace.
+        response: Response payload recorded in the trace.
+        tools: Tool call metadata extracted from child spans.
+    """
+
    prompt: List[Dict[str, Any]]
    completion: List[Dict[str, Any]]
    request: Dict[str, Any]
@@ -31,16 +51,20 @@ class _RawSpanInfo(TypedDict):


 def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
-    """
-    Convert a flat dict with keys like 'gen_ai.prompt.0.role'
-    into structured nested dicts or lists under the given prefix.
+    """Convert flattened trace attributes into nested structures.
+
+    Attributes emitted by the tracing pipeline often arrive as dotted paths (for example
+    `gen_ai.prompt.0.role`). This helper groups those keys into nested dictionaries or lists so that
+    downstream processing can operate on structured data.

    Args:
-        data: Flat dictionary (keys are dotted paths).
-        prefix: Top-level key to extract (e.g., 'gen_ai.prompt').
+        data: Flat dictionary whose keys are dotted paths.
+        prefix: Top-level key (for example `gen_ai.prompt`) that determines which attributes are
+            grouped.

    Returns:
-        A nested dict (if no index detected) or list (if indexed).
+        A nested dictionary (no numeric index detected) or list (numeric indices detected) containing
+        the grouped values.
    """
    result: Union[Dict[str, Any], List[Any]] = {}

@@ -80,12 +104,28 @@ def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any],


 def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]:
-    """
-    Convert raw tool call traces + prompt/completion list
-    into OpenAI fine-tuning JSONL format (tool calling style).
+    """Convert raw trace payloads into OpenAI-style chat messages.

-    https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/fine-tuning-functions
+    The function consumes an iterable produced by
+    [`TraceToMessages.adapt()`][agentlightning.TraceToMessages.adapt] and yields
+    structures that match the OpenAI fine-tuning JSONL schema, including tool definitions.
+
+    Args:
+        prompt_completion_list: Raw prompt/completion/tool payloads extracted from a trace.
+
+    Returns:
+        A generator that yields [`OpenAIMessages`][agentlightning.adapter.messages.OpenAIMessages]
+        entries compatible with the OpenAI Functions fine-tuning format.
    """
+
+    # Import locally to avoid legacy OpenAI version type import errors
+    from openai.types.chat import (
+        ChatCompletionAssistantMessageParam,
+        ChatCompletionFunctionToolParam,
+        ChatCompletionMessageFunctionToolCallParam,
+        ChatCompletionMessageParam,
+    )
+
    for pc_entry in prompt_completion_list:
        messages: List[ChatCompletionMessageParam] = []

@@ -157,25 +197,29 @@ def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Ge


 class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
-    """
-    Adapter that converts OpenTelemetry trace spans into OpenAI-compatible message format.
+    """Convert trace spans into OpenAI-compatible conversation messages.

-    This adapter processes trace spans containing LLM conversation data and transforms them
-    into structured OpenAI message format suitable for fine-tuning or analysis. It extracts
-    prompts, completions, tool calls, and function definitions from trace attributes and
-    reconstructs the conversation flow.
+    The adapter reconstructs prompts, completions, tool calls, and function definitions from
+    `gen_ai.*` span attributes. The resulting objects match the JSONL structure expected by the
+    OpenAI fine-tuning pipeline.

-    The adapter handles:
-    - Converting flat trace attributes into structured message objects
-    - Extracting and matching tool calls with their corresponding requests
-    - Building proper OpenAI ChatCompletionMessage objects with roles, content, and tool calls
-    - Generating function definitions for tools used in conversations
+    !!! warning
+        The adapter assumes all spans share a common trace and that tool call spans are direct
+        children of the associated completion span.
    """

    def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable[Dict[str, Any]]:
-        """Find tool calls in the trace. Returns a dict with the tool call id, name, and arguments.
+        """Yield tool call payloads for a completion span.

-        The spans that are direct children of the completion span are the tool calls.
+        Args:
+            completion: The completion span whose descendants should be inspected.
+            all_spans: The complete span list belonging to the trace.
+
+        Yields:
+            Dictionaries describing tool calls with identifiers, names, and arguments.
+
+        Raises:
+            ValueError: If a candidate tool span cannot be converted into a dictionary.
        """
        # Get all the spans that are children of the completion span
        children = [span for span in all_spans if span.parent_id == completion.span_id]
@@ -188,6 +232,15 @@ class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
                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:
@@ -3,6 +3,7 @@
 from __future__ import annotations

 import json
+import logging
 import re
 from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple, Union, cast
@@ -10,16 +11,22 @@ from typing import Any, Dict, List, Optional, Tuple, Union, cast
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel

-from agentlightning.types import SpanNames, Triplet
-from agentlightning.types.tracer import Span
+from agentlightning.types import Span, SpanNames, Triplet

 from .base import TraceAdapter

+logger = logging.getLogger(__name__)
+

 class Transition(BaseModel):
-    """
-    Transition class representing one transition in a trajectory.
-    State and action are a list of token IDs.
+    """A single transition within a reinforcement learning trajectory.
+
+    Attributes:
+        state: Token identifiers describing the model input state.
+        action: Token identifiers representing the model output.
+        response_id: Identifier of the LLM response used to deduplicate spans.
+        agent_name: Human-readable agent name captured from the trace.
+        reward: Scalar reward associated with the transition, if available.
    """

    state: List[int]
@@ -31,22 +38,27 @@ class Transition(BaseModel):


 class RewardMatchPolicy(str, Enum):
-    """How to find the reward for each transition from the trace.
-    In all cases, the reward must have data `{"type": "reward", "value": <float>|None}`,
-    as defined in `reward.py`.
+    """Strategies for matching rewards to LLM call spans.
+
+    !!! note
+        Each reward span must expose a payload shaped like `{"type": "reward", "value": <float>|None}`
+        as described in `reward.py`.
    """

    FIRST_SIBLING = "first_sibling"
-    """Use the first sibling in the current trace subtree as the reward, except another LLM call match is found."""
+    """Use the first sibling in the current trace subtree as the reward unless another LLM call match is found."""

    FIRST_OCCURRENCE = "first_occurrence"
-    """Use the first occurrence of the reward (in start time order) that occur after the current LLM call match.
-    """
+    """Use the first reward encountered in chronological order after the current LLM call match."""


 class TraceTree:
-    """
-    A trace item, along with its span and children.
+    """Tree representation of a trace span and its descendants.
+
+    Attributes:
+        id: Unique identifier for the span node.
+        span: [`Span`][agentlightning.Span] backing this node.
+        children: Child nodes connected to the current span.
    """

    def __init__(
@@ -80,10 +92,16 @@ class TraceTree:
        self.children.append(child)

    def visualize(self, filename: str, interested_span_match: str | None = None) -> None:
-        """
-        Visualize the trace tree using graphviz.
-        For debugging purposes only.
-        Use `interested_span_match` to filter the spans (and its ancesters) to be visualized.
+        """Render the trace tree with Graphviz for debugging purposes.
+
+        Args:
+            filename: Base filename for the generated `.png` diagram.
+            interested_span_match: Optional regular expression used to keep only matching spans
+                (and their ancestors) in the output.
+
+        !!! note
+            The method requires the optional `graphviz` dependency to be available in the runtime
+            environment.
        """
        import graphviz

@@ -125,9 +143,11 @@ class TraceTree:
        dot.render(filename, format="png", cleanup=True)  # type: ignore

    def names_tuple(self) -> Tuple[str, List[Any]]:
-        """Return the span name, and a list of children.
-        Each child is also a tuple of span name and a list of children.
-        Useful for debugging and testing.
+        """Return the span name alongside nested child names.
+
+        Returns:
+            A tuple of the current span name and a list of tuples for each child containing the
+            child name and its descendants.
        """
        name = self.span.name
        agent_name = self.agent_name()
@@ -140,15 +160,14 @@ class TraceTree:
        return name, children_names

    def traverse(self) -> List["TraceTree"]:
-        """
-        Traverse the trace tree and return a list of all spans.
-        """
+        """Traverse the tree depth first and return every node."""
        spans: List["TraceTree"] = [self]
        for child in self.children:
            spans.extend(child.traverse())
        return spans

    def to_json(self) -> dict[str, Any]:
+        """Convert the tree node into a JSON-serialisable structure."""
        if isinstance(self.span, ReadableSpan):
            span_data = json.loads(self.span.to_json())
        else:
@@ -161,10 +180,17 @@ class TraceTree:

    @classmethod
    def from_spans(cls, spans: List[Span]) -> "TraceTree":
-        """
-        Create a TraceTree from a list of spans.
-        All spans without parents found will be considered as candidate root spans.
-        If multiple root spans are found, a virtual root span will be created as the parent of all root spans.
+        """Construct a tree from a flat list of spans.
+
+        Args:
+            spans: Spans that collectively form a single trace segment.
+
+        Returns:
+            A [`TraceTree`][agentlightning.adapter.triplet.TraceTree] rooted at either the
+            discovered root span or a synthetic root when multiple roots are present.
+
+        Raises:
+            ValueError: If the span list is empty or no root span can be inferred.
        """

        if not spans:
@@ -245,8 +271,11 @@ class TraceTree:
            return root_span

    def agent_name(self) -> Optional[str]:
-        """Return the name of agent span. Return the agent or None (not an agent at all).
-        Extend this function to support more agent frameworks."""
+        """Return the agent name associated with the span, if any.
+
+        Returns:
+            Agent name extracted from known attributes, otherwise `None`.
+        """
        attributes = self.span.attributes
        if attributes is None:  # type: ignore
            return None
@@ -279,6 +308,11 @@ class TraceTree:
            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",
@@ -299,6 +333,11 @@ class TraceTree:
        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

@@ -312,12 +351,19 @@ class TraceTree:
        within_llm_call: Optional[bool] = None,
        existing_llm_call_response_ids: Optional[set[str]] = None,
    ) -> List[Tuple["TraceTree", str]]:
-        """Find all LLM calls in the trace tree.
+        """Find LLM call spans matching the supplied filters.

-        The LLM call is defined as a span with type = request and name matching `llm_call_match`.
-        If `agent_match` is not None, it must also reside in an agent span (type = agent) with name matched.
+        Args:
+            llm_call_match: Regular expression used to match span names that qualify as LLM calls.
+            agent_match: Optional regular expression that must match the enclosing agent span name.
+            within_matching_subtree: Marker propagated through recursive calls to record matching agents.
+            within_reward: When `True`, suppresses LLM matches under reward spans.
+            within_llm_call: When `True`, prevents duplicate matches for nested LLM calls.
+            existing_llm_call_response_ids: Known response identifiers used to deduplicate spans.

-        Return a list of traces and the agent names (why it's selected).
+        Returns:
+            A list of tuples pairing the matching node with the agent subtree label that triggered the
+            match.
        """
        llm_calls: List[Tuple[TraceTree, str]] = []

@@ -373,17 +419,15 @@ class TraceTree:
        return llm_calls

    def repair_hierarchy(self) -> None:
-        """
-        We find that sometimes the hierarchy is not correct, due to the way the spans are created.
-        The spans within the agent frameworks (e.g., OpenAI Agent SDK) and spans within the LLM frameworks
-        (e.g., Anthropic) are created in two systems.
-        So the inner LLM completion span does not necessarily have an agent span as a parent.
-        Rather they sometimes directly become children of the root span.
-        This becomes a problem when we want to select the LLM completion span with agent as filter.
-        To repair the hierarchy, for each children of the root span, we find a span over the whole tree,
-        with duration covering the current span and being closest to the current span.
+        """Repair missing parent-child relationships introduced by mixed tracing systems.

-        This function modifies the tree in place.
+        Some agent frameworks emit spans via multiple subsystems, which can cause LLM completion
+        spans to float directly under the root span instead of being nested under the correct agent.
+        The method re-parents those spans to the closest ancestor that fully envelopes the child in
+        time.
+
+        If we don't, when we want to select the LLM completion span with agent as filter.
+        We will never get the correct span underneath.
        """
        nodes_to_repair = list(self.children)
        for repair_node in nodes_to_repair:
@@ -410,7 +454,16 @@ class TraceTree:
                closest_parent.children.append(repair_node)

    def match_rewards(self, reward_match: str, llm_calls: List["TraceTree"]) -> dict[str, Optional[float]]:
-        """Match the rewards to the LLM calls."""
+        """Assign rewards to previously matched LLM calls.
+
+        Args:
+            reward_match: Strategy identifier from
+                [`RewardMatchPolicy`][agentlightning.adapter.triplet.RewardMatchPolicy].
+            llm_calls: Trace nodes representing LLM call spans.
+
+        Returns:
+            Mapping from span identifier to reward value or `None` when no reward is available.
+        """
        llm_call_ids = set([llm_call.id for llm_call in llm_calls])
        rewards: dict[str, Optional[float]] = {}

@@ -455,6 +508,30 @@ class TraceTree:

        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",
@@ -463,20 +540,21 @@ class TraceTree:
        dedup_llm_call: bool = True,
        reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
        final_reward: Optional[float] = None,
+        _skip_empty_token_spans: bool = False,
    ) -> List[Triplet]:
-        """Convert the trace tree to a trajectory.
+        """Convert the trace tree into a trajectory of [`Triplet`][agentlightning.Triplet] items.

-        First, we find all the LLM calls (span type = request, `llm_call_match` matching the span name).
-        If the agent match is set, we check, for each LLM call,
-        if it resides in an agent (span type = agent, `agent_match` matching the span name).
-        The above sets the basis for the trajectory, as we use the prompt token IDs and response token IDs for each LLM call,
-        as the state and action of each transition.
+        Args:
+            llm_call_match: Regular expression for LLM call span names.
+            agent_match: Optional regular expression for agent span names.
+            exclude_llm_call_in_reward: When `True`, prevents searching for rewards under the LLM
+                call subtree.
+            dedup_llm_call: When `True`, deduplicates spans using the LLM response identifier.
+            reward_match: Reward matching policy used to associate reward spans with LLM calls.
+            final_reward: Optional reward appended to the final transition when provided.

-        Then, we find the reward for each transition.
-        The reward is searched on the trace tree, after the LLM call,
-        until the next LLM call or the end of the tree depending on the policy.
-        It can be enforced to a sibling or the first occurrence in the time order, depending on the policy.
-        If a reward is never found for a transition, it is set to None.
+        Returns:
+            A list of [`Triplet`][agentlightning.Triplet] objects ordered by call sequence.
        """
        # Find all LLM calls
        llm_calls = self.find_llm_calls(
@@ -487,25 +565,23 @@ class TraceTree:
            within_llm_call=False if dedup_llm_call else None,
            existing_llm_call_response_ids=set(),
        )
-        id_transitions = [
-            (
-                llm_call.id,
-                Triplet(
-                    prompt={"token_ids": llm_call.span.attributes.get("prompt_token_ids", [])},  # type: ignore
-                    response={"token_ids": llm_call.span.attributes.get("response_token_ids", [])},  # type: ignore
-                    reward=None,
-                    metadata=dict(
-                        response_id=llm_call.span.attributes.get(  # type: ignore
-                            "gen_ai.response.id", None
-                        ),  # it works at least for OpenAI
-                        agent_name=agent_name,
-                    ),
-                ),
-            )
-            for llm_call, agent_name in llm_calls
-        ]

-        rewards = self.match_rewards(reward_match, [call for call, _ in llm_calls])
+        id_transitions: List[Tuple[str, Triplet]] = []
+        # We need to filter out the LLM calls with unrecorded token IDs
+        filtered_llm_calls: List[Tuple[TraceTree, str]] = []
+        for llm_call, agent_name in llm_calls:
+            triplet = self.span_to_triplet(llm_call.span, agent_name)
+            # This is a hot-fix for Tinker+CrewAI, which has some anonymous requests outside the trained agent.
+            # TODO: We might need to reconsider this.
+            if _skip_empty_token_spans and (
+                not triplet.prompt.get("token_ids") or not triplet.response.get("token_ids")
+            ):
+                logger.warning(f"Skipping LLM call with unrecorded token IDs: {triplet}")
+                continue
+            filtered_llm_calls.append((llm_call, agent_name))
+            id_transitions.append((llm_call.id, triplet))
+
+        rewards = self.match_rewards(reward_match, [call for call, _ in filtered_llm_calls])
        transitions = [
            transition.model_copy(update={"reward": rewards.get(id, None)}) for id, transition in id_transitions
        ]
@@ -522,22 +598,22 @@ class TraceTree:


 class TraceToTripletBase(TraceAdapter[List[Triplet]]):
-    """
-    Base class for trace triplet adapters.
-    """
+    """Base class for adapters that emit [`Triplet`][agentlightning.Triplet] trajectories."""


 class TracerTraceToTriplet(TraceToTripletBase):
-    """
-    An adapter to convert OpenTelemetry spans to triplet data.
+    """Convert tracer-emitted spans into triplet trajectories.

    Attributes:
-        repair_hierarchy: When `repair_hierarchy` is set to True, the trace will be repaired with the time information.
-            See `TraceTree.repair_hierarchy` for more details.
-        llm_call_match: Regular expression pattern to match LLM call span names.
-        agent_match: Optional regular expression pattern to match agent span names. If None, all agents are matched.
-        exclude_llm_call_in_reward: Whether to exclude LLM calls that occur within reward spans.
-        reward_match: Policy for matching rewards to LLM calls.
+        repair_hierarchy: When `True`, repair the span tree using
+            [`TraceTree.repair_hierarchy()`][agentlightning.adapter.triplet.TraceTree.repair_hierarchy]
+            before matching calls and rewards.
+        llm_call_match: Regular expression pattern that selects LLM call span names.
+        agent_match: Optional regular expression pattern for agent span names. When omitted, spans
+            from any agent are considered.
+        exclude_llm_call_in_reward: When `True`, ignore matches under reward spans while searching
+            for rewards.
+        reward_match: Strategy used to associate rewards with LLM calls.
    """

    def __init__(
@@ -547,12 +623,14 @@ class TracerTraceToTriplet(TraceToTripletBase):
        agent_match: Optional[str] = None,
        exclude_llm_call_in_reward: bool = True,
        reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
+        _skip_empty_token_spans: bool = False,
    ):
        self.repair_hierarchy = repair_hierarchy
        self.llm_call_match = llm_call_match
        self.agent_match = agent_match
        self.exclude_llm_call_in_reward = exclude_llm_call_in_reward
        self.reward_match = reward_match
+        self._skip_empty_token_spans = _skip_empty_token_spans

    def visualize(
        self,
@@ -561,16 +639,17 @@ class TracerTraceToTriplet(TraceToTripletBase):
        filename: str = "trace_tree",
        interested_span_match: str | None = None,
    ) -> TraceTree:
-        """
-        Visualize the trace tree.
+        """Visualize the trace tree built from the supplied spans.

        Args:
-            source (List[Span]): The list of OpenTelemetry spans to visualize.
-            filename (str): The base filename for the output visualization (default: "trace_tree").
-            interested_span_match (str | None): Optional regular expression pattern to highlight or focus on specific spans in the visualization.
+            source: Collection of Agent Lightning [`Span`][agentlightning.Span] objects
+                or raw `opentelemetry.sdk.trace.ReadableSpan` instances.
+            filename: Base filename for the generated image; `.png` is appended automatically.
+            interested_span_match: Optional regular expression used to highlight a subset of spans.

        Returns:
-            TraceTree: The constructed trace tree object.
+            The [`TraceTree`][agentlightning.adapter.triplet.TraceTree] built from the provided
+            spans.
        """
        source_normalized = [
            Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
@@ -583,7 +662,14 @@ class TracerTraceToTriplet(TraceToTripletBase):
        return trace_tree

    def adapt(self, source: Union[List[Span], List[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
-        """Convert OpenTelemetry spans to a list of Triplet objects."""
+        """Convert tracer spans into [`Triplet`][agentlightning.Triplet] trajectories.
+
+        Args:
+            source: Agent Lightning spans or raw OpenTelemetry spans that form a trace.
+
+        Returns:
+            Ordered list of trajectory transitions with prompt, response, and reward information.
+        """
        source_normalized = [
            Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
            for span in source
@@ -596,30 +682,29 @@ class TracerTraceToTriplet(TraceToTripletBase):
            agent_match=self.agent_match,
            exclude_llm_call_in_reward=self.exclude_llm_call_in_reward,
            reward_match=self.reward_match,
+            _skip_empty_token_spans=self._skip_empty_token_spans,
        )
        return trajectory


 class LlmProxyTraceToTriplet(TraceToTripletBase):
-    """
-    Converting telemetry data emitted by the LLM Proxy to triplet data.
-    This adapter is very experimental. Should only be used when the TracerTraceToTriplet does not work at all.
+    """Convert telemetry emitted by the LLM Proxy into triplet trajectories.

-    IMPORTANT: Do NOT rely on timestamps here. Proxy spans can be emitted from different
-    machines with unsynchronized clocks. We therefore treat `sequence_id` as the only
-    reliable ordering primitive and perform "first occurrence" reward matching using
-    sequence order only.
+    !!! warning
+        This adapter is experimental and might be merged with
+        [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] in the future.
+
+    !!! danger
+        Do not rely on timestamps when using this adapter. Proxy spans can originate on different
+        machines with unsynchronised clocks, so `sequence_id` is treated as the sole source of
+        ordering.

    Strategy:

-    1) Sort spans by (sequence_id, start_time).
-    2) Extract LLM calls that expose prompt/response token IDs from either:
-       - litellm_request         (sometimes only metadata, ignore if no token ids)
-       - raw_gen_ai_request      (llm.hosted_vllm.* stringified fields)
-    3) Extract rewards from spans whose attributes contain an AgentOps-style
-       reward payload or explicit REWARD span.
-    4) For each reward with sequence R, assign it to the most recent *unmatched* LLM call
-       with sequence < R. Ignore timestamps completely.
+    1. Sort spans by `(sequence_id, start_time)` for deterministic processing.
+    2. Extract token identifiers from `litellm_request` or `raw_gen_ai_request` spans.
+    3. Extract rewards from spans exposing AgentOps-style payloads or explicit reward spans.
+    4. Match each reward to the most recent unmatched LLM call whose sequence is smaller.
    """

    def _literal_eval_maybe(self, v: Any) -> Any:
@@ -681,9 +766,7 @@ class LlmProxyTraceToTriplet(TraceToTripletBase):
        return cast(List[int], prompt_ids), cast(List[int], resp_ids)

    def _maybe_reward_value(self, span: Span) -> Optional[float]:
-        """
-        Parse reward from typical AgentOps payload or explicit REWARD span.
-        """
+        """Parse reward from typical AgentOps payloads or explicit reward spans."""
        attrs = span.attributes or {}

        # AgentOps new/old keys
@@ -709,6 +792,14 @@ class LlmProxyTraceToTriplet(TraceToTripletBase):
        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,
@@ -4,7 +4,7 @@ from __future__ import annotations

 from typing import TYPE_CHECKING, Any

-from .base import BaseAlgorithm
+from .base import Algorithm
 from .decorator import algo
 from .fast import Baseline, FastAlgorithm

@@ -12,7 +12,7 @@ if TYPE_CHECKING:
    from .apo import APO as APOType
    from .verl import VERL as VERLType

-__all__ = ["BaseAlgorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]
+__all__ = ["Algorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]

 # Shortcuts for usages like algo.APO(...)

@@ -19,7 +19,7 @@ import poml
 from openai import AsyncOpenAI

 from agentlightning.adapter.messages import TraceToMessages
-from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.algorithm.base import Algorithm
 from agentlightning.reward import find_final_reward
 from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus

@@ -91,7 +91,7 @@ def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterat
                    current_batch = []


-class APO(BaseAlgorithm, Generic[T_task]):
+class APO(Algorithm, Generic[T_task]):
    """Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search.

    APO is an iterative prompt optimization algorithm that uses LLM-generated textual gradients
@@ -99,14 +99,16 @@ class APO(BaseAlgorithm, Generic[T_task]):
    computes critiques based on the results, and applies edits to generate improved prompts.

    The algorithm operates in rounds, where each round:
+
    1. Samples parent prompts from the current beam
    2. Generates new prompts by computing textual gradients and applying edits
    3. Evaluates all candidates on a validation set
    4. Selects the top-k prompts for the next round

    Based on the ideas from:
-    - ProTeGi: https://aclanthology.org/2023.emnlp-main.494.pdf
-    - TextGrad: https://github.com/zou-group/textgrad
+
+    - [ProTeGi](https://aclanthology.org/2023.emnlp-main.494.pdf)
+    - [TextGrad](https://github.com/zou-group/textgrad)
    """

    def __init__(
@@ -337,6 +339,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
        Generate an improved prompt by computing a textual gradient and applying an edit.

        This is the main optimization step that:
+
        1. Computes a critique (textual gradient) based on rollout performance
        2. Uses another LLM to apply the critique and generate an improved prompt

@@ -443,6 +446,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
        Evaluate a prompt on a batch of tasks by running rollouts and computing average reward.

        This method:
+
        1. Adds the prompt as a named resource to the store
        2. Enqueues rollouts for each task in the dataset
        3. Waits for rollouts to complete (with timeout)
@@ -587,6 +591,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
        Generate new candidate prompts from parents using textual gradients.

        For each parent prompt, generates branch_factor new candidates by:
+
        1. Evaluating the parent on a training batch
        2. Computing textual gradient
        3. Applying edit to generate improved prompt
@@ -814,6 +819,7 @@ class APO(BaseAlgorithm, Generic[T_task]):
        Execute the APO algorithm to optimize prompts through beam search with textual gradients.

        The algorithm performs iterative prompt optimization over multiple rounds:
+
        - Each round: samples parent prompts, generates new candidates via textual gradients,
          evaluates all candidates on validation data, and keeps the top performers
        - Tracks the historically best prompt across all rounds
@@ -22,7 +22,7 @@ if TYPE_CHECKING:
    from agentlightning.trainer import Trainer


-class BaseAlgorithm:
+class Algorithm:
    """Algorithm is the strategy, or tuner to train the agent."""

    _trainer_ref: weakref.ReferenceType[Trainer] | None = None
@@ -26,7 +26,7 @@ from agentlightning.types import Dataset, NamedResources
 if TYPE_CHECKING:
    from agentlightning.llm_proxy import LLMProxy

-from .base import BaseAlgorithm
+from .base import Algorithm

 # Algorithm function signature types
 # We've missed a lot of combinations here.
@@ -100,12 +100,13 @@ AsyncFlag = Literal[True, False]
 AF = TypeVar("AF", bound=AsyncFlag)


-class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
-    """A BaseAlgorithm that wraps a function-based algorithm implementation.
+class FunctionalAlgorithm(Algorithm, Generic[AF]):
+    """An algorithm wrapper built from a callable implementation.

-    This class allows users to define algorithm behavior using a simple function
-    that takes train_dataset and val_dataset parameters, rather than implementing
-    a full BaseAlgorithm subclass.
+    Functional algorithms let you provide an ordinary function instead of
+    subclassing [`Algorithm`][agentlightning.Algorithm]. The wrapper inspects
+    the callable signature to supply optional dependencies
+    such as the store, adapter, and LLM proxy.
    """

    @overload
@@ -115,13 +116,12 @@ class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
    def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...

    def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
-        """
-        Initialize the FunctionalAlgorithm with an algorithm function.
+        """Wrap a function that implements algorithm behaviour.

        Args:
-            algorithm_func: A function that defines the algorithm's behavior.
-                           Can be sync or async with signature:
-                           (train_dataset, val_dataset) -> None
+            algorithm_func: Sync or async callable implementing the algorithm
+                contract. Arguments are detected automatically based on the
+                function signature.
        """
        super().__init__()
        self._algorithm_func = algorithm_func
@@ -156,14 +156,20 @@ class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> Union[None, Awaitable[None]]:
-        """Execute the algorithm using the wrapped function.
+        """Execute the wrapped function with injected dependencies.

        Args:
-            train_dataset: The dataset to train on.
-            val_dataset: The dataset to validate on.
+            train_dataset: Optional training dataset passed through when the
+                callable declares a `train_dataset` parameter.
+            val_dataset: Optional validation dataset passed through when the
+                callable declares a `val_dataset` parameter.

        Returns:
-            None or Awaitable[None] if the function is async.
+            None for sync callables or an awaitable when the callable is async.
+
+        Raises:
+            TypeError: If a dataset is provided but the function signature does
+                not accept the corresponding argument.
        """
        kwargs: Dict[str, Any] = {}
        if "store" in self._sig.parameters:
@@ -217,40 +223,42 @@ def algo(
        AlgorithmFuncAsyncFallback,
    ],
 ) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
-    """Create a BaseAlgorithm from a function.
+    """Convert a callable into a [`FunctionalAlgorithm`][agentlightning.algorithm.decorator.FunctionalAlgorithm].

-    This decorator allows you to define an algorithm using a simple function
-    instead of creating a full BaseAlgorithm subclass. The returned FunctionalAlgorithm
-    instance is callable, preserving the original function's behavior.
+    The decorator inspects the callable signature to decide which dependencies
+    to inject at runtime, enabling concise algorithm definitions that still
+    leverage the full training runtime.

    Args:
-        func: A function that defines the algorithm's behavior with signature:
-              (train_dataset, val_dataset) -> None
-              Can be sync or async.
+        func: Function implementing the algorithm logic. May be synchronous or
+            asynchronous. The function can expect all of, or a subset of the following parameters:
+
+            - `store`: [`LightningStore`][agentlightning.store.base.LightningStore],
+            - `train_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `val_dataset`: [`Dataset`][agentlightning.Dataset],
+            - `llm_proxy`: [`LLMProxy`][agentlightning.LLMProxy],
+            - `adapter`: [`TraceAdapter`][agentlightning.TraceAdapter],
+            - `initial_resources`: [`NamedResources`][agentlightning.NamedResources],
+
+            If the function does not expect a parameter, the wrapper will not inject it into the call.
+            Using `*args` and `**kwargs` will not work and no parameters will be injected.

    Returns:
-        A callable FunctionalAlgorithm instance that preserves the original function's
-        type hints and behavior while providing all algorithm functionality.
+        FunctionalAlgorithm that proxies the callable while exposing the
+        `Algorithm` interface.

-    Example:
-        @algo
-        def my_algorithm(train_dataset, val_dataset):
-            # Algorithm logic here
-            for task in train_dataset:
-                # Process training tasks
-                pass
+    Examples:
+        ```python
+        from agentlightning.algorithm.decorator import algo

        @algo
-        async def my_async_algorithm(train_dataset, val_dataset):
-            # Async algorithm logic here
-            async for task in train_dataset:
-                # Process training tasks asynchronously
-                pass
+        def batching_algorithm(*, store, train_dataset, val_dataset):
+            for sample in train_dataset:
+                store.enqueue_rollout(input=sample, mode="train")

-        # Function is still callable with original behavior
-        my_algorithm(train_data, val_data)
-
-        # Algorithm methods are also available
-        my_algorithm.run(train_data, val_data)
+        @algo
+        async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
+            await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
+        ```
    """
    return FunctionalAlgorithm(func)
@@ -7,21 +7,21 @@ import logging
 from datetime import datetime
 from typing import Any, List, Literal, Optional

-from agentlightning.llm_proxy import ModelConfig
 from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span

-from .base import BaseAlgorithm
+from .base import Algorithm

 logger = logging.getLogger(__name__)

 __all__ = ["FastAlgorithm", "Baseline"]


-class FastAlgorithm(BaseAlgorithm):
-    """Algorithm that can run fast and qualify for dev mode.
+class FastAlgorithm(Algorithm):
+    """Base class for lightweight algorithms optimised for developer workflows.

-    Fast algorithms enable agent developers to quickly iterate on agent development
-    without waiting for a long training to complete.
+    Fast algorithms prioritise short feedback loops so an agent developer can run
+    small-scale experiments without waiting for long-running training jobs to
+    finish.
    """


@@ -30,24 +30,38 @@ def _timestamp_to_iso_str(timestamp: float) -> str:


 class Baseline(FastAlgorithm):
-    """A dummy implementation of algorithm interface that puts all dataset into the queue, and waits for all rollouts to complete.
+    """Reference implementation that streams the full dataset through the rollout queue.

-    Logs all collected spans and rewards.
+    The baseline algorithm batches task submissions, waits for each rollout to
+    finish, and logs every collected span and reward. It is primarily useful as
+    a smoke test for the platform plumbing rather than a performant trainer.

    Args:
-        model_list: Optional list of models to load into the llm proxy.
-            If both model_list and llm_proxy is provided, llm_proxy will be launched.
-            Not implemented yet.
-        n_epochs: Number of epochs to run through the dev dataset.
-        train_split: Fraction of dev dataset to use for training vs validation. Must be between 0 and 1.
-        polling_interval: Time interval (in seconds) to poll the store for queue length and for completed rollouts.
-        max_queue_length: Maximum number of rollouts to keep in the queue at any time.
+        n_epochs: Number of dataset passes to execute for both the train and val
+            splits during developer experiments.
+        train_split: Fraction of the concatenated dataset to treat as training
+            data. Must be strictly between 0 and 1.
+        polling_interval: Interval, in seconds, to poll the store for queue
+            depth and rollout completion.
+        max_queue_length: Number of rollouts allowed to wait in the queue before
+            throttling additional submissions.
+        span_verbosity: Level of detail to include when logging span metadata.
+
+    Raises:
+        ValueError: If `train_split` falls outside the `(0, 1)` interval.
+
+    Examples:
+        ```python
+        from agentlightning.algorithm.fast import Baseline
+
+        algorithm = Baseline(n_epochs=2, train_split=0.8, span_verbosity="key_values")
+        trainer.fit(algorithm, train_dataset=my_train, val_dataset=my_val)
+        ```
    """

    def __init__(
        self,
        *,
-        model_list: Optional[List[ModelConfig]] = None,
        n_epochs: int = 1,
        train_split: float = 0.5,
        polling_interval: float = 5.0,
@@ -66,6 +80,7 @@ class Baseline(FastAlgorithm):
        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 ""

@@ -85,6 +100,7 @@ class Baseline(FastAlgorithm):
        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
@@ -97,7 +113,12 @@ class Baseline(FastAlgorithm):
        attempts = await store.query_attempts(rollout_id)
        for attempt in attempts:
            logger.info(
-                f"[Rollout {rollout_id} | Attempt {attempt.sequence_id}] ID: {attempt.attempt_id}. Status: {attempt.status}. Worker: {attempt.worker_id}"
+                "[Rollout %s | Attempt %s] ID: %s. Status: %s. Worker: %s",
+                rollout_id,
+                attempt.sequence_id,
+                attempt.attempt_id,
+                attempt.status,
+                attempt.worker_id,
            )
            spans = await store.query_spans(rollout_id=rollout_id)
            for span in spans:
@@ -107,15 +128,18 @@ class Baseline(FastAlgorithm):
        # Attempts to adapt the spans using the adapter if provided
        try:
            adapter = self.get_adapter()
+        except ValueError:
+            logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
+            adapter = None
+        if adapter is not None:
            spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
            transformed_data = adapter.adapt(spans)
            logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}")
-        except ValueError:
-            logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")

    async def _enqueue_rollouts(
        self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
    ) -> None:
+        """Submit rollouts while respecting the maximum queue length."""
        store = self.get_store()

        for index in train_indices + val_indices:
@@ -129,6 +153,7 @@ class Baseline(FastAlgorithm):
            await asyncio.sleep(self.polling_interval)

    async def _harvest_rollout_spans(self, rollout_id: str):
+        """Poll rollout status updates until completion and log transitions."""
        store = self.get_store()
        last_status: Optional[RolloutStatus] = None

@@ -160,11 +185,12 @@ class Baseline(FastAlgorithm):
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> None:
+        """Execute the baseline loop across the provided datasets."""
        train_dataset_length = len(train_dataset) if train_dataset is not None else 0
        val_dataset_length = len(val_dataset) if val_dataset is not None else 0
        if train_dataset_length == 0 and val_dataset_length == 0:
            logger.error(
-                "MockAlgorithm requires at least a train_dataset or val_dataset to run. No train_dataset or val_dataset is provided. Exiting."
+                "MockAlgorithm requires at least one dataset. Provide train_dataset or val_dataset before running."
            )
            return

@@ -173,6 +199,8 @@ class Baseline(FastAlgorithm):
        ]
        train_indices = list(range(0, train_dataset_length))
        val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length))
+        logger.debug(f"Train indices: {train_indices}")
+        logger.debug(f"Val indices: {val_indices}")

        store = self.get_store()

@@ -190,19 +218,24 @@ class Baseline(FastAlgorithm):
            harvest_tasks: List[asyncio.Task[None]] = []
            logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.")
            for index in train_indices + val_indices:
-                queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
-                if len(queuing_rollouts) <= self.max_queue_length:
-                    # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
-                    sample = concatenated_dataset[index]
-                    mode = "train" if index in train_indices else "val"
-                    rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
-                    harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
-                    logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
-                else:
-                    # Sleep a bit and try again later.
-                    await asyncio.sleep(self.polling_interval)
+                logger.info(
+                    f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
+                )
+                while True:
+                    queuing_rollouts = await store.query_rollouts(status=["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
-            print(f"Waiting for {len(harvest_tasks)} 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)
@@ -5,23 +5,89 @@ from typing import Any, Optional
 from hydra import compose, initialize
 from omegaconf import OmegaConf

-from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.algorithm.base import Algorithm
 from agentlightning.client import AgentLightningClient
 from agentlightning.types import Dataset
 from agentlightning.verl.entrypoint import run_ppo  # type: ignore


-class VERL(BaseAlgorithm):
-    """Algorithm leveraging VERL as the backend framework.
+class VERL(Algorithm):
+    """VERL-powered algorithm that delegates training to the VERL PPO runner.

-    **Note on Customization:**
-
-    At present, we recommend copying the source code from VERL and modifying it as needed to suit your requirements.
-    Native support for customizing training logic will be provided in future releases.
+    !!! warning
+        Advanced customisation currently requires copying the VERL source and
+        modifying it directly. Native hooks for overriding training behaviour
+        will land in a future release.

    Args:
-        config: The VERL configuration, matching what is typically provided when running VERL via the command line.
-            This config will be merged with VERL's base configuration and processed by Hydra.
+        config: Dictionary mirroring the overrides passed to the VERL CLI. The
+            overrides are merged with VERL's packaged defaults via Hydra before
+            launching training.
+
+    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]):
@@ -40,6 +106,17 @@ class VERL(BaseAlgorithm):
        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:
@@ -66,5 +143,10 @@ class VERL(BaseAlgorithm):
            )

    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}")
@@ -6,12 +6,15 @@ 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")
@@ -22,7 +25,11 @@ def main(argv: Iterable[str] | None = None) -> int:

    store = InMemoryLightningStore()
    server = LightningStoreServer(store, host="0.0.0.0", port=args.port)
-    asyncio.run(server.run_forever())
+    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


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

-"""Legacy client for interacting with a legacy Agent Lightning server."""
+"""Utilities for interacting with legacy Agent Lightning servers.
+
+This module contains compatibility shims that speak the deprecated HTTP
+interface used by older Agent Lightning deployments. Modern code should prefer
+the store-based APIs exposed by `agentlightning.store`, but keeping these
+clients available makes it easier to migrate existing workflows incrementally.
+"""

 import asyncio
 import logging
@@ -18,13 +24,24 @@ logger = logging.getLogger(__name__)


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

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

    _next_task_uri = "/task"
@@ -33,12 +50,12 @@ class AgentLightningClient:
    _report_rollout_uri = "/rollout"

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        Returns:
-            The JSON response as a dictionary or None if the request fails.
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
        """
        try:
            response = requests.post(url, json=payload, timeout=self.timeout, headers=self._default_headers)
@@ -193,10 +214,11 @@ class AgentLightningClient:
            return None

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

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

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

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

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

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

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

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

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

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


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

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

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

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

        Args:
-            tasks: Either a List of TaskInput objects or a List of Task objects.
-            resources: Either NamedResources or ResourcesUpdate object.
-            **kwargs: Additional arguments passed to the parent AgentLightningClient.
+            tasks: Sequence of task inputs or preconstructed tasks that will be served in
+                order.
+            resources: Static resources returned for any `resources_id` query.
+            **kwargs: Additional keyword arguments forwarded to the parent client.
+
+        Raises:
+            ValueError: If no tasks are provided or both [`Task`][agentlightning.Task]
+                and [`TaskInput`][agentlightning.TaskInput] instances are mixed.
        """
        warnings.warn("DevTaskLoader is deprecated. Please use Trainer.dev instead.", DeprecationWarning)
        super().__init__(endpoint="local://", **kwargs)
@@ -307,17 +339,18 @@ class DevTaskLoader(AgentLightningClient):

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

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

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

        Returns:
-            The next Task object from the local task list.
+            Next task to execute.
        """
        if self._task_index >= len(self._tasks):
            self._task_index = 0
@@ -83,12 +83,17 @@ def _str_to_bool(v: str) -> bool:


 def _get_param_type_details(param_annotation: Any) -> Tuple[Any, bool, bool]:
-    """
-    Determines the core type, if it's Optional, and if it's a List.
-    Returns: (core_type, is_optional, is_list)
-    - For Optional[T]: (T, True, is_list_status_of_T)
-    - For List[T]: (List[T], is_optional_status_of_List, True)
-    - For Optional[List[T]]: (List[T], True, True)
+    """Normalize an annotation into its core type, optionality, and list status.
+
+    Args:
+        param_annotation: The annotation to inspect.
+
+    Returns:
+        A tuple ``(core_type, is_optional, is_list)`` describing the normalized type.
+
+        - For ``Optional[T]`` → ``(T, True, is_list_status_of_T)``
+        - For ``List[T]`` → ``(List[T], is_optional_status_of_List, True)``
+        - For ``Optional[List[T]]`` → ``(List[T], True, True)``
    """
    is_optional = False
    is_list = False
@@ -13,7 +13,15 @@ logger = logging.getLogger(__name__)


 def emit_exception(exception: BaseException) -> None:
-    """Emit an exception as a span."""
+    """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
@@ -10,10 +10,14 @@ logger = logging.getLogger(__name__)


 def emit_message(message: str) -> None:
-    """Emit a string message as a span.
+    """Emit a textual message as an OpenTelemetry span.

-    OpenTelemetry has a dedicated design of logs by design, but we can also use spans to emit messages.
-    So that it can all be unified in the data store and analyzed together.
+    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.")
@@ -12,7 +12,15 @@ logger = logging.getLogger(__name__)


 def emit_object(object: Any) -> None:
-    """Emit any object as a span. Make sure the object is JSON serializable."""
+    """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):
@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.

+"""Helpers for emitting reward spans and integrating with AgentOps telemetry."""
+
 import asyncio
 import inspect
 import json
@@ -47,20 +49,29 @@ FnType = TypeVar("FnType", bound=Callable[..., Any])


 def _agentops_initialized() -> bool:
-    """Check if AgentOps is initialized in the current context."""
+    """Return `True` when the AgentOps client has been configured."""
    return agentops.get_client().initialized


 def reward(fn: FnType) -> FnType:
-    """
-    A decorator to wrap a function that computes rewards.
-    It will automatically handle the input and output of the function.
+    """Decorate a reward function so its outputs are tracked as spans.
+
+    The decorator integrates with AgentOps when it is available and falls back to
+    the built-in telemetry otherwise. Both synchronous and asynchronous functions
+    are supported transparently.
+
+    Deprecated:
+        This decorator is deprecated. Use [`emit_reward`][agentlightning.emit_reward] instead.
+
+    Args:
+        fn: Callable that produces a numeric reward.
+
+    Returns:
+        Wrapped callable that preserves the original signature.
    """

    def wrap_result(result: Optional[float]) -> RewardSpanData:
-        """
-        Wrap the result of the function in a dict.
-        """
+        """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
@@ -119,8 +130,18 @@ def reward(fn: FnType) -> FnType:


 def emit_reward(reward: float) -> ReadableSpan:
-    """
-    Record a new reward as a new span.
+    """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)):
@@ -128,6 +149,7 @@ def emit_reward(reward: float) -> ReadableSpan:
    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
@@ -139,8 +161,13 @@ def emit_reward(reward: float) -> ReadableSpan:


 def get_reward_value(span: SpanLike) -> Optional[float]:
-    """
-    Get the reward value from a span.
+    """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
@@ -178,35 +205,31 @@ def get_reward_value(span: SpanLike) -> Optional[float]:


 def is_reward_span(span: SpanLike) -> bool:
-    """
-    Check if a span is a reward span.
-    """
+    """Return ``True`` when the provided span encodes a reward value."""
    maybe_reward = get_reward_value(span)
    return maybe_reward is not None


 def find_reward_spans(spans: Sequence[SpanLike]) -> List[SpanLike]:
-    """
-    Find all reward spans in the given list of spans.
+    """Return all reward spans in the provided sequence.

    Args:
-        spans: A list of spans (either ReadableSpan or Span).
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.

    Returns:
-        A list of spans whose name matches the reward span name.
+        List of spans that could be parsed as rewards.
    """
    return [span for span in spans if is_reward_span(span)]


 def find_final_reward(spans: Sequence[SpanLike]) -> Optional[float]:
-    """
-    Get the last reward value from a list of spans.
+    """Return the last reward value present in the provided spans.

    Args:
-        spans: A list of spans (either ReadableSpan or Span).
+        spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.

    Returns:
-        The reward value from the last reward span, or None if not found.
+        Reward value from the latest reward span, or `None` when none are found.
    """
    for span in reversed(spans):
        reward = get_reward_value(span)
@@ -1,19 +1,19 @@
 # Copyright (c) Microsoft. All rights reserved.

-"""Common utilities for the emitter module."""
+"""Utilities shared across emitter implementations."""

 import opentelemetry.trace as trace_api
 from opentelemetry.trace import get_tracer_provider


 def get_tracer() -> trace_api.Tracer:
-    """Return the tracer used for AgentLightning spans.
-
-    Raises:
-        RuntimeError: If the tracer is not initialized.
+    """Resolve the OpenTelemetry tracer configured for Agent Lightning.

    Returns:
-        The AgentLightning tracer instance.
+        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.")
@@ -1,6 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.

+from __future__ import annotations
+
 import logging
+import os
 from typing import Protocol

 from agentlightning.store.base import LightningStore
@@ -10,28 +13,94 @@ 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:
-        """Initalization and execution logic."""
+        """Execute algorithm logic using ``store`` until completion or stop."""


 class RunnerBundle(Protocol):
+    """Callable bundle wrapping runner setup and the worker loop, as opposed to the
+    [`AlgorithmBundle`][agentlightning.AlgorithmBundle]."""
+
    async def __call__(self, store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
-        """Initalization and execution logic."""
+        """Execute runner logic for ``worker_id`` using ``store`` and ``event``."""


 class ExecutionStrategy:
-    """When trainer has created the executable of algorithm and runner in two bundles,
-    the execution strategy defines how to run them together, and how many parallel runners to run.
+    """Coordinate algorithm and runner bundles within a single process abstraction.

-    The store is the centric place for the two bundles to communicate.
+    Strategies decide how many worker bundles to launch, whether to communicate
+    through shared memory or an HTTP boundary, and how to react to shutdown
+    signals. They intentionally avoid inspecting the bundle internals; instead,
+    each bundle remains responsible for its own scheduling semantics.

-    The algorithm and runner's behavior (whether runner should perform one step or run forever,
-    whether the algo would send out the tasks or not) are defined inside the bundle,
-    and does not belong to the execution strategy.
-
-    The execute should support Ctrl+C to exit gracefully.
+    !!! note
+        Implementations must honor the [execute()][agentlightning.ExecutionStrategy.execute]
+        contract by propagating `KeyboardInterrupt` and ensuring resources are
+        released when an error occurs on either side of the algorithm/runner
+        pair.
    """

    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        """Run the provided bundles using the configured orchestration model.
+
+        Args:
+            algorithm: Callable bundle responsible for algorithm execution.
+            runner: Callable bundle for runner workers.
+            store: Concrete [`LightningStore`][agentlightning.LightningStore]
+                shared across bundles.
+
+        Raises:
+            NotImplementedError: Subclasses must provide the orchestration
+                implementation.
+        """
+
        raise NotImplementedError()
@@ -12,52 +12,51 @@ 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
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
 from .events import ExecutionEvent, MultiprocessingEvent

 logger = logging.getLogger(__name__)


 class ClientServerExecutionStrategy(ExecutionStrategy):
-    """Run algorithm (server) and runners (clients) as separate processes over HTTP.
+    """Run algorithm and runner bundles as separate processes over HTTP.

-    **Execution Roles:**
+    Execution Roles:

-    - "algorithm": Start the HTTP server (`LightningStoreServer`) in-process and run the
-      algorithm bundle against it.
-    - "runner": Connect to an already running server via `LightningStoreClient` and
-      execute runner bundles (optionally in multiple processes).
-    - "both": Spawn the runner processes first, then launch the algorithm/server
-      bundle on the main process. This mode orchestrates the full loop locally.
+    - `"algorithm"`: Start [`LightningStoreServer`][agentlightning.LightningStoreServer]
+      in-process and execute the algorithm bundle against it.
+    - `"runner"`: Connect to an existing server with
+      [`LightningStoreClient`][agentlightning.LightningStoreClient] and run the
+      runner bundle locally (spawning multiple processes when requested).
+    - `"both"`: Spawn runner processes first, then execute the algorithm and
+      server on the same machine. This mode orchestrates the full loop locally.

-    When role == "both", you may choose which side runs on the main process via
-    `main_process` (debug helper). Running the runner bundle on the main process
-    is only supported with `n_runners == 1`.
+    When `role == "both"` you may choose which side runs on the main process
+    via `main_process`. The runner-on-main option is limited to
+    `n_runners == 1` because each additional runner requires its own event
+    loop and process.

-    Important: When `main_process == "runner"`, the algorithm runs in a subprocess
-    with the LightningStore server. This means any state modifications made during
-    execution remain in that subprocess and are NOT reflected in the original store
-    object passed to `execute()`. The main process runner accesses the store only
-    through the HTTP client interface.
+    !!! warning
+        When `main_process == "runner"` the algorithm and HTTP server execute
+        in a child process. Store mutations remain isolated inside that process,
+        so the original store instance passed to
+        [execute()][agentlightning.ExecutionStrategy.execute] is not updated.

-    **Abort / Stop Model (four-step escalation):**
+    Abort Model (four-step escalation):

-    1. Cooperative stop:
-       A shared :class:`~agentlightning.execution.events.MultiprocessingEvent`
-       (`stop_evt`) is passed to *all* bundles. Bundles should check it to exit.
-       Any crash (algorithm or runner) sets `stop_evt` so the other side can
-       stop cooperatively. Ctrl+C on the main process also flips the event.
-    2. KeyboardInterrupt synth:
-       Remaining subprocesses receive `SIGINT` to trigger `KeyboardInterrupt`
-       handlers.
-    3. Termination:
-       Stubborn subprocesses get `terminate()` (SIGTERM on POSIX).
-    4. Kill:
-        As a last resort we call `kill()` (SIGKILL on POSIX).
+    1. Cooperative stop. Every bundle receives a shared
+       [`MultiprocessingEvent`][agentlightning.MultiprocessingEvent] (`stop_evt`).
+       Any failure flips the event so peers can exit cleanly. Ctrl+C on the main
+       process also sets the flag.
+    2. KeyboardInterrupt synthesis. Remaining subprocesses receive ``SIGINT`` to
+       trigger `KeyboardInterrupt` handlers.
+    3. Termination. Stubborn processes are asked to ``terminate()``
+       (`SIGTERM` on POSIX).
+    4. Kill. As a last resort `kill()` is invoked (`SIGKILL` on POSIX).

-    Notes:
-        This mirrors the semantics implemented in :mod:`shared_memory`, but adapted
-        to multiple processes and the HTTP client/server boundary.
+    This mirrors the semantics implemented in
+    [`SharedMemoryExecutionStrategy`][agentlightning.SharedMemoryExecutionStrategy]
+    but adapts them to multiple processes and the HTTP client/server boundary.
    """

    alias: str = "cs"
@@ -71,17 +70,18 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
        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
-                :envvar:`AGL_CURRENT_ROLE` environment variable is used.
+                `AGL_CURRENT_ROLE` environment variable is used.
            server_host: Interface the HTTP server binds to when running the
-                algorithm bundle locally. Defaults to :envvar:`AGL_SERVER_HOST`
-                or ``"localhost"`` if unset.
+                algorithm bundle locally. Defaults to `AGL_SERVER_HOST`
+                or `"localhost"` if unset.
            server_port: Port for the HTTP server in "algorithm"/"both" modes.
-                Defaults to :envvar:`AGL_SERVER_PORT` or ``4747`` if unset.
+                Defaults to `AGL_SERVER_PORT` or `4747` if unset.
            n_runners: Number of runner processes to spawn in "runner"/"both".
            graceful_timeout: How long to wait (seconds) after setting the stop
                event before escalating to signals.
@@ -90,14 +90,20 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
            main_process: Which bundle runs on the main process when
                `role == "both"`. `"runner"` requires `n_runners == 1` and is
                primarily intended for debugging.
+            managed_store: When `True` (default) the strategy constructs
+                LightningStore client/server wrappers automatically. When
+                `False` the provided `store` is passed directly to the
+                bundles, allowing callers to manage store wrappers manually.
        """
        if role is None:
            role_env = os.getenv("AGL_CURRENT_ROLE")
            if role_env is None:
-                raise ValueError("role must be provided via argument or AGL_CURRENT_ROLE env var")
-            if role_env not in ("algorithm", "runner", "both"):
+                # 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'")
-            role = role_env
+            else:
+                role = role_env

        if server_host is None:
            server_host = os.getenv("AGL_SERVER_HOST", "localhost")
@@ -126,19 +132,26 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
            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:
-        logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
-        server_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
-        server_started = False
+        wrapper_store: LightningStore | None = None
+        if self.managed_store:
+            logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
+            wrapper_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
+            server_started = False
+        else:
+            wrapper_store = store
+            server_started = False

        try:
-            await server_store.start()
-            server_started = True
-            logger.debug("Algorithm bundle starting against endpoint %s", server_store.endpoint)
-            await algorithm(server_store, stop_evt)
+            if self.managed_store and isinstance(wrapper_store, LightningStoreServer):
+                await wrapper_store.start()
+                server_started = True
+                logger.debug("Algorithm bundle starting against endpoint %s", wrapper_store.endpoint)
+            await algorithm(wrapper_store, stop_evt)
            logger.debug("Algorithm bundle completed successfully")
        except KeyboardInterrupt:
            logger.warning("Algorithm received KeyboardInterrupt; signaling stop event")
@@ -149,18 +162,31 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
            stop_evt.set()
            raise
        finally:
-            if server_started:
+            if self.managed_store and isinstance(wrapper_store, LightningStoreServer) and server_started:
                try:
-                    await server_store.stop()
+                    await wrapper_store.stop()
                except Exception:
                    logger.exception("Error stopping LightningStore server")
                else:
                    logger.debug("LightningStore server shutdown completed")

-    async def _execute_runner(self, runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
-        client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
+    async def _execute_runner(
+        self,
+        runner: RunnerBundle,
+        worker_id: int,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+    ) -> None:
+        if self.managed_store:
+            # If managed, we actually do not use the provided store
+            client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
+        else:
+            client_store = store
        try:
-            logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
+            if self.managed_store:
+                logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
+            else:
+                logger.debug("Runner %s executing with provided store", worker_id)
            await runner(client_store, worker_id, stop_evt)
            logger.debug("Runner %s completed successfully", worker_id)
        except KeyboardInterrupt:
@@ -172,16 +198,18 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
            stop_evt.set()
            raise
        finally:
-            try:
-                await client_store.close()
-            except Exception:
-                logger.exception("Error closing LightningStore client for runner %s", worker_id)
-            else:
-                logger.debug("Runner %s closed LightningStore client", worker_id)
+            if self.managed_store and isinstance(client_store, LightningStoreClient):
+                try:
+                    await client_store.close()
+                except Exception:
+                    logger.exception("Error closing LightningStore client for runner %s", worker_id)
+                else:
+                    logger.debug("Runner %s closed LightningStore client", worker_id)

    def _spawn_runners(
        self,
        runner: RunnerBundle,
+        store: LightningStore,
        stop_evt: ExecutionEvent,
        *,
        ctx: BaseContext,
@@ -189,15 +217,15 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
        """Used when `role == "runner"` or `role == "both"` and `n_runners > 1`."""
        processes: list[multiprocessing.Process] = []

-        def _runner_sync(runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
+        def _runner_sync(runner: RunnerBundle, worker_id: int, store: LightningStore, stop_evt: ExecutionEvent) -> None:
            # Runners are executed in child processes; each process owns its own
            # event loop to keep the asyncio scheduler isolated.
-            asyncio.run(self._execute_runner(runner, worker_id, stop_evt))
+            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, stop_evt), name=f"runner-{i}"),  # type: ignore
+                ctx.Process(target=_runner_sync, args=(runner, i, store, stop_evt), name=f"runner-{i}"),  # type: ignore
            )
            process.start()
            logger.debug("Spawned runner process %s (pid=%s)", process.name, process.pid)
@@ -341,10 +369,10 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
            elif self.role == "runner":
                if self.n_runners == 1:
                    logger.info("Running runner solely...")
-                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+                    asyncio.run(self._execute_runner(runner, 0, store, stop_evt))
                else:
                    logger.info("Spawning runner processes...")
-                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
                    # Wait for the processes to finish naturally.
                    for process in processes:
                        process.join()
@@ -352,7 +380,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
            elif self.role == "both":
                if self.main_process == "algorithm":
                    logger.info("Spawning runner processes...")
-                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
                    try:
                        logger.info("Running algorithm...")
                        asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
@@ -373,7 +401,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                    # the background process spawned above (the provided
                    # store must therefore be picklable when using spawn).
                    logger.info("Running runner...")
-                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+                    asyncio.run(self._execute_runner(runner, 0, store, stop_evt))

                    # Wait for the algorithm process to finish.
                    algorithm_process.join()
@@ -7,15 +7,18 @@ from typing import Optional, Protocol


 class ExecutionEvent(Protocol):
-    """
-    A minimal protocol similar to threading.Event.
+    """Protocol capturing the cooperative stop contract shared by strategies.
+
+    Implementations mirror the API of ``threading.Event`` and
+    ``multiprocessing.Event`` so the rest of the execution layer can remain
+    agnostic to the underlying concurrency primitive.

    Methods:
-        set(): Signal event like a cancellation (idempotent).
-        clear(): Reset to the non-set state.
-        is_set() -> bool: True if event has been signaled.
-        wait(timeout: Optional[float] = None) -> bool:
-            Block until event is set or timeout. Returns True if event has signaled.
+
+        set: Signal cancellation. The call must be idempotent.
+        clear: Reset the event to the unsignaled state.
+        is_set: Return ``True`` when cancellation has been requested.
+        wait: Block until the event is signaled or an optional timeout elapses.
    """

    def set(self) -> None: ...
@@ -25,11 +28,7 @@ class ExecutionEvent(Protocol):


 class ThreadingEvent:
-    """
-    An Event implementation using threading.Event.
-
-    Provides a thread-safe event object for signaling between threads.
-    """
+    """Thread-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""

    __slots__ = ("_evt",)

@@ -50,12 +49,7 @@ class ThreadingEvent:


 class MultiprocessingEvent:
-    """
-    An Event implementation using multiprocessing.Event.
-
-    Provides a process-safe event object for signaling between processes.
-    Optionally accepts a multiprocessing context for custom process start methods.
-    """
+    """Process-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""

    __slots__ = ("_evt",)

@@ -4,6 +4,12 @@ from .base import ExecutionStrategy


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

    alias: str = "ipc"

@@ -10,28 +10,33 @@ 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
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
 from .events import ExecutionEvent, ThreadingEvent

 logger = logging.getLogger(__name__)


 class SharedMemoryExecutionStrategy(ExecutionStrategy):
-    """Run algorithm and runners in a single process with threads sharing memory.
+    """Execute bundles in a single process with cooperative worker threads.

-    Termination & abort model:
+    Stop Model:

-    - One shared ThreadingEvent (`stop_evt`) is passed to *all* bundles.
-    - The main thread (only) receives KeyboardInterrupt on Ctrl+C; we set `stop_evt` there.
-    - If any bundle raises, we set `stop_evt` from that thread to stop the rest.
-    - After the main-thread bundle finishes normally:
-      - If main_thread is "algorithm", we also set `stop_evt` to stop the runners.
-      - If main_thread is "runner", we do not set `stop_evt` to stop the algorithm.
-        We instead wait for the algorithm to finish naturally.
-    - Background threads are daemons; we join briefly and log any stragglers.
+    - All bundles share one [`ThreadingEvent`][agentlightning.ThreadingEvent]
+      named `stop_evt`.
+    - Only the main thread receives `KeyboardInterrupt`. When Ctrl+C occurs we
+      set `stop_evt`.
+    - Any exception raised inside a bundle sets `stop_evt` so other threads can
+      unwind cooperatively.
+    - Once the bundle running on the main thread exits successfully the
+      treatment depends on `main_thread`:
+        - `"algorithm"`: the runners are asked to stop by setting `stop_evt`.
+        - `"runner"`: the algorithm keeps running until it exits naturally.
+    - Background threads are marked as daemons. We join them briefly and log any
+      stragglers before shutting down.

-    Notes: Signals other than SIGINT (e.g., SIGTERM) are not intercepted; we respect
-    Python's default behavior for them.
+    !!! note
+        Signals other than `SIGINT` (such as `SIGTERM`) are not intercepted;
+        Python's default behavior for those signals is preserved.
    """

    alias: str = "shm"
@@ -43,32 +48,39 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
        join_timeout: float = 15.0,
        graceful_delay: float = 5.0,
        poll_interval: float = 0.05,
+        managed_store: bool | None = None,
    ) -> None:
        if main_thread not in ("algorithm", "runner"):
            raise ValueError("main_thread must be 'algorithm' or 'runner'")
        if main_thread == "runner" and n_runners != 1:
-            raise ValueError("When main_thread is 'runner', n_runners must be 1")
+            raise ValueError(
+                "When main_thread is 'runner', n_runners must be 1. "
+                "Either use 'algorithm' on the main thread or set n_runners to 1."
+            )
        self.n_runners = n_runners
        self.main_thread = main_thread
        self.join_timeout = join_timeout
        self.graceful_delay = graceful_delay
        self.poll_interval = poll_interval
+        self.managed_store = resolve_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) Start a watcher task that waits for `stop_evt` *without blocking* the loop
-             by periodically polling the threading event.
-          3) When the stop event flips:
-               a) Give the bundle *graceful_delay* seconds to finish on its own,
-                  because well-behaved bundles will check the event and return.
-               b) If still running after the grace period, cancel the bundle task.
-          4) Ensure both tasks are awaited; swallow `CancelledError` where appropriate.
+
+        1. Start the bundle coroutine as `task`.
+        2. Launch a watcher that polls `stop_evt` without blocking the loop.
+        3. When the stop event flips:
+            a. Give the bundle `graceful_delay` seconds to finish on its own,
+               because well-behaved bundles will check the event and return.
+            b. Cancel the bundle task if it is still running after the grace
+               period.
+        4. Await both tasks and swallow `CancelledError` where appropriate.

        This is a *backup* mechanism for bundles that might not poll the event
-        frequently; cooperative shutdown (checking `stop_evt` yourself) is still preferred.
+        frequently; cooperative shutdown (checking `stop_evt` inside the
+        bundle) remains the preferred approach.
        """
        task: asyncio.Task[Any] = asyncio.create_task(coro)  # type: ignore
        task_exception: Optional[BaseException] = None
@@ -191,7 +203,10 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):

        # Create stop event and thread-safe store.
        stop_evt = ThreadingEvent()
-        thread_safe_store = LightningStoreThreaded(store)
+        if self.managed_store:
+            thread_safe_store = LightningStoreThreaded(store)
+        else:
+            thread_safe_store = store

        thread_exceptions: SimpleQueue[BaseException] = SimpleQueue()
        raised_from_thread: Optional[BaseException] = None
@@ -2,14 +2,16 @@

 from __future__ import annotations

+import json
 import logging
 import multiprocessing
 import signal
 import socket
 import time
-from typing import Any, Callable
+from typing import Any, Callable, no_type_check

 import flask
+import requests
 import setproctitle

 logger = logging.getLogger(__name__)
@@ -39,41 +41,67 @@ def _patch_new_agentops():

    _original_handle_chat_attributes = handle_chat_attributes  # type: ignore

+    @no_type_check
    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):  # type: ignore
-        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)  # type: ignore
-        if return_value is not None and hasattr(return_value, "prompt_token_ids"):  # type: ignore
-            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)  # type: ignore
-        if return_value is not None and hasattr(return_value, "response_token_ids"):  # type: ignore
-            attributes["response_token_ids"] = list(return_value.response_token_ids[0])  # type: ignore
+        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
+        if (
+            return_value is not None
+            and hasattr(return_value, "prompt_token_ids")
+            and return_value.prompt_token_ids is not None
+        ):
+            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
+        if (
+            return_value is not None
+            and hasattr(return_value, "response_token_ids")
+            and return_value.response_token_ids is not None
+        ):
+            attributes["response_token_ids"] = list(return_value.response_token_ids[0])

        # For LiteLLM Proxy (v0.2) with vLLM return_token_ids, response_token_ids now lives in choices
        if (
-            not attributes.get("response_token_ids")
-            and return_value is not None
-            and hasattr(return_value, "choices")  # type: ignore
-            and return_value.choices  # type: ignore
-            and isinstance(return_value.choices, list)  # type: ignore
+            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]  # type: ignore
-            if hasattr(first_choice, "token_ids"):  # type: ignore
-                attributes["response_token_ids"] = list(first_choice.token_ids)  # type: ignore
-            # newer versions of OpenAI client SDK
-            elif hasattr(first_choice, "provider_specific_fields") and "token_ids" in first_choice.provider_specific_fields:  # type: ignore
-                attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])  # type: ignore
+            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]
+                    )

        # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
        if (
            return_value is not None
-            and hasattr(return_value, "http_response")  # type: ignore
-            and return_value.http_response is not None  # type: ignore
-            and hasattr(return_value.http_response, "json")  # type: ignore
+            and hasattr(return_value, "http_response")
+            and return_value.http_response is not None
+            and hasattr(return_value.http_response, "json")
        ):
-            json_data = return_value.http_response.json()  # type: ignore
+            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"])  # type: ignore
-                if "response_token_ids" in json_data:
-                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])  # type: ignore
+                if json_data.get("prompt_token_ids") is not None:
+                    attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])
+                if json_data.get("response_token_ids") is not None:
+                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])

        return attributes

@@ -238,7 +266,19 @@ class AgentOpsServerManager:
        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
+        for attempt in range(20):  # 10 seconds total
+            time.sleep(0.5)  # Brief wait for server to start up
+            try:
+                result = requests.get(f"http://127.0.0.1:{self.server_port}/")
+                if result.status_code == 200:
+                    break
+            except Exception as e:
+                logger.debug(f"Error checking AgentOps server: {e}")
+            logger.warning(f"AgentOps still not ready after {attempt} attempts. Retrying...")
+        else:
+            logger.error(f"AgentOps local server failed to start or exited prematurely.")
+            return
+
        if not self.server_process.is_alive():
            logger.error(f"AgentOps local server failed to start or exited prematurely.")

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

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

 from typing import Any, Optional
@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.

+"""Convenience decorators for building lightweight `LitAgent` implementations."""
+
 from __future__ import annotations

 import functools
@@ -90,24 +92,25 @@ class FunctionalLitAgentFunc(Protocol[T_contra]):


 class FunctionalLitAgent(LitAgent[T]):
-    """A specialized LitAgent that wraps a function-based rollout that accepts
-    dynamically a task input and a configured resource (LLM / prompt template / ...).
+    """Adapter that turns plain rollout functions into [`LitAgent`][agentlightning.LitAgent] instances.

-    This class allows users to define agent behavior using a simple function
-    that takes task input and a resource, rather than implementing a full
-    LitAgent subclass.
+    The helper inspects the wrapped function to determine which resources to
+    inject, allowing both synchronous and asynchronous callables to participate
+    in the training loop without writing a dedicated subclass.
    """

    def __init__(self, rollout_func: FunctionalLitAgentFunc[T], *, strip_proxy: bool = True) -> None:
-        """
-        Initialize the FunctionalLitAgent with a functional rollout function.
+        """Initialize the wrapper around a rollout function.

        Args:
-            rollout_func: A function that defines the agent's behavior.
-                          Can be sync or async, and can optionally accept a Rollout parameter.
-                          The function signature determines which resources are injected (llm, prompt_template, etc.).
-            strip_proxy: Whether to strip the ProxyLLM resource into a LLM resource when the function accepts an llm parameter.
-                         Defaults to True.
+            rollout_func: Callable that implements the rollout. It may be synchronous
+                or asynchronous and can optionally receive a
+                [`Rollout`][agentlightning.Rollout] alongside resources such as
+                `llm` or `prompt_template`.
+            strip_proxy: When ``True``, convert
+                [`ProxyLLM`][agentlightning.ProxyLLM] inputs into
+                [`LLM`][agentlightning.LLM] instances before calling the
+                rollout function. Defaults to `True`.
        """
        super().__init__()
        self._rollout_func = rollout_func
@@ -138,12 +141,15 @@ class FunctionalLitAgent(LitAgent[T]):
        """Execute a synchronous rollout using the wrapped function.

        Args:
-            task: The task input data.
-            resources: Dictionary of named resources including LLMs.
-            rollout: The rollout object with metadata.
+            task: Task input data.
+            resources: Mapping of named resources available to the agent.
+            rollout: Rollout metadata provided by the runtime.

        Returns:
-            The result from the wrapped rollout function.
+            Result produced by the wrapped rollout function.
+
+        Raises:
+            RuntimeError: If the wrapped function is asynchronous.
        """
        if self._is_async:
            raise RuntimeError(f"{self._rollout_func} is asynchronous. Use rollout_async instead.")
@@ -155,12 +161,15 @@ class FunctionalLitAgent(LitAgent[T]):
        """Execute an asynchronous rollout using the wrapped function.

        Args:
-            task: The task input data.
-            resources: Dictionary of named resources including LLMs.
-            rollout: The rollout object with metadata.
+            task: Task input data.
+            resources: Mapping of named resources available to the agent.
+            rollout: Rollout metadata provided by the runtime.

        Returns:
-            The result from the wrapped rollout function.
+            Result produced by the wrapped rollout coroutine.
+
+        Raises:
+            RuntimeError: If the wrapped function is synchronous.
        """
        if not self._is_async:
            raise RuntimeError(f"{self._rollout_func} is synchronous. Use rollout instead.")
@@ -169,18 +178,19 @@ class FunctionalLitAgent(LitAgent[T]):
        return await self._rollout_func(task, **kwargs)  # type: ignore

    def _get_kwargs(self, resources: NamedResources, rollout: Rollout) -> Dict[str, Any]:
-        """Extract the kwargs needed for the rollout function based on its signature.
+        """Prepare keyword arguments expected by the wrapped rollout function.

-        Dynamically builds the kwargs dictionary by inspecting the function signature and
+
+        It dynamically builds the `kwargs` dictionary by inspecting the function signature and
        including only the parameters the function accepts. This allows flexible function
        signatures that can request any combination of: rollout, llm, and/or prompt_template.

        Args:
-            resources: Dictionary of named resources available for the rollout.
-            rollout: The rollout object with metadata.
+            resources: Mapping of named resources available for the rollout.
+            rollout: Rollout metadata provided by the runtime.

        Returns:
-            A dictionary of kwargs to pass to the rollout function.
+            Dictionary of keyword arguments to forward to the rollout function.
        """

        kwargs: Dict[str, Any] = {}
@@ -194,19 +204,19 @@ class FunctionalLitAgent(LitAgent[T]):
        return kwargs

    def _get_llm_resource(self, resources: NamedResources, rollout: Rollout) -> LLM:
-        """Extract the first LLM resource from the resources dictionary.
+        """Retrieve the first LLM resource from the available resources.

        Strip the ProxyLLM resource into a LLM resource if needed.

        Args:
-            resources: Dictionary of named resources.
-            rollout: The rollout object with metadata.
+            resources: Mapping of named resources.
+            rollout: Rollout metadata used when stripping proxy endpoints.

        Returns:
-            The first LLM resource found.
+            First [`LLM`][agentlightning.LLM] resource encountered.

        Raises:
-            ValueError: If no LLM resource is found.
+            ValueError: If no LLM resource is present.
        """
        resource_found: LLM | None = None
        for name, resource in resources.items():
@@ -225,17 +235,17 @@ class FunctionalLitAgent(LitAgent[T]):
        return resource_found

    def _get_prompt_template_resource(self, resources: NamedResources, rollout: Rollout) -> PromptTemplate:
-        """Extract the first PromptTemplate resource from the resources dictionary.
+        """Retrieve the first prompt template resource from the available resources.

        Args:
-            resources: Dictionary of named resources.
-            rollout: The rollout object with metadata. Not used in this method.
+            resources: Mapping of named resources.
+            rollout: Rollout metadata (unused).

        Returns:
-            The first PromptTemplate resource found.
+            First [`PromptTemplate`][agentlightning.PromptTemplate] resource encountered.

        Raises:
-            ValueError: If no PromptTemplate resource is found.
+            ValueError: If no prompt template resource is present.
        """
        resource_found: PromptTemplate | None = None
        for name, resource in resources.items():
@@ -253,21 +263,22 @@ class FunctionalLitAgent(LitAgent[T]):
        return resource_found

    def _strip_proxy_helper(self, proxy_llm: LLM, rollout: Rollout) -> LLM:
-        """Strip the ProxyLLM resource into a concrete LLM resource.
+        """Convert [`ProxyLLM`][agentlightning.ProxyLLM] instances into concrete LLMs.

-        This method resolves ProxyLLM instances to their concrete LLM implementation
+        It resolves ProxyLLM instances to their concrete LLM implementation
        by attaching the attempted rollout context. This is only used when the function
-        signature accepts an 'llm' parameter and strip_proxy is True.
+        signature accepts an `llm` parameter and strip_proxy is True.

        Args:
-            proxy_llm: The LLM resource, which may be a ProxyLLM.
-            rollout: The rollout object with metadata.
+            proxy_llm: Candidate LLM resource.
+            rollout: Rollout metadata that provides rollout and attempt identifiers.

        Returns:
-            The concrete LLM resource.
+            [`LLM`][agentlightning.LLM] with rollout context baked into the endpoint.

        Raises:
-            ValueError: If the rollout is not an AttemptedRollout (required for stripping ProxyLLM).
+            ValueError: If the rollout is not an
+                [`AttemptedRollout`][agentlightning.AttemptedRollout].
        """

        if not isinstance(proxy_llm, ProxyLLM):
@@ -293,41 +304,37 @@ def llm_rollout(*, strip_proxy: bool = True) -> Callable[[LlmRolloutFunc[T]], Fu
 def llm_rollout(
    func: LlmRolloutFunc[T] | None = None, *, strip_proxy: bool = True
 ) -> FunctionalLitAgent[T] | Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]:
-    """Create a FunctionalLitAgent from a function that takes (task, llm[, rollout]).
-
-    This decorator allows you to define an agent using a simple function
-    instead of creating a full LitAgent subclass. The returned FunctionalLitAgent
-    instance is callable, preserving the original function's behavior.
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for LLM-based rollouts.

    Args:
-        func: A function that defines the agent's behavior. Can be:
-              - sync: (task, llm) -> result
-              - sync with rollout: (task, llm, rollout) -> result
-              - async: async (task, llm) -> result
-              - async with rollout: async (task, llm, rollout) -> result
-        strip_proxy: Whether to strip the ProxyLLM resource into a LLM resource.
-                     Defaults to True.
+        func: Callable defining the agent's behaviour. Supported signatures include:
+
+            * `(task, llm) -> result`
+            * `(task, llm, rollout) -> result`
+            * `async (task, llm) -> result`
+            * `async (task, llm, rollout) -> result`
+
+        strip_proxy: When `True`, convert proxy resources into concrete
+            [`LLM`][agentlightning.LLM] instances before calling the
+            function. Defaults to `True`.

    Returns:
-        A callable FunctionalLitAgent instance that preserves the original function's
-        type hints and behavior while providing all agent functionality.
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.

-    Example:
+    Examples:
+        ```python
        @llm_rollout
        def my_agent(task, llm):
-            # Agent logic here
-            return response
+            return llm.endpoint

        @llm_rollout(strip_proxy=False)
        def my_agent_no_strip(task, llm):
-            # Agent logic here
-            return response
+            return llm.model

-        # Function is still callable with original behavior
        result = my_agent(task, llm)
-
-        # Agent methods are also available
        result = my_agent.rollout(task, resources, rollout)
+        ```
    """

    def decorator(f: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]:
@@ -343,19 +350,20 @@ def llm_rollout(


 def _validate_llm_rollout_func(func: Any) -> TypeGuard[LlmRolloutFunc[Any]]:
-    """Validate the function signature of a LLM rollout function.
+    """Validate the function signature of an LLM rollout function.

    Ensures the function follows the expected pattern for LLM-based rollouts:
+
    - Must have at least 2 parameters
    - First parameter must be named 'task'
    - Must have a parameter named 'llm'
    - Optionally can have a 'rollout' parameter

    Args:
-        func: The function to validate.
+        func: Function to inspect.

    Returns:
-        True if the function signature is valid.
+        `True` when the signature matches the supported patterns.

    Raises:
        ValueError: If the function signature does not match the expected pattern.
@@ -383,36 +391,34 @@ def prompt_rollout() -> Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]:
 def prompt_rollout(
    func: PromptRolloutFunc[T] | None = None,
 ) -> FunctionalLitAgent[T] | Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]:
-    """Create a FunctionalLitAgent from a function that takes (task, prompt_template[, rollout]).
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for prompt-based rollouts.

    This decorator is designed for agents that work with tunable prompt templates. It enables
    a workflow where algorithms manage and optimize the prompt template, while agents consume
    the template to perform rollouts. This is particularly useful for prompt optimization scenarios.

    Args:
-        func: A function that defines the agent's behavior. Can be:
-              - sync: (task, prompt_template) -> result
-              - sync with rollout: (task, prompt_template, rollout) -> result
-              - async: async (task, prompt_template) -> result
-              - async with rollout: async (task, prompt_template, rollout) -> result
+        func: Callable defining the agent's behavior. Supported signatures include:
+
+            * `(task, prompt_template) -> result`
+            * `(task, prompt_template, rollout) -> result`
+            * `async (task, prompt_template) -> result`
+            * `async (task, prompt_template, rollout) -> result`

    Returns:
-        A callable FunctionalLitAgent instance that preserves the original function's
-        type hints and behavior while providing all agent functionality.
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.

-    Example:
+    Examples:
+        ```python
        @prompt_rollout
        def my_agent(task, prompt_template):
-            # Use the prompt template to generate a response
            messages = prompt_template.format(task=task.input)
-            # ... perform rollout with the formatted prompt
-            return response
+            return messages

-        # Function is still callable with original behavior
        result = my_agent(task, prompt_template)
-
-        # Agent methods are also available
        result = my_agent.rollout(task, resources, rollout)
+        ```
    """

    def decorator(f: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]:
@@ -429,16 +435,17 @@ def _validate_prompt_rollout_func(func: Any) -> TypeGuard[PromptRolloutFunc[Any]
    """Validate the function signature of a prompt rollout function.

    Ensures the function follows the expected pattern for prompt-template-based rollouts:
+
    - Must have at least 2 parameters
    - First parameter must be named 'task'
    - Must have a parameter named 'prompt_template'
    - Optionally can have a 'rollout' parameter

    Args:
-        func: The function to validate.
+        func: Function to inspect.

    Returns:
-        True if the function signature is valid.
+        `True` when the signature matches the supported patterns.

    Raises:
        ValueError: If the function signature does not match the expected pattern.
@@ -456,23 +463,30 @@ def _validate_prompt_rollout_func(func: Any) -> TypeGuard[PromptRolloutFunc[Any]


 def rollout(func: Union[LlmRolloutFunc[T], PromptRolloutFunc[T], Callable[..., Any]]) -> FunctionalLitAgent[T]:
-    """Create a LitAgent from a function, automatically detecting the appropriate type.
+    """Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] from an arbitrary rollout function.

    This function inspects the provided callable and creates the appropriate
    agent type based on its signature. It supports both LLM-based and prompt-template-based
    agents. The returned agent instance is callable, preserving the original function's
    behavior and type hints.

+    See [`llm_rollout`][agentlightning.litagent.decorator.llm_rollout] and
+    [`prompt_rollout`][agentlightning.litagent.decorator.prompt_rollout] for more details.
+
    Args:
-        func: A function that defines the agent's behavior. Supported signatures:
-              - (task, llm[, rollout]) for LLM-based agents
-              - (task, prompt_template[, rollout]) for prompt-template-based agents
+        func: Callable that implements the rollout. Supported signatures:
+
+            - `[async ](task, llm[, rollout])` for LLM-based agents
+            - `[async ](task, prompt_template[, rollout])` for prompt-template-based agents
+
+            The supported output types of `func` is same as the return type of [`rollout`][agentlightning.LitAgent.rollout].

    Returns:
-        A callable FunctionalLitAgent instance that preserves the original function's
-        type hints and behavior while providing all agent functionality.
+        [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
+        wraps the supplied function.

-    Example:
+    Examples:
+        ```python
        # LLM-based agent
        @rollout
        def my_llm_agent(task, llm):
@@ -495,6 +509,7 @@ def rollout(func: Union[LlmRolloutFunc[T], PromptRolloutFunc[T], Callable[..., A

        # Agent methods are also available
        result = my_llm_agent.rollout(task, resources, rollout)
+        ```

    Raises:
        NotImplementedError: If the function signature doesn't match any known patterns.
@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.

+"""Base abstractions for building agents that plug into Agent Lightning."""
+
 from __future__ import annotations

 import inspect
@@ -11,8 +13,8 @@ from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar
 from agentlightning.types import NamedResources, Rollout, RolloutRawResult, Task

 if TYPE_CHECKING:
-    from agentlightning.runner import BaseRunner
-    from agentlightning.tracer import BaseTracer
+    from agentlightning.runner import Runner
+    from agentlightning.tracer import Tracer
    from agentlightning.trainer import Trainer


@@ -26,32 +28,38 @@ __all__ = [


 def is_v0_1_rollout_api(func: Callable[..., Any]) -> bool:
-    """Check if the rollout API is v0.1.
-    Inspect the function signature to see if it has a rollout_id parameter.
+    """Return `True` when the rollout function uses the deprecated v0.1 signature.
+
+    The helper inspects the callable's signature to detect whether a `rollout_id`
+    parameter is present, which indicates the legacy API.

    Args:
-        func: The function to check.
+        func: Function to analyze.
+
+    Returns:
+        `True` if the callable exposes a `rollout_id` parameter.
    """
    return "rollout_id" in inspect.signature(func).parameters


 class LitAgent(Generic[T]):
-    """Base class for the training and validation logic of an agent.
+    """Base class for implementing agent rollouts.

-    Developers should subclass this class and implement the rollout methods
-    to define the agent's behavior for a single task. The agent's logic
-    is completely decoupled from the server communication and training
-    infrastructure.
+    Subclasses override the rollout methods to process tasks while the trainer and
+    runner infrastructure manages orchestration, tracing, and persistence.
    """

    def __init__(self, *, trained_agents: Optional[str] = None) -> None:  # FIXME: str | None won't work for cli
-        """
-        Initialize the LitAgent.
+        """Initialize the agent instance.

        Args:
-            trained_agents: Optional string representing the trained agents.
-                            This can be used to track which agents have been trained by this instance.
-                            Deprecated. Configure `agent_match` in adapter instead.
+            trained_agents: Optional identifier used by legacy tooling to mark trained
+                agents.
+
+        !!! warning "Deprecated"
+            The `trained_agents` flag is deprecated. Configure `agent_match` in the adapter
+            layer instead. See [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet]
+            for more details.
        """
        if trained_agents is not None:
            warnings.warn(
@@ -62,15 +70,12 @@ class LitAgent(Generic[T]):
        self.trained_agents = trained_agents

        self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
-        self._runner_ref: weakref.ReferenceType[BaseRunner[T]] | None = None
+        self._runner_ref: weakref.ReferenceType[Runner[T]] | None = None

    def is_async(self) -> bool:
-        """
-        Check if the agent implements asynchronous rollout methods.
-        Override this property for customized async detection logic.
+        """Return `True` when the agent overrides any asynchronous rollout methods.

-        Returns:
-            True if the agent has custom async rollout methods, False otherwise.
+        Override this method for customized async detection logic.
        """
        return (
            (
@@ -85,21 +90,15 @@ class LitAgent(Generic[T]):
        )

    def set_trainer(self, trainer: Trainer) -> None:
-        """
-        Set the trainer for this agent.
+        """Attach the trainer responsible for orchestration.

        Args:
-            trainer: The Trainer instance that will handle training and validation.
+            trainer: [`Trainer`][agentlightning.Trainer] that manages the agent.
        """
        self._trainer_ref = weakref.ref(trainer)

    def get_trainer(self) -> Trainer:
-        """
-        Get the trainer for this agent.
-
-        Returns:
-            The Trainer instance associated with this agent.
-        """
+        """Return the trainer associated with this agent."""
        if self._trainer_ref is None:
            raise ValueError("Trainer has not been set for this agent.")
        trainer = self._trainer_ref()
@@ -109,42 +108,31 @@ class LitAgent(Generic[T]):

    @property
    def trainer(self) -> Trainer:
-        """Convenient shortcut of self.get_trainer()."""
+        """Return the trainer associated with this agent."""
        return self.get_trainer()

-    def get_tracer(self) -> BaseTracer:
-        """
-        Get the tracer for this agent.
-
-        Returns:
-            The BaseTracer instance associated with this agent.
-        """
+    def get_tracer(self) -> Tracer:
+        """Return the tracer configured for this agent."""
        if hasattr(self.runner, "tracer"):
            return self.runner.tracer  # type: ignore
        else:
            return self.trainer.tracer

    @property
-    def tracer(self) -> BaseTracer:
-        """Convenient shortcut of self.get_tracer()."""
+    def tracer(self) -> Tracer:
+        """Return the tracer configured for this agent."""
        return self.get_tracer()

-    def set_runner(self, runner: BaseRunner[T]) -> None:
-        """
-        Set the runner for this agent.
+    def set_runner(self, runner: Runner[T]) -> None:
+        """Attach the runner responsible for executing rollouts.

        Args:
-            runner: The runner instance that will handle the execution of rollouts.
+            runner: [`Runner`][agentlightning.Runner] coordinating execution.
        """
        self._runner_ref = weakref.ref(runner)

-    def get_runner(self) -> BaseRunner[T]:
-        """
-        Get the runner for this agent.
-
-        Returns:
-            The runner instance associated with this agent.
-        """
+    def get_runner(self) -> Runner[T]:
+        """Return the runner responsible for executing rollouts."""
        if self._runner_ref is None:
            raise ValueError("Runner has not been set for this agent.")
        runner = self._runner_ref()
@@ -153,159 +141,111 @@ class LitAgent(Generic[T]):
        return runner

    @property
-    def runner(self) -> BaseRunner[T]:
-        """Convenient shortcut of self.get_runner()."""
+    def runner(self) -> Runner[T]:
+        """Return the runner responsible for executing rollouts."""
        return self.get_runner()

-    def on_rollout_start(self, task: Task, runner: BaseRunner[T], tracer: BaseTracer) -> None:
-        """Hook called immediately before a rollout begins.
+    def on_rollout_start(self, task: Task, runner: Runner[T], tracer: Tracer) -> None:
+        """Hook invoked immediately before a rollout begins.

-        Deprecated in favor of `on_rollout_start` in the `Hook` interface.
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource setup. The default implementation is a no-op.

        Args:
-            task: The :class:`Task` object that will be processed.
-            runner: The :class:`BaseRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
+            task: [`Task`][agentlightning.Task] that will be processed.
+            runner: [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.

-        Subclasses can override this method to implement custom logic such as
-        logging, metric collection, or resource setup. By default, this is a
-        no-op.
+        !!! warning "Deprecated"
+            Override [`Hook.on_rollout_start`][agentlightning.Hook.on_rollout_start]
+            instead of this method when extending agents.
        """

-    def on_rollout_end(self, task: Task, rollout: Rollout, runner: BaseRunner[T], tracer: BaseTracer) -> None:
-        """Hook called after a rollout completes.
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: Runner[T], tracer: Tracer) -> None:
+        """Hook invoked after a rollout completes.

-        Deprecated in favor of `on_rollout_end` in the `Hook` interface.
+        Subclasses can override this method for cleanup or additional logging. The default
+        implementation is a no-op.

        Args:
-            task: The :class:`Task` object that was processed.
-            rollout: The resulting :class:`Rollout` object.
-            runner: The :class:`BaseRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
+            task: [`Task`][agentlightning.Task] that was processed.
+            rollout: Resulting [`Rollout`][agentlightning.Rollout].
+            runner: [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.

-        Subclasses can override this method for cleanup or additional
-        logging. By default, this is a no-op.
+        !!! warning "Deprecated"
+            Override [`Hook.on_rollout_end`][agentlightning.Hook.on_rollout_end]
+            instead of this method when extending agents.
        """

    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Main entry point for executing a rollout.
+        """Execute a rollout synchronously.

-        This method determines whether to call the synchronous or
-        asynchronous rollout method based on the agent's implementation.

        If you don't wish to implement both training rollout and validation
        rollout separately, you can just implement `rollout` which will work for both.

        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-            rollout: The full rollout object, please avoid from directly modifying it.
-                     Most agents should only use `task` and `resources`. Use `rollout`
-                     only if you need to access metadata like `rollout_id`.
+            task: Task payload provided by the scheduler.
+            resources: Mapping of named resources (for example LLMs or prompt templates).
+            rollout: Rollout metadata. Avoid mutating this object directly unless a
+                subclass needs to override defaults.

        Returns:
-            The result of the rollout, which can be one of:
-            - None. The tracing should be handled by the agent runner.
-            - A float representing the final reward.
-            - A list of `Triplet` objects for detailed, step-by-step feedback.
-            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
-            - A list of dictionaries for any trace spans.
-            - A complete `Rollout` object for full control over reporting.
+            One of the following values:
+
+            * `None` when tracing is handled by the runner.
+            * `float` representing the final reward.
+            * `List[ReadableSpan]` with OpenTelemetry spans.
+            * `List[Span]` with Agent Lightning spans.
        """
        raise NotImplementedError("Agents must implement the `rollout` method.")

    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Asynchronous version of the main rollout method.
-
-        This method determines whether to call the synchronous or
-        asynchronous rollout method based on the agent's implementation.
+        """Execute a rollout asynchronously.

        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-            rollout: The full rollout object, please avoid from directly modifying it.
-                     Most agents should only use `task` and `resources`. Use `rollout`
-                     only if you need to access metadata like `rollout_id`.
+            task: Task payload provided by the scheduler.
+            resources: Mapping of named resources (for example LLMs or prompt templates).
+            rollout: Rollout metadata. Avoid mutating this object directly unless a
+                subclass needs to override defaults.

        Returns:
-            The result of the rollout, which can be one of:
-            - None. The tracing should be handled by the agent runner.
-            - A float representing the final reward.
-            - A list of `Triplet` objects for detailed, step-by-step feedback.
-            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
-            - A list of dictionaries for any trace spans.
-            - A complete `Rollout` object for full control over reporting.
+            Same possible return values as
+            [`rollout`][agentlightning.LitAgent.rollout].
        """
        raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")

    def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Defines the agent's behavior for a single training task.
+        """Process a single training task synchronously.

-        This method should contain the logic for how the agent processes an
-        input, uses the provided resources (like LLMs or prompts), and
-        produces a result.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-            rollout: The full rollout object, please avoid from directly modifying it.
+        By default, this method delegates to
+        [`rollout`][agentlightning.LitAgent.rollout].
        """
        return self.rollout(task, resources, rollout)

    def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Defines the agent's behavior for a single validation task.
+        """Process a single validation task synchronously.

-        By default, this method redirects to `training_rollout`. Override it
-        if the agent should behave differently during validation.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            resources: A dictionary of named resources for the agent to use.
-            rollout: The full rollout object, avoid from modifying it.
-
-        Returns:
-            The result of the validation rollout. See `rollout` for
-            possible return types.
+        Override this method when validation should differ from training. The default
+        implementation delegates to
+        [`training_rollout`][agentlightning.LitAgent.training_rollout].
        """
        return self.rollout(task, resources, rollout)

    async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Asynchronous version of `training_rollout`.
+        """Process a single training task asynchronously.

-        This method should be implemented by agents that perform asynchronous
-        operations (e.g., non-blocking I/O, concurrent API calls).
-
-        Args:
-            task: The task object received from the server.
-            resources: A dictionary of named resources for the agent to use.
-            rollout: The full rollout object, avoid from modifying it.
-
-        Returns:
-            The result of the asynchronous training rollout. See `rollout` for
-            possible return types.
+        By default, this method delegates to
+        [`rollout_async`][agentlightning.LitAgent.rollout_async].
        """
        return await self.rollout_async(task, resources, rollout)

    async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
-        """Asynchronous version of `validation_rollout`.
+        """Process a single validation task asynchronously.

-        By default, this method redirects to `training_rollout_async`.
-        Override it for different asynchronous validation behavior.
-
-        Args:
-            task: The task object received from the server.
-            resources: A dictionary of named resources for the agent to use.
-            rollout: The full rollout object, avoid from modifying it.
-
-        Returns:
-            The result of the asynchronous validation rollout. See `rollout` for
-            possible return types.
+        Override this method when validation should differ from training. The default
+        implementation delegates to
+        [`training_rollout_async`][agentlightning.LitAgent.training_rollout_async].
        """
        return await self.rollout_async(task, resources, rollout)
@@ -23,6 +23,7 @@ from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfi
 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

@@ -38,7 +39,7 @@ __all__ = [
 class ModelConfig(TypedDict):
    """LiteLLM model registration entry.

-    This mirrors the items in LiteLLM's ``model_list`` section.
+    This mirrors the items in LiteLLM's `model_list` section.

    Attributes:
        model_name: Logical model name exposed by the proxy.
@@ -53,8 +54,8 @@ class ModelConfig(TypedDict):
 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.
+    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.
@@ -77,98 +78,53 @@ def _get_pre_call_data(args: Any, kwargs: Any) -> Dict[str, Any]:
    return cast(Dict[str, Any], data)


-# We need global state because litellm is based on a global app.
-# Repeatedly initializing the app with different stores will cause errors.
-_initialized: bool = False
-_global_store: LightningStore | None = None
+def _reset_litellm_logging_worker() -> None:
+    """Reset LiteLLM's global logging worker to the current event loop.

-
-def get_global_store() -> LightningStore:
-    """Return the globally registered LightningStore.
-
-    Used by components that are initialized without an explicit store
-    (e.g., exporter created inside OpenTelemetry).
-
-    Returns:
-        LightningStore: The active global store.
-
-    Raises:
-        ValueError: If the global store has not been set by ``LLMProxy.start()``.
+    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.
    """
-    if _global_store is None:
-        raise ValueError("Global store is not initialized. Please start a LLMProxy first.")
-    return _global_store
+
+    # ``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 initialize() -> None:
-    """Initialize global middleware and LiteLLM callbacks once.
+def _reset_litellm_logging_callback_manager() -> None:
+    """Reset LiteLLM's global callback manager.

-    Idempotent. Installs:
+    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.

-    * 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.
-
-    This function does not start any server. It only wires global hooks.
+    It does not respect existing input/output callbacks.
    """
-    global _initialized
-    if _initialized:
-        return

-    # Add middleware here because it relies on the global store.
-    @app.middleware("http")
-    async def rollout_attempt_middleware(  # pyright: ignore[reportUnusedFunction]
-        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()
-
-            # Allocate a monotonic sequence id per (rollout, attempt).
-            sequence_id = await get_global_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()),
-            ]
-
-        response = await call_next(request)
-        return response
-
-    # Register callbacks once on the global LiteLLM callback list.
-    litellm.callbacks.extend(  # pyright: ignore[reportUnknownMemberType]
-        [
-            AddReturnTokenIds(),
-            LightningOpenTelemetry(),
-        ]
-    )
-
-    _initialized = True
+    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``
+    This mutates the outgoing request payload to include `return_token_ids=True`
    for backends that support token id return (e.g., vLLM).

-    See:
-        https://github.com/vllm-project/vllm/pull/22587
+    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]]]:
@@ -205,15 +161,13 @@ class LightningSpanExporter(SpanExporter):
    * 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.
-
-    Args:
-        store: Optional explicit LightningStore. If None, uses ``get_global_store()``.
    """

-    def __init__(self, store: Optional[LightningStore] = None):
-        self._store = store
+    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.
@@ -227,6 +181,7 @@ class LightningSpanExporter(SpanExporter):
        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)
@@ -239,22 +194,25 @@ class LightningSpanExporter(SpanExporter):
        Returns:
            threading.RLock: The initialized lock.
        """
+        self._clear_loop_and_lock()
        if self._lock is None:
            self._lock = threading.RLock()
        return self._lock

-    def _get_store(self) -> LightningStore:
-        """Return the LightningStore to use.
+    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.

-        Returns:
-            LightningStore: Explicit store if provided, else the global store.
-
-        Raises:
-            ValueError: If no global store is configured and no explicit store was given.
+        This should only happen in CI.
        """
-        if self._store is None:
-            return get_global_store()
-        return self._store
+        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."""
@@ -324,10 +282,10 @@ class LightningSpanExporter(SpanExporter):
            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.
+            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)
+            `x-rollout-id` (str), `x-attempt-id` (str), `x-sequence-id` (str of int)

        Raises:
            None directly. Logs and skips malformed spans.
@@ -339,6 +297,11 @@ class LightningSpanExporter(SpanExporter):
            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] = {}

@@ -353,6 +316,9 @@ class LightningSpanExporter(SpanExporter):
                        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)
@@ -390,14 +356,14 @@ class LightningSpanExporter(SpanExporter):

            # Persist each span in the subtree with the resolved identifiers.
            for span in subtree_spans:
-                await self._get_store().add_otel_span(
+                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``.
+        A root span is defined as one with `parent is None`.

        Yields:
            int: Span id for each root span found.
@@ -409,7 +375,7 @@ class LightningSpanExporter(SpanExporter):
                    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``.
+        """Yield span_ids in the subtree rooted at `root_span_id`.

        Depth-first traversal over the current buffer.

@@ -459,25 +425,63 @@ class LightningOpenTelemetry(OpenTelemetry):

    * Ensures each request is annotated with a per-attempt sequence id so spans
      are ordered deterministically even with clock skew across nodes.
-    * Uses ``LightningSpanExporter`` to persist spans for analytics and training.
-
-    Args:
-        store: Optional explicit LightningStore for the exporter.
+    * Uses [`LightningSpanExporter`][agentlightning.llm_proxy.LightningSpanExporter] to persist spans for analytics and training.
    """

-    def __init__(self, store: LightningStore | None = None):
-        config = OpenTelemetryConfig(exporter=LightningSpanExporter(store))
+    def __init__(self):
+        config = OpenTelemetryConfig(exporter=LightningSpanExporter())

        # Check for tracer initialization
-        if (
-            hasattr(trace_api, "_TRACER_PROVIDER")
-            and trace_api._TRACER_PROVIDER is not None  # pyright: ignore[reportPrivateUsage]
-        ):
+        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.

@@ -486,25 +490,34 @@ class LLMProxy:
    * 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.
+    * Writes a LiteLLM worker config file with `model_list` and settings.

    Lifecycle:

-    * ``start()`` writes config, starts uvicorn server in a thread, and waits until ready.
-    * ``stop()`` tears down the server and removes the temp config file.
-    * ``restart()`` convenience wrapper to stop then start.
+    * [`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.

    Usage 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).

+    !!! 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.
-        model_list: LiteLLM ``model_list`` entries.
+        model_list: LiteLLM `model_list` entries.
        store: LightningStore used for span sequence and persistence.
        host: Publicly reachable host used in resource endpoints. Defaults to best-guess IPv4.
-        litellm_config: Extra LiteLLM proxy config merged with ``model_list``.
-        num_retries: Default LiteLLM retry count injected into ``litellm_settings``.
+        litellm_config: Extra LiteLLM proxy config merged with `model_list`.
+        num_retries: Default LiteLLM retry count injected into `litellm_settings`.
    """

    def __init__(
@@ -515,6 +528,7 @@ class LLMProxy:
        host: str | None = None,
        litellm_config: Dict[str, Any] | None = None,
        num_retries: int = 0,
+        _add_return_token_ids: bool = True,
    ):
        self.store = store
        self.host = host or _get_default_ipv4_address()
@@ -531,6 +545,16 @@ class LLMProxy:
        self._uvicorn_server = None
        self._ready_event = threading.Event()

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

@@ -578,13 +602,56 @@ class LLMProxy:
                break
            time.sleep(0.01)

+    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()
+
    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.
+        * 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.
        """
@@ -595,12 +662,8 @@ class LLMProxy:
        if not self.store:
            raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")

-        global _global_store
-
-        _global_store = self.store
-
-        # Initialize global middleware and callbacks once.
-        initialize()
+        # 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
@@ -670,7 +733,7 @@ class LLMProxy:
    def restart(self, *, _port: int | None = None) -> None:
        """Restart the proxy if running, else start it.

-        Convenience wrapper calling ``stop()`` followed by ``start()``.
+        Convenience wrapper calling `stop()` followed by `start()`.
        """
        logger.info("Restarting LLMProxy server...")
        if self.is_running():
@@ -694,31 +757,37 @@ class LLMProxy:
        model: str | None = None,
        sampling_parameters: Dict[str, Any] | None = None,
    ) -> LLM:
-        """Create an ``LLM`` resource pointing at this proxy with rollout context.
+        """Create an `LLM` resource pointing at this proxy with rollout context.

        The returned endpoint is:
-            ``http://{host}:{port}/rollout/{rollout_id}/attempt/{attempt_id}``
+            `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, that model is used.
+                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.
+            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:
-                raise ValueError(
-                    f"Multiple or zero models found in model_list: {self.model_list}. Please specify the model."
-                )
+                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(
@@ -736,6 +805,79 @@ class LLMProxy:
            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 _get_default_ipv4_address() -> str:
    """Determine the default outbound IPv4 address for this machine.

@@ -744,7 +886,7 @@ def _get_default_ipv4_address() -> str:
        selection, then inspects the socket's local address. No packets are sent.

    Returns:
-        str: Best-guess IPv4 like ``192.168.x.y``. Falls back to ``127.0.0.1``.
+        str: Best-guess IPv4 like `192.168.x.y`. Falls back to `127.0.0.1`.
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
@@ -763,3 +905,19 @@ def _check_port(host: str, port: int) -> bool:
        s.settimeout(1)
        result = s.connect_ex((host, port))
        return result != 0  # True if unavailable
+
+
+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,11 +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,11 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.

 from .agent import LitAgentRunner
-from .base import BaseRunner
+from .base import Runner
 from .legacy import LegacyAgentRunner

 __all__ = [
-    "BaseRunner",
+    "Runner",
    "LegacyAgentRunner",
    "LitAgentRunner",
 ]
@@ -20,7 +20,7 @@ from agentlightning.litagent import LitAgent
 from agentlightning.reward import emit_reward, find_final_reward
 from agentlightning.store.base import LightningStore
 from agentlightning.tracer.agentops import AgentOpsTracer
-from agentlightning.tracer.base import BaseTracer
+from agentlightning.tracer.base import Tracer
 from agentlightning.types import (
    AttemptedRollout,
    Hook,
@@ -34,33 +34,32 @@ from agentlightning.types import (
 if TYPE_CHECKING:
    from agentlightning.execution.events import ExecutionEvent

-from .base import BaseRunner
+from .base import Runner

 T_task = TypeVar("T_task")

 logger = logging.getLogger(__name__)


-class LitAgentRunner(BaseRunner[T_task]):
-    """Runner implementation for executing agent tasks with distributed support.
+class LitAgentRunner(Runner[T_task]):
+    """Execute [`LitAgent`][agentlightning.LitAgent] tasks with tracing support.

    This runner manages the complete lifecycle of agent rollout execution,
    including task polling, resource management, tracing, and hooks. It supports
    both continuous iteration over tasks from the store and single-step execution.

    Attributes:
-        worker_id: The unique identifier for this worker process.
+        worker_id: Identifier for the active worker process, if any.
    """

-    def __init__(self, tracer: BaseTracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0) -> None:
+    def __init__(self, tracer: Tracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0) -> None:
        """Initialize the agent runner.

        Args:
-            tracer: The tracer instance for recording execution traces and spans.
-            max_rollouts: Maximum number of tasks to process in iter() mode. If None,
-                the runner will continue indefinitely until interrupted.
-            poll_interval: Time in seconds to wait between polling attempts when
-                no tasks are available in the store.
+            tracer: [`Tracer`][agentlightning.Tracer] used for rollout spans.
+            max_rollouts: Optional cap on iterations processed by
+                [`iter`][agentlightning.LitAgentRunner.iter].
+            poll_interval: Seconds to wait between store polls when no work is available.
        """
        super().__init__()
        self._tracer = tracer
@@ -80,10 +79,9 @@ class LitAgentRunner(BaseRunner[T_task]):
        initializes the tracer.

        Args:
-            agent: The LitAgent instance to be managed by this runner.
-            hooks: Optional sequence of Hook objects to be called at various
-                lifecycle stages (on_trace_start, on_trace_end, on_rollout_start,
-                on_rollout_end).
+            agent: [`LitAgent`][agentlightning.LitAgent] instance executed by the runner.
+            hooks: Optional sequence of [`Hook`][agentlightning.Hook]
+                callbacks invoked around tracing and rollout boundaries.
            **kwargs: Additional initialization arguments (currently unused).
        """
        self._agent = agent
@@ -100,7 +98,8 @@ class LitAgentRunner(BaseRunner[T_task]):

        Args:
            worker_id: Unique identifier for this worker process.
-            store: The LightningStore instance for task coordination and data persistence.
+            store: [`LightningStore`][agentlightning.LightningStore]
+                used for task coordination and persistence.
            **kwargs: Additional worker-specific initialization arguments (currently unused).
        """
        self._store = store
@@ -131,7 +130,7 @@ class LitAgentRunner(BaseRunner[T_task]):
        This method cleans up worker-specific resources and resets the worker ID.

        Args:
-            worker_id: The unique identifier of the worker being torn down.
+            worker_id: Unique identifier of the worker being torn down.
            *args: Additional teardown arguments (currently unused).
            **kwargs: Additional teardown keyword arguments (currently unused).
        """
@@ -140,11 +139,11 @@ class LitAgentRunner(BaseRunner[T_task]):
        self._tracer.teardown_worker(worker_id)

    @property
-    def tracer(self) -> BaseTracer:
+    def tracer(self) -> Tracer:
        """Get the tracer instance.

        Returns:
-            The BaseTracer instance used by this runner.
+            The Tracer instance used by this runner.
        """
        return self._tracer

@@ -155,7 +154,7 @@ class LitAgentRunner(BaseRunner[T_task]):
            The LitAgent instance managed by this runner.

        Raises:
-            ValueError: If the agent has not been initialized via init().
+            ValueError: If the agent has not been initialized via [`init`][agentlightning.LitAgentRunner.init].
        """
        if self._agent is None:
            raise ValueError("Agent not initialized. Call init() first.")
@@ -168,7 +167,7 @@ class LitAgentRunner(BaseRunner[T_task]):
            The LightningStore instance for this worker.

        Raises:
-            ValueError: If the store has not been initialized via init_worker().
+            ValueError: If the store has not been initialized via [`init_worker`][agentlightning.LitAgentRunner.init_worker].
        """
        if self._store is None:
            raise ValueError("Store not initialized. Call init_worker() first.")
@@ -312,7 +311,7 @@ class LitAgentRunner(BaseRunner[T_task]):
        and return early if the event is set.

        Args:
-            event: Optional ExecutionEvent object that can be used to interrupt the sleep.
+            event: Optional [`ExecutionEvent`][agentlightning.ExecutionEvent] object that can be used to interrupt the sleep.
                If set during the sleep period, the method returns immediately.
        """
        if event is None:
@@ -364,7 +363,7 @@ class LitAgentRunner(BaseRunner[T_task]):
            await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)

            start_time = time.time()
-            with self._tracer.trace_context(
+            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(
@@ -435,6 +434,7 @@ class LitAgentRunner(BaseRunner[T_task]):
        """Run the runner, continuously iterating over tasks in the store.

        This method polls the store for new rollouts and executes them until:
+
        - The event is set (if provided)
        - The max_rollouts limit is reached (if configured)
        - No more tasks are available
@@ -497,7 +497,8 @@ class LitAgentRunner(BaseRunner[T_task]):
        """Execute a single task directly, bypassing the task queue.

        This method creates a new rollout for the given input and executes it
-        immediately. Unlike iter(), exceptions are propagated to the caller.
+        immediately. Unlike [`iter()`][agentlightning.LitAgentRunner.iter],
+        exceptions are propagated to the caller.

        Args:
            input: The task input to be processed by the agent.
@@ -1,11 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.

-"""Base runner interface for executing agent tasks.
-
-This module defines the abstract base class for all runner implementations
-in the agent-lightning framework. Runners are responsible for managing the
-execution lifecycle of agents and coordinating with the store.
-"""
+"""Abstract runner interface for executing agent tasks."""

 from __future__ import annotations

@@ -27,90 +22,75 @@ T_task = TypeVar("T_task")
 logger = logging.getLogger(__name__)


-class BaseRunner(ParallelWorkerBase, Generic[T_task]):
-    """Base class for all runners.
+class Runner(ParallelWorkerBase, Generic[T_task]):
+    """Abstract base class for long-running agent executors.

-    This abstract base class defines the interface that all runner implementations
-    must follow. Runners are responsible for executing agent tasks, managing the
-    execution lifecycle, and coordinating with the store.
+    Runner implementations coordinate [`LitAgent`][agentlightning.LitAgent]
+    instances, acquire work from a [`LightningStore`][agentlightning.LightningStore],
+    and emit [`Rollout`][agentlightning.Rollout] objects. Subclasses decide how
+    to schedule work (polling, streaming, etc.) while this base class provides a
+    minimal lifecycle contract.
    """

    def init(self, agent: LitAgent[T_task], **kwargs: Any) -> None:
-        """Initialize the runner with the agent.
+        """Prepare the runner to execute tasks for `agent`.

-        This method is called once during setup to configure the runner with
-        the agent it will execute.
+        This method is called only once during the setup for all workers, not for each worker.

        Args:
-            agent: The LitAgent instance to be managed by this runner.
-            **kwargs: Additional initialization arguments specific to the runner implementation.
+            agent: Agent instance providing task-specific logic.
+            **kwargs: Optional runner-specific configuration.

        Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must supply the initialization
+                routine.
        """
        raise NotImplementedError()

    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
-        """Initialize the runner for each worker with worker_id and store.
+        """Configure worker-local state before processing tasks.

-        This method is called once per worker process in a distributed setup.
-        It provides the worker with its unique ID and the store instance for
-        task coordination.
+        This method is called for **each** worker during the setup.

        Args:
-            worker_id: Unique identifier for this worker process.
-            store: The LightningStore instance for task coordination and data persistence.
-            **kwargs: Additional worker-specific initialization arguments.
+            worker_id: Unique identifier for this worker process or thread.
+            store: Shared [`LightningStore`][agentlightning.LightningStore]
+                backing task coordination.
+            **kwargs: Optional worker-specific configuration.

        Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must prepare per-worker resources.
        """
        raise NotImplementedError()

    def run(self, *args: Any, **kwargs: Any) -> None:
-        """Undefined method - use iter() or step() instead.
+        """Deprecated synchronous entry point.

-        This method is intentionally not implemented as the execution behavior
-        should be defined through iter() for continuous execution or step()
-        for single-task execution.
-
-        Args:
-            *args: Unused positional arguments.
-            **kwargs: Unused keyword arguments.
+        Use [`iter()`][agentlightning.Runner.iter] or [`step()`][agentlightning.Runner.step] instead.

        Raises:
-            RuntimeError: Always raised to indicate this method should not be used.
+            RuntimeError: Always raised to direct callers to
+                [iter()][agentlightning.Runner.iter] or
+                [step()][agentlightning.Runner.step].
        """
        raise RuntimeError("The behavior of run() of Runner is undefined. Use iter() or step() instead.")

    def teardown(self, *args: Any, **kwargs: Any) -> None:
-        """Clean up runner resources and reset state.
-
-        This method is called once during shutdown to clean up any resources
-        allocated during initialization and reset the runner state.
-
-        Args:
-            *args: Additional teardown arguments.
-            **kwargs: Additional teardown keyword arguments.
+        """Release resources acquired during [`init()`][agentlightning.Runner.init].

        Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must implement the shutdown routine.
        """
        raise NotImplementedError()

    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
-        """Clean up worker-specific resources.
-
-        This method is called once per worker during shutdown to clean up
-        any resources specific to that worker.
+        """Release per-worker resources allocated by [`init_worker()`][agentlightning.Runner.init_worker].

        Args:
-            worker_id: The unique identifier of the worker being torn down.
-            *args: Additional teardown arguments.
-            **kwargs: Additional teardown keyword arguments.
+            worker_id: Identifier of the worker being torn down.

        Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses must implement the shutdown routine.
        """
        raise NotImplementedError()

@@ -122,18 +102,21 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
        store: LightningStore,
        hooks: Optional[Sequence[Hook]] = None,
        worker_id: Optional[int] = None,
-    ) -> Iterator[BaseRunner[T_task]]:
-        """Context manager for quickly init and teardown the runner,
-        so that you can debug the runner without a trainer environment.
+    ) -> Iterator[Runner[T_task]]:
+        """Initialize and tear down a runner within a simple context manager.
+
+        The helper is primarily intended for debugging runner implementations
+        outside of a full [`Trainer`][agentlightning.Trainer] stack.

        Args:
-            agent: The LitAgent instance to be managed by this runner.
-                   It should be the same agent that is to be run within the context.
-            store: The LightningStore instance for task coordination and data persistence.
-                   If you don't have one, you can easily create one with `InMemoryLightningStore()`.
-            hooks: Optional sequence of Hook instances to be used by the runner.
-                   Only some runners support hooks.
-            worker_id: Optional worker ID to be used by the runner.
+            agent: Agent executed by this runner.
+            store: Backing [`LightningStore`][agentlightning.LightningStore].
+                If you don't have one, you can easily create one with
+                [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
+            hooks: Optional sequence of hooks recognised by the runner.
+                Not all runners support hooks.
+            worker_id: Override the worker identifier used during setup. Defaults
+                to `0`.
        """
        _initialized: bool = False
        _worker_initialized: bool = False
@@ -163,12 +146,11 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
        them until interrupted by the event or when no more tasks are available.

        Args:
-            event: Optional ExecutionEvent object that can be used to signal the runner
-                to stop gracefully. When set, the runner should finish its current
-                task and exit the iteration loop.
+            event: Cooperative stop signal. When set, the runner should complete
+                the current unit of work and exit the loop.

        Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses provide the iteration behavior.
        """
        raise NotImplementedError()

@@ -186,18 +168,15 @@ class BaseRunner(ParallelWorkerBase, Generic[T_task]):
        directly, bypassing the store's task queue.

        Args:
-            input: The task input to be processed by the agent.
-            resources: Optional named resources to be used for this specific task.
-                If not provided, the latest resources from the store will be used.
-            mode: Optional rollout mode (e.g., "train", "test"). If not provided,
-                the default mode will be used.
-            event: Optional ExecutionEvent object to signal interruption. When set, the
-                runner may abort the current execution.
+            input: Task payload consumed by the agent.
+            resources: Optional named resources scoped to this invocation.
+            mode: Optional rollout mode such as `"train"` or `"eval"`.
+            event: Cooperative stop signal for long-running tasks.

        Returns:
-            The completed rollout.
+            Completed rollout produced by the agent.

        Raises:
-            NotImplementedError: Must be implemented by subclasses.
+            NotImplementedError: Subclasses provide the execution behavior.
        """
        raise NotImplementedError()
@@ -11,10 +11,10 @@ from agentlightning.adapter import TracerTraceToTriplet
 from agentlightning.client import AgentLightningClient
 from agentlightning.litagent import LitAgent
 from agentlightning.litagent.litagent import is_v0_1_rollout_api
-from agentlightning.tracer.base import BaseTracer
+from agentlightning.tracer.base import Tracer
 from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet

-from .base import BaseRunner
+from .base import Runner

 logger = logging.getLogger(__name__)

@@ -23,7 +23,7 @@ __all__ = [
 ]


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

    This class orchestrates the interaction between the agent (`LitAgent`) and
@@ -43,7 +43,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
        self,
        agent: LitAgent[Any],
        client: AgentLightningClient,
-        tracer: BaseTracer,
+        tracer: Tracer,
        triplet_exporter: TracerTraceToTriplet,
        worker_id: Optional[int] = None,
        max_tasks: Optional[int] = None,
@@ -58,7 +58,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
        self.worker_id = worker_id
        self.max_tasks = max_tasks

-    # These methods are overridden by BaseRunner, getting them back to old behavior.
+    # These methods are overridden by Runner, getting them back to old behavior.
    def init(self, *args: Any, **kwargs: Any) -> None:
        pass

@@ -180,7 +180,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
            except Exception:
                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")

-            with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
+            with self.tracer._trace_context_sync(name=f"rollout_{rollout_id}"):  # pyright: ignore[reportPrivateUsage]
                start_time = time.time()
                rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
                # Pass the task input, not the whole task object
@@ -257,7 +257,7 @@ class LegacyAgentRunner(BaseRunner[Any]):
            except Exception:
                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")

-            with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
+            async with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
                start_time = time.time()
                rollout_method = (
                    self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async
@@ -1,6 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.

-"""Legacy server for the Agent Lightning framework. Deprecated in favor of agentlightning.store."""
+"""Legacy HTTP server compatible with the original Agent Lightning protocol.
+
+The implementation in this module predates the modern store-powered runtime and
+is kept for backwards compatibility with older deployments. New applications
+should migrate to the store architecture where possible.
+"""

 from __future__ import annotations

@@ -29,9 +34,15 @@ logger = logging.getLogger(__name__)


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

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

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

    async def update_resources(self, update: ResourcesUpdate):
-        """
-        Safely stores a new version of named resources and sets it as the latest.
+        """Persist a new resource bundle and mark it as the latest version.
+
+        Args:
+            update: Resource payload received from a client.
        """
        # TODO: evict old resources if necessary.
        async with self._resources_lock:
@@ -105,8 +130,14 @@ 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)
@@ -115,16 +146,16 @@ class ServerDataStore:
            return None

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

    async def store_rollout(self, rollout: RolloutLegacy):
-        """
-        Safely stores a completed rollout from a client.
+        """Persist a completed rollout for later inspection.
+
+        Args:
+            rollout: Rollout returned by a client.
        """
        async with self._results_lock:
            self._processing_tasks.pop(rollout.rollout_id, None)
@@ -132,27 +163,31 @@ class ServerDataStore:
            logger.info(f"Rollout received and stored: {rollout.rollout_id}")

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

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

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

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


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

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

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

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

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

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

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

            if not self._store:
@@ -253,7 +289,7 @@ class AgentLightningServer:

        @self._app.get("/resources/latest", response_model=ResourcesUpdate)
        async def fetch_latest_resources() -> ResourcesUpdate:  # type: ignore
-            """Endpoint for clients to poll for the latest available resources."""
+            """Return the most recent resource bundle published to the server."""
            if not self._store:
                raise HTTPException(status_code=503, detail="Server not fully initialized.")
            resources_update = await self._store.get_latest_resources()
@@ -266,7 +302,7 @@ class AgentLightningServer:
        async def fetch_resources_by_id(  # type: ignore
            resource_id: str = Path(..., description="The unique identifier for the resource version.")
        ) -> ResourcesUpdate:
-            """Endpoint for clients to fetch a specific version of resources."""
+            """Return a specific version of resources by identifier."""
            if not self._store:
                raise HTTPException(status_code=503, detail="Server not fully initialized.")
            resources_update = await self._store.get_resources_by_id(resource_id)
@@ -277,7 +313,7 @@ class AgentLightningServer:

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

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

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

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

    async def queue_task(
@@ -314,17 +347,13 @@ 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()}"
@@ -333,16 +362,20 @@ class AgentLightningServer:
        return resources_id

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

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

    async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
-        """
-        Retrieves all available completed trajectories and clears the internal store.
-        """
+        """Return every completed rollout and clear the internal buffer."""
        if not self._store:
            raise RuntimeError("Store not initialized. The server may not be running.")
        return await self._store.retrieve_completed_rollouts()
@@ -53,12 +53,25 @@ Unset = _UnsetType  # Alias for convenience


 class LightningStore:
-    """
-    A centralized, thread-safe, async, data store for the lightning's state.
-    This holds the task queue, versioned resources, and completed rollouts.
+    """Contract for the persistent control-plane that coordinates training rollouts.

-    The store has a built-in clock and it should be responsible for tracking the times.
-    All the time-based operations like retry, timeout, etc. should be handled by the store.
+    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`.
    """

    async def start_rollout(
@@ -66,19 +79,38 @@ class LightningStore:
        input: TaskInput,
        mode: Literal["train", "val", "test"] | None = None,
        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> AttemptedRollout:
-        """
-        Add one incomplete rollout to the store, and get an attempt created for it.
-        This will immediately sets the rollout to a preparing state, and should be
-        used by whoever is going to execute the rollout.
+        """Register a rollout and immediately create its first attempt.

-        Return a special rollout with attempt object. Do not update it directly.
+        !!! note
+            Use [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout] when the
+            caller only wants to submit work for later scheduling.

-        But if the rollout fails or timeouts, it's still possible that the watchdog
-        sends it back to the queue for retry.
+        The rollout must be persisted with `status="preparing"` and an initial attempt
+        with `sequence_id == 1` so the caller can begin execution without visiting the
+        public queue. Implementations are expected to:

-        To enqueue a rollout to the task queue, use `enqueue_rollout` instead.
+        1. Generate a unique `rollout_id` and `attempt_id`.
+        2. Record `start_time` for both rollout and attempt based on the current clock.
+        3. Copy `config` and `metadata` so later mutations do not leak shared references.
+        4. Resolve `resources_id` to the latest resource snapshot when `None` is supplied.
+
+        Args:
+            input: Arbitrary task payload supplied by an algorithm.
+            mode: Optional semantic mode for downstream analytics (`"train"`, `"val"`, `"test"`).
+            resources_id: Concrete resource snapshot to execute against; defaults to the latest stored snapshot.
+            config: Rollout retry/timeout policy. Should default to a fresh [`RolloutConfig`][agentlightning.RolloutConfig].
+            metadata: Free-form metadata persisted verbatim with the rollout.
+
+        Returns:
+            The fully-populated [`AttemptedRollout`][agentlightning.AttemptedRollout] including
+            the just-created attempt.
+
+        Raises:
+            NotImplementedError: Subclasses must provide durable storage for the rollout.
+            ValueError: Implementations should raise when `resources_id` does not exist.
        """
        raise NotImplementedError()

@@ -87,34 +119,99 @@ class LightningStore:
        input: TaskInput,
        mode: Literal["train", "val", "test"] | None = None,
        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> Rollout:
-        """
-        Adds a new task to the queue with specific metadata and
-        returns the rollout object with its unique ID.
+        """Persist a rollout in `queuing` state so runners can claim it later.
+
+        !!! note
+            Different from [`start_rollout()`][agentlightning.LightningStore.start_rollout],
+            this method is called when the caller only wants to submit work for later scheduling.
+
+        Implementations must generate a unique `rollout_id`, stamp `start_time` with
+        the current time, default `config` to a fresh [`RolloutConfig`][agentlightning.RolloutConfig],
+        and insert the rollout at the tail of the scheduling queue. No attempt is created yet.
+
+        Args:
+            input: Arbitrary task payload supplied by an algorithm.
+            mode: Optional semantic mode indicator (`"train"`, `"val"`, `"test"`).
+            resources_id: Resource snapshot used when a runner eventually executes the rollout.
+            config: Fine-grained retry/timeout parameters to persist with the rollout.
+            metadata: Free-form metadata stored verbatim with the rollout record.
+
+        Returns:
+            The stored [`Rollout`][agentlightning.Rollout] in `queuing` status.
+
+        Raises:
+            NotImplementedError: Subclasses must persist the rollout.
+            ValueError: Implementations should raise when `resources_id` does not exist.
        """
        raise NotImplementedError()

    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        """
-        Retrieves the next task from the queue without blocking.
-        Returns None if the queue is empty.
+        """Claim the oldest queued rollout and transition it to `preparing`.

-        Will set the rollout status 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 new attempt for a given rollout ID and return the attempt details.
+        """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:
-        """
-        Add a span to the store.
+        """Persist a pre-constructed span emitted during rollout execution.

-        This method is responsible for updating the rollout/attempt status to "running" if needed.
+        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()

@@ -125,88 +222,221 @@ class LightningStore:
        readable_span: ReadableSpan,
        sequence_id: int | None = None,
    ) -> Span:
-        """
-        Add an opentelemetry span to the store.
+        """Convert and persist an OpenTelemetry span for a particular attempt.

-        If sequence_id is not provided, it will be fetched from `get_next_span_sequence_id` and assigned automatically.
+        Implementations must transform the `readable_span` into a [`Span`][agentlightning.Span]
+        (typically via [`Span.from_opentelemetry()`][agentlightning.Span.from_opentelemetry]),
+        assign a strictly increasing `sequence_id` when one is not provided, and persist it
+        using the same semantics as [`add_span()`][agentlightning.LightningStore.add_span].
+
+        Args:
+            rollout_id: Identifier of the rollout that produced the span.
+            attempt_id: Attempt identifier the span belongs to.
+            readable_span: OpenTelemetry span in SDK form.
+            sequence_id: Optional explicit ordering hint. When omitted, call
+                [`get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id]
+                automatically.
+
+        Returns:
+            The stored span record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement span persistence.
+            ValueError: Implementations must raise when the rollout or attempt is unknown.
        """
        raise NotImplementedError()

    async def query_rollouts(
        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
    ) -> List[Rollout]:
-        """
-        Query and retrieve rollouts filtered by their status.
-        If no status is provided, returns all rollouts.
+        """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]:
-        """
-        Query and retrieve all attempts associated with a specific rollout ID.
-        Returns an empty list if no attempts are found.
+        """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[Rollout]:
-        """
-        Safely retrieves a specific rollout by its ID.
+        """Fetch a rollout by identifier without mutating its state.
+
+        Args:
+            rollout_id: Identifier to retrieve.
+
+        Returns:
+            The rollout when found, otherwise `None`.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
        """
        raise NotImplementedError()

    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
-        """
-        Safely retrieves the latest attempt for a given rollout ID.
+        """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 get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves a specific version of named resources by its ID.
+        """Return a specific named resource snapshot by identifier.
+
+        Args:
+            resources_id: Identifier of the snapshot.
+
+        Returns:
+            The stored [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or `None` when missing.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
        """
        raise NotImplementedError()

    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves the latest version of named resources.
+        """Fetch the latest resource snapshot marked as the global default.
+
+        Returns:
+            The current latest [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or `None` when
+            no resources have been registered yet.
+
+        Raises:
+            NotImplementedError: Subclasses must implement retrieval.
        """
        raise NotImplementedError()

    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
-        """
-        Get the next span sequence ID for a given rollout and attempt.
-        This should be used to assign a unique sequence ID to each span within an attempt.
+        """Allocate the next strictly increasing sequence number used to order spans.

-        Recommend getting the ID before the operation even begins to avoid racing conditions.
+        Implementations must retain counters so repeated calls return `1, 2, ...` without
+        gaps unless spans were explicitly inserted with a custom `sequence_id`. The
+        counter may be scoped per rollout or per attempt, but the sequence must be
+        strictly increasing for spans emitted by the specified attempt so traces remain
+        totally ordered.
+
+        See [Distributed Tracing][distributed-tracing] for detailed motivations.
+
+        Args:
+            rollout_id: Identifier of the rollout emitting spans.
+            attempt_id: Attempt identifier for the upcoming span.
+
+        Returns:
+            The next integer sequence identifier, unique within the attempt.
+
+        Raises:
+            NotImplementedError: Subclasses must provide the allocator.
+            ValueError: Implementations must raise when the rollout or attempt does not exist.
        """
        raise NotImplementedError()

    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
-        """
-        Wait for specified rollouts to complete with a timeout.
-        Returns the completed rollouts, potentially incomplete if timeout is reached.
+        """Block until the targeted rollouts reach a terminal status or the timeout expires.

-        TODO: Add support for waiting for 20 new rollouts, or wait until 80% of the pending ids are completed.
+        Terminal statuses are `"succeeded"`, `"failed"`, and `"cancelled"`. When the timeout
+        elapses, implementations should return the subset of rollouts that are already terminal
+        and omit the rest.
+
+        !!! warning
+            It's dangerous and might be event-loop blocking to call this function
+            with a long timeout. It's a good idea to poll for the method to check
+            if new completed rollouts can coming. Be careful in implementing the sleep logic
+            to avoid busy-waiting.
+
+        Args:
+            rollout_ids: Identifiers of rollouts to watch.
+            timeout: Maximum time in seconds to wait. `None` waits indefinitely.
+
+        Returns:
+            Rollouts that finished before the deadline, in arbitrary order.
+
+        Raises:
+            NotImplementedError: Subclasses must implement waiting semantics.
+            ValueError: Implementations must raise when a rollout identifier is unknown.
        """
        raise NotImplementedError()

    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
-        """
-        Query and retrieve all spans associated with a specific rollout ID.
-        Returns an empty list if no spans are found.
+        """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:
-        """
-        Safely stores a new version of named resources and sets it as the latest.
-        Not implemented by many stores yet.
+        """Persist a new immutable snapshot of named resources and mark it as latest.
+
+        Implementations must assign a fresh `resources_id` and ensure subsequent calls to
+        [`get_latest_resources()`][agentlightning.LightningStore.get_latest_resources] return the
+        snapshot produced here.
+
+        Args:
+            resources: Mapping of resource names to their serialized payloads.
+
+        Returns:
+            The stored [`ResourcesUpdate`][agentlightning.ResourcesUpdate] including its generated id.
+
+        Raises:
+            NotImplementedError: Subclasses must implement resource persistence.
        """
        raise NotImplementedError()

    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
-        """
-        Safely stores a new version or updates an existing version of named resources and sets it as the latest.
+        """Overwrite or extend an existing resource snapshot and mark it as latest.
+
+        This API is typically used by algorithms that maintain mutable resources (e.g., model
+        checkpoints) under a stable identifier.
+
+        Args:
+            resources_id: Identifier of the snapshot to replace.
+            resources: Updated mapping of resource names to payloads.
+
+        Returns:
+            The persisted [`ResourcesUpdate`][agentlightning.ResourcesUpdate].
+
+        Raises:
+            NotImplementedError: Subclasses must implement resource persistence.
+            ValueError: Implementations must raise when `resources_id` does not exist.
        """
        raise NotImplementedError()

@@ -220,22 +450,31 @@ class LightningStore:
        config: RolloutConfig | Unset = UNSET,
        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
    ) -> Rollout:
-        """
-        Update the rollout status and related metadata.
+        """Update rollout metadata and, when provided, drive status transitions.

-        Not-listed fields here either cannot be updated, or should be auto-updated (e.g., end_time).
+        Parameters default to the sentinel [`UNSET`][agentlightning.store.base.UNSET] to
+        distinguish omitted fields from explicit `None` assignments. Implementations must:

-        When status is updated to a finished / problematic state, other states like task
-        queues will be updated accordingly.
+        * Validate the rollout exists before mutating it.
+        * Replace each property when a concrete value (including `None`) is supplied.
+        * When the status switches into a terminal state, set `end_time` and signal any waiters.
+        * When the status re-enters a queueing state, ensure the rollout is enqueued exactly once.

        Args:
-            rollout_id: Unique identifier for the rollout to update
-            input: New input data for the rollout. If set, will be updated. Can be updated to None
-            mode: New mode for the rollout. If set, will be updated. Can be updated to None
-            resources_id: New resources ID for the rollout. If set, will be updated. Can be updated to None
-            status: New status for the rollout. If set, will be updated
-            config: New config for the rollout. If set, will be updated
-            metadata: Dictionary of additional metadata to update. If set, will replace the existing metadata
+            rollout_id: Identifier of the rollout to update.
+            input: Replacement task payload; pass `None` to explicitly clear the input.
+            mode: Replacement rollout mode.
+            resources_id: Replacement resources snapshot reference.
+            status: Target rollout status.
+            config: Replacement retry/timeout configuration.
+            metadata: Replacement metadata dictionary.
+
+        Returns:
+            The updated rollout record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement mutation logic.
+            ValueError: Implementations must raise when the rollout is unknown or the update is invalid.
        """
        raise NotImplementedError()

@@ -248,18 +487,29 @@ class LightningStore:
        last_heartbeat_time: float | Unset = UNSET,
        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
    ) -> Attempt:
-        """
-        Update a specific or latest attempt for a given rollout.
+        """Update attempt bookkeeping such as status, worker ownership, and heartbeats.

-        Update the latest attempt will NOT affect the corresponding rollout status.
+        When `attempt_id` is `"latest"` the update must target the attempt with the highest
+        `sequence_id`; otherwise it must target the specific attempt. Implementations should
+        propagate status changes to the rollout (for example via [`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: Unique identifier for the rollout
-            attempt_id: Unique identifier for the attempt
-            status: Status to set for the attempt, update if provided
-            worker_id: Worker identifier, update if provided
-            last_heartbeat_time: Timestamp of the last heartbeat from the worker
-            metadata: Dictionary of additional metadata to update, will replace the existing metadata
+            rollout_id: Identifier of the rollout whose attempt will be updated.
+            attempt_id: Attempt identifier or `"latest"` as a convenience.
+            status: Replacement attempt status. Terminal statuses must set `end_time`.
+            worker_id: Identifier for the worker currently processing the attempt.
+            last_heartbeat_time: Wall-clock timestamp (seconds) of the latest heartbeat/span.
+            metadata: Replacement metadata dictionary.
+
+        Returns:
+            The updated attempt record.
+
+        Raises:
+            NotImplementedError: Subclasses must implement mutation logic.
+            ValueError: Implementations must raise when the rollout or attempt is unknown.
        """
        raise NotImplementedError()
@@ -6,13 +6,31 @@ import asyncio
 import functools
 import hashlib
 import logging
+import sys
 import threading
 import time
 import uuid
+import weakref
 from collections import deque
-from typing import Any, Callable, Counter, Dict, List, Literal, Optional, Sequence, TypeVar, cast
+from collections.abc import Iterable
+from collections.abc import Mapping as MappingABC
+from typing import (
+    Any,
+    Callable,
+    Counter,
+    Dict,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Set,
+    TypeVar,
+    cast,
+)

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

 from agentlightning.types import (
    Attempt,
@@ -35,6 +53,61 @@ T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 logger = logging.getLogger(__name__)


+class _LoopAwareAsyncLock:
+    """Async lock that transparently rebinds to the current event loop.
+
+    The lock intentionally remains *thread-unsafe*: callers must only use it from
+    one thread at a time. If multiple threads interact with the store, each
+    thread gets its own event loop specific lock.
+    """
+
+    def __init__(self) -> None:
+        self._locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = weakref.WeakKeyDictionary()
+
+    # When serializing and deserializing, we don't need to serialize the locks.
+    # Because another process will have its own set of event loops and its own lock.
+    def __getstate__(self) -> dict[str, Any]:
+        return {}
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self._locks = weakref.WeakKeyDictionary()
+
+    def _get_lock_for_current_loop(self) -> asyncio.Lock:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._locks[loop] = lock
+        return lock
+
+    async def __aenter__(self) -> asyncio.Lock:
+        lock = self._get_lock_for_current_loop()
+        await lock.acquire()
+        return lock
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> None:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None or not lock.locked():
+            raise RuntimeError("Lock released without being acquired")
+        lock.release()
+
+
+def estimate_model_size(obj: Any) -> int:
+    """Rough recursive size estimate for Pydantic BaseModel instances."""
+
+    if isinstance(obj, BaseModel):
+        values = cast(Iterable[Any], obj.__dict__.values())
+        return sum(estimate_model_size(value) for value in values) + sys.getsizeof(cast(object, obj))
+    if isinstance(obj, MappingABC):
+        mapping = cast(Mapping[Any, Any], obj)
+        return sum(estimate_model_size(value) for value in mapping.values()) + sys.getsizeof(cast(object, obj))
+    if isinstance(obj, (list, tuple, set)):
+        iterable = cast(Iterable[Any], obj)
+        return sum(estimate_model_size(value) for value in iterable) + sys.getsizeof(cast(object, obj))
+    return sys.getsizeof(cast(object, obj))
+
+
 def _healthcheck_wrapper(func: T_callable) -> T_callable:
    """
    Decorator to run the watchdog healthcheck **before** executing the decorated method.
@@ -82,6 +155,19 @@ def _generate_attempt_id() -> str:
    return "at-" + short_id


+def _detect_total_memory_bytes() -> int:
+    """Best-effort detection of the total available system memory in bytes."""
+
+    try:
+        import psutil
+
+        return int(psutil.virtual_memory().total)
+    except ImportError:
+        # Fallback to 8GB if memory cannot be detected.
+        logger.error("psutil is not installed. Falling back to 8GB of memory in total.")
+        return 8 * 1024**3
+
+
 class InMemoryLightningStore(LightningStore):
    """
    In-memory implementation of LightningStore using Python data structures.
@@ -89,10 +175,24 @@ class InMemoryLightningStore(LightningStore):

    The methods in this class should generally not call each other,
    especially those that are locked.
+
+    Args:
+        eviction_memory_threshold: The threshold for evicting spans in bytes.
+            By default, it's 70% of the total VRAM available.
+        safe_memory_threshold: The threshold for safe memory usage in bytes.
+            By default, it's 80% of the eviction threshold.
+        span_size_estimator: A function to estimate the size of a span in bytes.
+            By default, it's a simple size estimator that uses sys.getsizeof.
    """

-    def __init__(self):
-        self._lock = asyncio.Lock()
+    def __init__(
+        self,
+        *,
+        eviction_memory_threshold: float | int | None = None,
+        safe_memory_threshold: float | int | None = None,
+        span_size_estimator: Callable[[Span], int] | None = None,
+    ):
+        self._lock = _LoopAwareAsyncLock()

        # Task queue and rollouts storage
        self._task_queue: deque[Rollout] = deque()
@@ -105,6 +205,36 @@ class InMemoryLightningStore(LightningStore):
        # Spans storage
        self._spans: Dict[str, List[Span]] = {}  # rollout_id -> list of spans
        self._span_sequence_ids: Dict[str, int] = Counter()  # rollout_id -> sequence_id
+        self._span_bytes_by_rollout: Dict[str, int] = Counter()
+        self._total_span_bytes: int = 0
+        self._evicted_rollout_span_sets: Set[str] = set()
+
+        self._memory_capacity_bytes = _detect_total_memory_bytes()
+        if self._memory_capacity_bytes <= 0:
+            raise ValueError("Detected memory capacity must be positive")
+
+        self._eviction_threshold_bytes = self._resolve_memory_threshold(
+            eviction_memory_threshold,
+            default_ratio=0.7,
+            capacity_bytes=self._memory_capacity_bytes,
+            name="eviction_memory_threshold",
+            minimum=1,
+        )
+
+        if safe_memory_threshold is None:
+            safe_memory_threshold = max(int(self._eviction_threshold_bytes * 0.8), 0)
+
+        self._safe_threshold_bytes = self._resolve_memory_threshold(
+            safe_memory_threshold,
+            default_ratio=self._eviction_threshold_bytes / self._memory_capacity_bytes,
+            capacity_bytes=self._memory_capacity_bytes,
+            name="safe_memory_threshold",
+            minimum=0,
+        )
+
+        if not (0 <= self._safe_threshold_bytes < self._eviction_threshold_bytes):
+            raise ValueError("safe_memory_threshold must be smaller than eviction_memory_threshold")
+        self._custom_span_size_estimator = span_size_estimator

        # Attempt tracking
        self._attempts: Dict[str, List[Attempt]] = {}  # rollout_id -> list of attempts
@@ -118,15 +248,20 @@ class InMemoryLightningStore(LightningStore):
        input: TaskInput,
        mode: Literal["train", "val", "test"] | None = None,
        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> AttemptedRollout:
-        """
-        Notify the store that I'm about to run a rollout.
+        """Notify the store that I'm about to run a rollout.
+
+        See [`LightningStore.start_rollout()`][agentlightning.LightningStore.start_rollout] for semantics.
        """
        async with self._lock:
            rollout_id = _generate_rollout_id()
            current_time = time.time()

+            rollout_config = config.model_copy(deep=True) if config is not None else RolloutConfig()
+            rollout_metadata = dict(metadata) if metadata is not None else {}
+
            rollout = Rollout(
                rollout_id=rollout_id,
                input=input,
@@ -134,7 +269,8 @@ class InMemoryLightningStore(LightningStore):
                resources_id=resources_id or self._latest_resources_id,
                start_time=current_time,
                status="preparing",
-                metadata=metadata or {},
+                config=rollout_config,
+                metadata=rollout_metadata,
            )

            # Create the initial attempt
@@ -161,15 +297,20 @@ class InMemoryLightningStore(LightningStore):
        input: TaskInput,
        mode: Literal["train", "val", "test"] | None = None,
        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> Rollout:
-        """
-        Adds a new task to the queue with specific metadata and returns its unique ID.
+        """Adds a new task to the queue with specific metadata and returns the rollout.
+
+        See [`LightningStore.enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout] for semantics.
        """
        async with self._lock:
            rollout_id = _generate_rollout_id()
            current_time = time.time()

+            rollout_config = config.model_copy(deep=True) if config is not None else RolloutConfig()
+            rollout_metadata = dict(metadata) if metadata is not None else {}
+
            rollout = Rollout(
                rollout_id=rollout_id,
                input=input,
@@ -177,7 +318,8 @@ class InMemoryLightningStore(LightningStore):
                resources_id=resources_id or self._latest_resources_id,
                start_time=current_time,
                status="queuing",  # should be queuing
-                metadata=metadata or {},
+                config=rollout_config,
+                metadata=rollout_metadata,
            )

            self._rollouts[rollout.rollout_id] = rollout
@@ -188,11 +330,12 @@ class InMemoryLightningStore(LightningStore):

    @_healthcheck_wrapper
    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        """
-        Retrieves the next task from the queue without blocking.
-        Returns None if the queue is empty.
+        """Retrieves the next task from the queue without blocking.
+        Returns `None` if the queue is empty.

        Will set the rollout status to preparing and create a new attempt.
+
+        See [`LightningStore.dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] for semantics.
        """
        async with self._lock:
            # Keep looking until we find a rollout that's still in queuing status
@@ -236,8 +379,9 @@ class InMemoryLightningStore(LightningStore):

    @_healthcheck_wrapper
    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
-        """
-        Create a new attempt for a given rollout ID and return the attempt details.
+        """Creates a new attempt for a given rollout ID and return the attempt details.
+
+        See [`LightningStore.start_attempt()`][agentlightning.LightningStore.start_attempt] for semantics.
        """
        async with self._lock:
            # Get the rollout
@@ -277,9 +421,10 @@ class InMemoryLightningStore(LightningStore):
    async def query_rollouts(
        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
    ) -> List[Rollout]:
-        """
-        Query and retrieve rollouts filtered by their status and rollout ids.
+        """Retrieves rollouts filtered by their status and rollout ids.
        If no status is provided, returns all rollouts.
+
+        See [`LightningStore.query_rollouts()`][agentlightning.LightningStore.query_rollouts] for semantics.
        """
        async with self._lock:
            rollouts = list(self._rollouts.values())
@@ -298,25 +443,28 @@ class InMemoryLightningStore(LightningStore):

    @_healthcheck_wrapper
    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
-        """
-        Safely retrieves a specific rollout by its ID.
+        """Retrieves a specific rollout by its ID.
+
+        See [`LightningStore.get_rollout_by_id()`][agentlightning.LightningStore.get_rollout_by_id] for semantics.
        """
        async with self._lock:
            return self._rollouts.get(rollout_id)

    @_healthcheck_wrapper
    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
-        """
-        Query and retrieve all attempts associated with a specific rollout ID.
+        """Retrieves all attempts associated with a specific rollout ID.
        Returns an empty list if no attempts are found.
+
+        See [`LightningStore.query_attempts()`][agentlightning.LightningStore.query_attempts] for semantics.
        """
        async with self._lock:
            return self._attempts.get(rollout_id, [])

    @_healthcheck_wrapper
    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
-        """
-        Safely retrieves the latest attempt for a given rollout ID.
+        """Retrieves the latest attempt for a given rollout ID.
+
+        See [`LightningStore.get_latest_attempt()`][agentlightning.LightningStore.get_latest_attempt] for semantics.
        """
        async with self._lock:
            attempts = self._attempts.get(rollout_id, [])
@@ -326,8 +474,9 @@ class InMemoryLightningStore(LightningStore):

    @_healthcheck_wrapper
    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
-        """
-        Safely stores a new version of named resources and sets it as the latest.
+        """Stores a new version of named resources and sets it as the latest.
+
+        See [`LightningStore.add_resources()`][agentlightning.LightningStore.add_resources] for semantics.
        """
        resources_id = _generate_resources_id()
        async with self._lock:
@@ -340,6 +489,8 @@ class InMemoryLightningStore(LightningStore):
    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
        """
        Safely stores a new version of named resources and sets it as the latest.
+
+        See [`LightningStore.update_resources()`][agentlightning.LightningStore.update_resources] for semantics.
        """
        async with self._lock:
            update = ResourcesUpdate(resources_id=resources_id, resources=resources)
@@ -349,16 +500,18 @@ class InMemoryLightningStore(LightningStore):

    @_healthcheck_wrapper
    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves a specific version of named resources by its ID.
+        """Retrieves a specific version of named resources by its ID.
+
+        See [`LightningStore.get_resources_by_id()`][agentlightning.LightningStore.get_resources_by_id] for semantics.
        """
        async with self._lock:
            return self._resources.get(resources_id)

    @_healthcheck_wrapper
    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        """
-        Safely retrieves the latest version of named resources.
+        """Retrieves the latest version of named resources.
+
+        See [`LightningStore.get_latest_resources()`][agentlightning.LightningStore.get_latest_resources] for semantics.
        """
        async with self._lock:
            if self._latest_resources_id:
@@ -366,17 +519,21 @@ class InMemoryLightningStore(LightningStore):
            return None

    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
-        """
-        Get the next span sequence ID for a given rollout and attempt.
+        """Get the next span sequence ID for a given rollout and attempt.
        The number is strictly increasing for each rollout.
        The store will not issue the same sequence ID twice.
+
+        See [`LightningStore.get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id] for semantics.
        """
        async with self._lock:
            self._span_sequence_ids[rollout_id] += 1
            return self._span_sequence_ids[rollout_id]

    async def add_span(self, span: Span) -> Span:
-        """Persist a pre-converted span."""
+        """Persist a pre-converted span.
+
+        See [`LightningStore.add_span()`][agentlightning.LightningStore.add_span] for semantics.
+        """
        async with self._lock:
            self._span_sequence_ids[span.rollout_id] = max(self._span_sequence_ids[span.rollout_id], span.sequence_id)
            return await self._add_span_unlocked(span)
@@ -384,7 +541,10 @@ class InMemoryLightningStore(LightningStore):
    async def add_otel_span(
        self, rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None
    ) -> Span:
-        """Add an opentelemetry span to the store."""
+        """Add an opentelemetry span to the store.
+
+        See [`LightningStore.add_otel_span()`][agentlightning.LightningStore.add_otel_span] for semantics.
+        """
        async with self._lock:
            if sequence_id is None:
                # Issue a new sequence ID for the rollout
@@ -416,10 +576,12 @@ class InMemoryLightningStore(LightningStore):
        if span.rollout_id not in self._spans:
            self._spans[span.rollout_id] = []
        self._spans[span.rollout_id].append(span)
+        self._account_span_size(span)
+        self._maybe_evict_spans()

        # Update attempt heartbeat
        current_attempt.last_heartbeat_time = time.time()
-        if current_attempt.status in ["preparing", "unresponsive", "timeout"]:
+        if current_attempt.status in ["preparing", "unresponsive"]:
            current_attempt.status = "running"

        # If the status has already timed out or failed, do not change it
@@ -439,13 +601,85 @@ class InMemoryLightningStore(LightningStore):

        return span

+    @staticmethod
+    def _resolve_memory_threshold(
+        value: float | int | None,
+        *,
+        default_ratio: float,
+        capacity_bytes: int,
+        name: str,
+        minimum: int,
+    ) -> int:
+        if value is None:
+            resolved = int(capacity_bytes * default_ratio)
+        elif isinstance(value, float):
+            if minimum == 0:
+                if not (0 <= value <= 1):
+                    raise ValueError(f"{name} ratio must be between 0 and 1 inclusive")
+            else:
+                if not (0 < value <= 1):
+                    raise ValueError(f"{name} ratio must be greater than 0 and at most 1")
+            resolved = int(capacity_bytes * value)
+        else:
+            value_int = value
+            if value_int < 0:
+                raise ValueError(f"{name} must be non-negative")
+            resolved = value_int
+
+        if resolved < minimum:
+            raise ValueError(f"{name} must be at least {minimum} bytes")
+
+        return resolved
+
+    def _account_span_size(self, span: Span) -> int:
+        if self._custom_span_size_estimator is not None:
+            size = max(int(self._custom_span_size_estimator(span)), 0)
+        else:
+            size = estimate_model_size(span)
+
+        self._span_bytes_by_rollout[span.rollout_id] += size
+        self._total_span_bytes += size
+        return size
+
+    def _maybe_evict_spans(self) -> None:
+        if self._total_span_bytes <= self._eviction_threshold_bytes:
+            return
+
+        candidates: List[tuple[float, str]] = []
+        for rollout_id, spans in self._spans.items():
+            if not spans:
+                continue
+            rollout = self._rollouts.get(rollout_id)
+            start_time = rollout.start_time if rollout is not None else (spans[0].start_time or 0.0)
+            candidates.append((start_time, rollout_id))
+
+        candidates.sort(key=lambda item: item[0])
+
+        logger.info(f"Evicting spans for {len(candidates)} rollouts to free up memory...")
+        memory_consumed_before = self._total_span_bytes
+        for _, rollout_id in candidates:
+            if self._total_span_bytes <= self._safe_threshold_bytes:
+                break
+            logger.debug(f"Evicting spans for rollout {rollout_id} to free up memory...")
+            self._evict_spans_for_rollout(rollout_id)
+        logger.info(f"Freed up {memory_consumed_before - self._total_span_bytes} bytes of memory")
+
+    def _evict_spans_for_rollout(self, rollout_id: str) -> None:
+        spans = self._spans.pop(rollout_id, [])
+        if not spans:
+            return
+        removed_bytes = self._span_bytes_by_rollout.pop(rollout_id, 0)
+        self._total_span_bytes = max(self._total_span_bytes - removed_bytes, 0)
+        self._evicted_rollout_span_sets.add(rollout_id)
+
    @_healthcheck_wrapper
    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
-        """
-        Wait for specified rollouts to complete with a timeout.
+        """Wait for specified rollouts to complete with a timeout.
        Returns the completed rollouts, potentially incomplete if timeout is reached.

        This method does not change the state of the store.
+
+        See [`LightningStore.wait_for_rollouts()`][agentlightning.LightningStore.wait_for_rollouts] for semantics.
        """
        completed_rollouts: List[Rollout] = []

@@ -498,8 +732,12 @@ class InMemoryLightningStore(LightningStore):
        """
        Query and retrieve all spans associated with a specific rollout ID.
        Returns an empty list if no spans are found.
+
+        See [`LightningStore.query_spans()`][agentlightning.LightningStore.query_spans] for semantics.
        """
        async with self._lock:
+            if rollout_id in self._evicted_rollout_span_sets:
+                raise RuntimeError(f"Spans for rollout {rollout_id} have been evicted")
            spans = self._spans.get(rollout_id, [])
            if attempt_id is None:
                return spans
@@ -523,8 +761,9 @@ class InMemoryLightningStore(LightningStore):
        config: RolloutConfig | Unset = UNSET,
        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
    ) -> Rollout:
-        """
-        Update the rollout status and related metadata.
+        """Update the rollout status and related metadata.
+
+        See [`LightningStore.update_rollout()`][agentlightning.LightningStore.update_rollout] for semantics.
        """
        async with self._lock:
            return await self._update_rollout_unlocked(
@@ -547,8 +786,9 @@ class InMemoryLightningStore(LightningStore):
        last_heartbeat_time: float | Unset = UNSET,
        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
    ) -> Attempt:
-        """
-        Update a specific or latest attempt for a given rollout.
+        """Update a specific or latest attempt for a given rollout.
+
+        See [`LightningStore.update_attempt()`][agentlightning.LightningStore.update_attempt] for semantics.
        """
        async with self._lock:
            attempt = await self._update_attempt_unlocked(
@@ -40,20 +40,22 @@ class LightningStoreThreaded(LightningStore):
        input: TaskInput,
        mode: Literal["train", "val", "test"] | None = None,
        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> AttemptedRollout:
        with self._lock:
-            return await self.store.start_rollout(input, mode, resources_id, metadata)
+            return await self.store.start_rollout(input, mode, resources_id, config, metadata)

    async def enqueue_rollout(
        self,
        input: TaskInput,
        mode: Literal["train", "val", "test"] | None = None,
        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> Rollout:
        with self._lock:
-            return await self.store.enqueue_rollout(input, mode, resources_id, metadata)
+            return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)

    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
        with self._lock:
@@ -57,6 +57,7 @@ async def healthcheck(
    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)
@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.

 from .agentops import AgentOpsTracer
-from .base import BaseTracer
+from .base import Tracer
 from .otel import OtelTracer

-__all__ = ["AgentOpsTracer", "BaseTracer", "OtelTracer"]
+__all__ = ["AgentOpsTracer", "Tracer", "OtelTracer"]
@@ -6,20 +6,21 @@ import asyncio
 import logging
 import os
 import threading
-from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Awaitable, Iterator, List, Optional
+from contextlib import asynccontextmanager, contextmanager
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Iterator, List, Optional

 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 agentlightning.instrumentation import instrument_all, uninstrument_all
 from agentlightning.instrumentation.agentops import AgentOpsServerManager
 from agentlightning.store.base import LightningStore

-from .base import BaseTracer
+from .base import Tracer

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


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

    This tracer provides functionality to capture execution details using the
@@ -152,15 +153,15 @@ class AgentOpsTracer(BaseTracer):
            self.uninstrument(worker_id)
            logger.info(f"[Worker {worker_id}] Instrumentation removed.")

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

@@ -171,8 +172,23 @@ class AgentOpsTracer(BaseTracer):
            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.")

@@ -224,42 +240,11 @@ class AgentOpsTracer(BaseTracer):
    get_langchain_callback_handler = get_langchain_handler  # alias


-async def heartbeat(name="exporter-loop", period=0.5):
-    import asyncio
-    import time
-
-    last = time.perf_counter()
-    while True:
-        await asyncio.sleep(period)
-        now = time.perf_counter()
-        dt = now - last
-        last = now
-        if dt > period * 4:  # e.g., >2s if period=0.5s
-            print("!!!!!!! [%s] loop stall detected: slept %.3fs (expected %.3fs)" % (name, dt, period))
-
-
-import asyncio
-import logging
-
-# logging.basicConfig(level=logging.DEBUG)
-# asyncio.get_event_loop().set_debug(True)
-import time
-
-
-def debug_dump(loop):
-    while True:
-        try:
-            print("=== Pending tasks ===")
-            for t in asyncio.all_tasks(loop):
-                if not t.done():
-                    print(t, "awaiting", t.get_coro())
-                    t.print_stack()
-        except Exception:
-            pass
-        time.sleep(5)
-
-
 class LightningSpanProcessor(SpanProcessor):
+    """Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
+    to a [`LightningStore`][agentlightning.LightningStore].
+    """
+
    def __init__(self):
        self._spans: List[ReadableSpan] = []

@@ -279,13 +264,8 @@ class LightningSpanProcessor(SpanProcessor):
    def _loop_runner(self):
        loop = asyncio.new_event_loop()
        self._loop = loop
-        self._loop.set_debug(True)
        asyncio.set_event_loop(loop)
        self._loop_ready.set()
-
-        thread = threading.Thread(target=debug_dump, args=(loop,), daemon=True)
-        thread.start()
-        # asyncio.create_task(heartbeat())
        loop.run_forever()
        loop.close()

@@ -374,10 +354,6 @@ class LightningSpanProcessor(SpanProcessor):
        Args:
            span: The span that has ended.
        """
-        import traceback
-
-        # print("ON_END")
-        # print(traceback.format_stack())
        # Skip if span is not sampled
        if not span.context or not span.context.trace_flags.sampled:
            return
@@ -385,27 +361,11 @@ class LightningSpanProcessor(SpanProcessor):
        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
-                print("!!! before,")
-                print("Ready callbacks:", self._loop._ready)
-                print("Scheduled callbacks:", len(self._loop._scheduled))
-                if self._loop._scheduled:
-                    print("First in the queue:", self._loop._scheduled[0])
-                print("..... Current thread: ", threading.current_thread())
-                print("..... Loop thread: ", self._loop_thread)
-                if self._loop_thread.ident == threading.current_thread().ident:
-                    traceback.print_stack()
-                    print("Span content: ", span.attributes)
-                from opentelemetry.instrumentation.utils import suppress_instrumentation
-
                with suppress_instrumentation():
                    self._await_in_loop(
                        self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
-                        timeout=30.0,
+                        timeout=60.0,
                    )
-                print("!!! after,")
-                print("All tasks")
-                print("Ready callbacks:", self._loop._ready)
-                print("Scheduled callbacks:", self._loop._scheduled)
            except Exception:
                # log; on_end MUST NOT raise
                logger.exception(f"Error adding span to store: {span.name}")
@@ -3,8 +3,7 @@
 from __future__ import annotations

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

 from opentelemetry.sdk.trace import ReadableSpan

@@ -12,12 +11,12 @@ from agentlightning.store.base import LightningStore
 from agentlightning.types import ParallelWorkerBase

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

 logger = logging.getLogger(__name__)


-class BaseTracer(ParallelWorkerBase):
+class Tracer(ParallelWorkerBase):
    """
    An abstract base class for tracers.

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

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

@@ -52,7 +51,6 @@ class BaseTracer(ParallelWorkerBase):
    ```
    """

-    @contextmanager
    def trace_context(
        self,
        name: Optional[str] = None,
@@ -60,14 +58,14 @@ class BaseTracer(ParallelWorkerBase):
        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
-    ) -> Iterator[Any]:
+    ) -> 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.

@@ -79,6 +77,17 @@ class BaseTracer(ParallelWorkerBase):
        """
        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.
@@ -92,6 +101,8 @@ class BaseTracer(ParallelWorkerBase):
        """
        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.
@@ -100,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[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.
@@ -115,10 +128,10 @@ 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]:
+    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.
@@ -5,8 +5,8 @@ import logging
 import multiprocessing
 import queue
 import uuid
-from contextlib import contextmanager
-from typing import Any, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple
 from urllib.parse import urlparse

 from httpdbg.hooks.all import httprecord
@@ -19,12 +19,12 @@ from opentelemetry.trace.span import (
    TraceState,
 )

-from .base import BaseTracer
+from .base import Tracer

 logger = logging.getLogger(__name__)


-class HttpTracer(BaseTracer):
+class HttpTracer(Tracer):
    """
    A tracer implementation that captures HTTP requests using httpdbg.

@@ -78,8 +78,19 @@ class HttpTracer(BaseTracer):
        super().init_worker(worker_id)
        logger.info(f"[Worker {worker_id}] HttpTracer initialized.")

+    @asynccontextmanager
+    async def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> AsyncGenerator[HTTPRecords, None]:
+        """
+        Starts a new HTTP tracing context. This should be used as a context manager.
+
+        Args:
+            name: Optional name for the tracing context.
+        """
+        with self._trace_context_sync(name=name, **kwargs) as records:
+            yield records
+
    @contextmanager
-    def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]:
+    def _trace_context_sync(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]:
        """
        Starts a new HTTP tracing context. This should be used as a context manager.

@@ -3,8 +3,8 @@
 from __future__ import annotations

 import logging
-from contextlib import contextmanager
-from typing import Iterator, List, Optional
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator, List, Optional

 import opentelemetry.trace as trace_api
 from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
@@ -12,12 +12,12 @@ from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
 from agentlightning.store.base import LightningStore

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

 logger = logging.getLogger(__name__)


-class OtelTracer(BaseTracer):
+class OtelTracer(Tracer):
    """Tracer that provides a basic OpenTelemetry tracer provider.

    You should be able to collect agent-lightning signals like rewards with this tracer,
@@ -49,15 +49,15 @@ class OtelTracer(BaseTracer):
        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
        self._tracer_provider = None

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

@@ -9,11 +9,11 @@ import warnings
 from typing import Any, List, Optional, TypeVar, Union

 from agentlightning.adapter import TraceAdapter, TracerTraceToTriplet
-from agentlightning.algorithm import BaseAlgorithm
+from agentlightning.algorithm import Algorithm
 from agentlightning.client import AgentLightningClient
 from agentlightning.litagent import LitAgent
 from agentlightning.runner import LegacyAgentRunner
-from agentlightning.tracer.base import BaseTracer
+from agentlightning.tracer.base import Tracer
 from agentlightning.types import Dataset, ParallelWorkerBase

 logger = logging.getLogger(__name__)
@@ -31,8 +31,8 @@ class TrainerLegacy(ParallelWorkerBase):
        It won't be used in practice.
        """
        self._dev = kwargs.pop("dev", False)
-        self.algorithm: Optional[BaseAlgorithm] = kwargs.pop("algorithm", None)
-        self.tracer: BaseTracer = kwargs.pop("tracer", None)
+        self.algorithm: Optional[Algorithm] = kwargs.pop("algorithm", None)
+        self.tracer: Tracer = kwargs.pop("tracer", None)
        self.n_workers: int = kwargs.pop("n_workers", None)
        self.max_tasks: Optional[int] = kwargs.pop("max_tasks", None)
        self.daemon: bool = kwargs.pop("daemon", True)
@@ -7,18 +7,18 @@ import warnings
 from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, Union

 from agentlightning.adapter import TraceAdapter, TracerTraceToTriplet
-from agentlightning.algorithm import BaseAlgorithm, Baseline, FastAlgorithm
+from agentlightning.algorithm import Algorithm, Baseline, FastAlgorithm
 from agentlightning.client import AgentLightningClient
 from agentlightning.execution.base import ExecutionStrategy
 from agentlightning.execution.client_server import ClientServerExecutionStrategy
 from agentlightning.execution.events import ExecutionEvent
 from agentlightning.litagent import LitAgent
 from agentlightning.llm_proxy import LLMProxy
-from agentlightning.runner import BaseRunner, LitAgentRunner
+from agentlightning.runner import LitAgentRunner, Runner
 from agentlightning.store.base import LightningStore
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.tracer.agentops import AgentOpsTracer
-from agentlightning.tracer.base import BaseTracer
+from agentlightning.tracer.base import Tracer
 from agentlightning.types import Dataset, Hook, NamedResources

 from .init_utils import build_component, instantiate_component
@@ -34,44 +34,89 @@ ComponentSpec = Union[T, type[T], Callable[[], T], str, Dict[str, Any], None]


 class Trainer(TrainerLegacy):
-    """Orchestrates the distributed execution of agent rollouts.
+    """High-level orchestration layer that wires Algorithm <-> Runner <-> Store.

-    The Trainer is responsible for launching one or more worker processes
-    that run the agent's execution loop. It manages multiprocessing,
-    handles graceful shutdown, and serves as the main entry point for
-    running a client-side agent fleet.
+    A [`Trainer`][agentlightning.Trainer] packages the moving parts of Agent-Lightning's
+    training loop into a single entry point:

-    Attributes:
-        algorithm: An instance of `BaseAlgorithm` to use for training.
-        store: An instance of `LightningStore` to use for storing tasks and traces.
-        runner: An instance of `BaseRunner` to use for running the agent.
-        initial_resources: An instance of `Resources` to use for bootstrapping the fit/dev process.
-            The resources will be handed over to the algorithm.
-            Note that not all algorithms support seeding resources.
-        n_runners: Number of agent runners to run in parallel.
-        max_rollouts: Maximum number of rollouts to process per runner. If None,
-                      workers run until no more rollouts are available.
-        strategy: An instance of `ExecutionStrategy` to use for spawning the algorithm and runners.
-        tracer: A tracer instance, or a string pointing to the class full name or a dictionary with a 'type' key
-                that specifies the class full name and other initialization parameters.
-                If None, a default `AgentOpsTracer` will be created with the current settings.
-        hooks: A sequence of `Hook` instances to be called at various lifecycle stages (e.g., on_trace_start,
-               on_trace_end, on_rollout_start, on_rollout_end).
-        adapter: An instance of `TracerTraceToTriplet` to export data consumble by algorithms from traces.
-        llm_proxy: An instance of `LLMProxy` to use for intercepting the LLM calls.
-                   If not provided, algorithm will create one on its own.
-        n_workers: Number of agent workers to run in parallel. Deprecated in favor of `n_runners`.
-        max_tasks: Maximum number of tasks to process per runner. Deprecated in favor of `max_rollouts`.
-        daemon: Whether worker processes should be daemons. Daemon processes
-                are terminated automatically when the main process exits. Deprecated.
-                Only have effect with `fit_v0`.
-        triplet_exporter: An instance of `TracerTraceToTriplet` to export triplets from traces,
-                          or a dictionary with the initialization parameters for the exporter.
-                          Deprecated. Use `adapter` instead.
-        dev: If True, rollouts are run against the dev endpoint provided in `fit`.
-             Deprecated in favor of `dev()` method.
+    * **Algorithm lifecycle:** Instantiates or accepts an [`Algorithm`][agentlightning.Algorithm],
+      attaches the current [`LightningStore`][agentlightning.LightningStore], adapter, and
+      initial resources, then executes the algorithm role inside the configured execution strategy.
+    * **Runner fleet:** Spawns one or more [`Runner`][agentlightning.Runner] instances (defaulting
+      to [`LitAgentRunner`][agentlightning.LitAgentRunner]) that hydrate a [`LitAgent`][agentlightning.LitAgent],
+      claim rollouts, stream spans, and respect graceful termination signals from the execution strategy.
+    * **Execution strategy:** Delegates process management to an
+      [`ExecutionStrategy`][agentlightning.ExecutionStrategy] (shared memory, client/server, etc.),
+      so advanced users can swap orchestration backends without changing trainer code.
+    * **Telemetry plumbing:** Ensures tracers, adapters, and optional [`LLMProxy`][agentlightning.LLMProxy]
+      are wired into both algorithm and runners so telemetry flows back into the store.
+
+    The trainer exposes two convenience entry points:
+    [`fit()`][agentlightning.Trainer.fit] for full training and
+    [`dev()`][agentlightning.Trainer.dev] for fast, reproducible dry-runs. See the
+    [Train the First Agent](../how-to/train-first-agent.md) and
+    [Write the First Algorithm](../how-to/write-first-algorithm.md) tutorials for the broader context.
    """

+    algorithm: Optional[Algorithm]
+    """An instance of [`Algorithm`][agentlightning.Algorithm] to use for training."""
+
+    store: LightningStore
+    """An instance of [`LightningStore`][agentlightning.LightningStore] to use for storing tasks and traces."""
+
+    runner: Runner[Any]
+    """An instance of [`Runner`][agentlightning.Runner] to use for running the agent."""
+
+    initial_resources: Optional[NamedResources]
+    """An instance of [`NamedResources`][agentlightning.NamedResources] to use for bootstrapping the fit/dev process.
+
+    The resources will be handed over to the algorithm. Note that not all algorithms support seeding resources.
+    """
+
+    n_runners: int
+    """Number of agent runners to run in parallel."""
+
+    max_rollouts: Optional[int]
+    """Maximum number of rollouts to process per runner. If None, workers run until no more rollouts are available."""
+
+    strategy: ExecutionStrategy
+    """An instance of [`ExecutionStrategy`][agentlightning.ExecutionStrategy] to use for spawning the algorithm and runners."""
+
+    tracer: Tracer
+    """A tracer instance, or a string pointing to the class full name or a dictionary with a 'type' key
+    that specifies the class full name and other initialization parameters.
+    If None, a default [`AgentOpsTracer`][agentlightning.AgentOpsTracer] will be created with the current settings."""
+
+    hooks: Sequence[Hook]
+    """A sequence of [`Hook`][agentlightning.Hook] instances to be called at various lifecycle stages (e.g., `on_trace_start`,
+    `on_trace_end`, `on_rollout_start`, `on_rollout_end`)."""
+
+    adapter: TraceAdapter[Any]
+    """An instance of [`TraceAdapter`][agentlightning.TraceAdapter] to export data consumble by algorithms from traces."""
+
+    llm_proxy: Optional[LLMProxy]
+    """An instance of [`LLMProxy`][agentlightning.LLMProxy] to use for intercepting the LLM calls.
+    If not provided, algorithm may create one on its own."""
+
+    n_workers: int
+    """Number of agent workers to run in parallel. Deprecated in favor of `n_runners`."""
+
+    max_tasks: Optional[int]
+    """Maximum number of tasks to process per runner. Deprecated in favor of `max_rollouts`."""
+
+    daemon: bool
+    """Whether worker processes should be daemons. Daemon processes
+    are terminated automatically when the main process exits. Deprecated.
+    Only have effect with `fit_v0`."""
+
+    triplet_exporter: TraceAdapter[Any]
+    """An instance of [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] to export triplets from traces,
+    or a dictionary with the initialization parameters for the exporter.
+    Deprecated. Use [`adapter`][agentlightning.Trainer.adapter] instead."""
+
+    port: Optional[int]
+    """Port forwarded to [`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy]."""
+
    def __init__(
        self,
        *,
@@ -79,12 +124,13 @@ class Trainer(TrainerLegacy):
        n_runners: Optional[int] = None,
        max_rollouts: Optional[int] = None,
        initial_resources: Optional[NamedResources] = None,
-        tracer: ComponentSpec[BaseTracer] = None,
+        tracer: ComponentSpec[Tracer] = None,
        adapter: ComponentSpec[TraceAdapter[Any]] = None,
        store: ComponentSpec[LightningStore] = None,
-        runner: ComponentSpec[BaseRunner[Any]] = None,
+        runner: ComponentSpec[Runner[Any]] = None,
        strategy: ComponentSpec[ExecutionStrategy] = None,
-        algorithm: ComponentSpec[BaseAlgorithm] = None,
+        port: Optional[int] = None,
+        algorithm: ComponentSpec[Algorithm] = None,
        llm_proxy: ComponentSpec[LLMProxy] = None,
        n_workers: Optional[int] = None,
        max_tasks: Optional[int] = None,
@@ -92,6 +138,16 @@ class Trainer(TrainerLegacy):
        triplet_exporter: ComponentSpec[TracerTraceToTriplet] = None,
        hooks: Optional[Union[Hook, Sequence[Hook]]] = None,
    ):
+        """Configure the trainer and resolve user-provided component specifications.
+
+        Each keyword accepts either a concrete instance, a class, a callable factory, a
+        registry string, or a lightweight configuration dictionary (see
+        [`build_component()`][agentlightning.trainer.init_utils.build_component]).
+
+        When ``port`` is provided it is forwarded to
+        [`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy]
+        instances constructed (or supplied) for the trainer.
+        """
        # Do not call super().__init__() here.
        # super().__init__() will call TrainerLegacy's initialization, which is not intended.
        self.worker_id: Optional[int] = None
@@ -161,7 +217,13 @@ class Trainer(TrainerLegacy):
        self.store = self._make_store(store)
        self.runner = self._make_runner(runner)

-        self.strategy = self._make_strategy(strategy, n_runners=self.n_runners)
+        self.port = port
+
+        self.strategy = self._make_strategy(
+            strategy,
+            n_runners=self.n_runners,
+            port=port,
+        )
        if hasattr(self.strategy, "n_runners"):
            strategy_runners = getattr(self.strategy, "n_runners")
            if isinstance(strategy_runners, int) and strategy_runners > 0:
@@ -179,8 +241,8 @@ class Trainer(TrainerLegacy):
                "The cleanup must be handled manually."
            )

-    def _make_tracer(self, tracer: ComponentSpec[BaseTracer]) -> BaseTracer:
-        """Creates a tracer instance based on the provided configuration."""
+    def _make_tracer(self, tracer: ComponentSpec[Tracer]) -> Tracer:
+        """Resolve the tracer component from user input, falling back to AgentOpsTracer."""
        default_factory = lambda: AgentOpsTracer(
            agentops_managed=True,
            instrument_managed=True,
@@ -188,26 +250,27 @@ class Trainer(TrainerLegacy):
        )
        return build_component(
            tracer,
-            expected_type=BaseTracer,
+            expected_type=Tracer,
            spec_name="tracer",
            default_factory=default_factory,
            dict_requires_type=True,
-            invalid_spec_error_fmt="Invalid tracer type: {actual_type}. Expected BaseTracer, str, dict, or None.",
-            type_error_fmt="Tracer factory returned {type_name}, which is not a BaseTracer subclass.",
+            invalid_spec_error_fmt="Invalid tracer type: {actual_type}. Expected Tracer, str, dict, or None.",
+            type_error_fmt="Tracer factory returned {type_name}, which is not a Tracer subclass.",
        )

-    def _make_algorithm(self, algorithm: ComponentSpec[BaseAlgorithm]) -> Optional[BaseAlgorithm]:
-        """Creates an algorithm instance based on the provided configuration."""
+    def _make_algorithm(self, algorithm: ComponentSpec[Algorithm]) -> Optional[Algorithm]:
+        """Resolve the algorithm component, allowing `None` for dev-mode dry runs."""
        return build_component(
            algorithm,
-            expected_type=BaseAlgorithm,
+            expected_type=Algorithm,
            spec_name="algorithm",
            allow_none=True,
-            invalid_spec_error_fmt="Invalid algorithm type: {actual_type}. Expected BaseAlgorithm, str, dict, or None.",
-            type_error_fmt="Algorithm factory returned {type_name}, which is not a BaseAlgorithm subclass.",
+            invalid_spec_error_fmt="Invalid algorithm type: {actual_type}. Expected Algorithm, str, dict, or None.",
+            type_error_fmt="Algorithm factory returned {type_name}, which is not a Algorithm subclass.",
        )

    def _make_adapter(self, adapter: ComponentSpec[TraceAdapter[Any]]) -> TraceAdapter[Any]:
+        """Resolve the adapter used to transform spans into algorithm-ready payloads."""
        return build_component(
            adapter,
            expected_type=TraceAdapter,
@@ -220,6 +283,7 @@ class Trainer(TrainerLegacy):
        )

    def _make_store(self, store: ComponentSpec[LightningStore]) -> LightningStore:
+        """Resolve the store implementation backing rollouts, attempts, spans, and resources."""
        return build_component(
            store,
            expected_type=LightningStore,
@@ -234,13 +298,21 @@ class Trainer(TrainerLegacy):
        strategy: ComponentSpec[ExecutionStrategy],
        *,
        n_runners: int,
+        port: Optional[int] = None,
    ) -> ExecutionStrategy:
+        """Resolve the execution strategy and seed defaults such as `n_runners`."""
        if isinstance(strategy, ExecutionStrategy):
+            if port is not None and isinstance(strategy, ClientServerExecutionStrategy):
+                strategy.server_port = port
            return strategy
        optional_defaults: Dict[str, Callable[[], Any]] = {"n_runners": lambda: n_runners}
+        if port is not None:
+            optional_defaults["server_port"] = lambda: port

        def default_factory() -> ExecutionStrategy:
-            return ClientServerExecutionStrategy(n_runners=n_runners, role="both")
+            if port is not None:
+                return ClientServerExecutionStrategy(n_runners=n_runners, server_port=port)
+            return ClientServerExecutionStrategy(n_runners=n_runners)

        return build_component(
            strategy,
@@ -259,6 +331,7 @@ class Trainer(TrainerLegacy):
        *,
        store: LightningStore,
    ) -> Optional[LLMProxy]:
+        """Resolve an optional LLM proxy and ensure it shares the trainer's store instance."""
        if isinstance(llm_proxy, LLMProxy):
            return llm_proxy

@@ -277,25 +350,27 @@ class Trainer(TrainerLegacy):
            type_error_fmt="llm_proxy factory returned {type_name}, which is not an LLMProxy subclass.",
        )

-    def _make_runner(self, runner: ComponentSpec[BaseRunner[Any]]) -> BaseRunner[Any]:
+    def _make_runner(self, runner: ComponentSpec[Runner[Any]]) -> Runner[Any]:
+        """Resolve the runner responsible for executing the agent inside each worker."""
        optional_defaults: Dict[str, Callable[[], Any]] = {"tracer": lambda: self.tracer}
        if self.max_rollouts is not None:
            optional_defaults["max_rollouts"] = lambda: self.max_rollouts

-        def default_runner_factory() -> BaseRunner[Any]:
+        def default_runner_factory() -> Runner[Any]:
            return instantiate_component(LitAgentRunner, optional_defaults=optional_defaults)

        return build_component(
            runner,
-            expected_type=BaseRunner,
+            expected_type=Runner,
            spec_name="runner",
            default_factory=default_runner_factory,
            optional_defaults=optional_defaults,
-            invalid_spec_error_fmt="Invalid runner type: {actual_type}. Expected BaseRunner, callable, str, dict, or None.",
-            type_error_fmt="Runner factory returned {type_name}, which is not a BaseRunner subclass.",
+            invalid_spec_error_fmt="Invalid runner type: {actual_type}. Expected Runner, callable, str, dict, or None.",
+            type_error_fmt="Runner factory returned {type_name}, which is not a Runner subclass.",
        )

    def _normalize_hooks(self, hooks: Optional[Union[Hook, Sequence[Hook]]]) -> Sequence[Hook]:
+        """Coerce hook inputs into an immutable sequence for runner initialization."""
        if hooks is None:
            return ()
        if isinstance(hooks, Hook):
@@ -309,13 +384,33 @@ class Trainer(TrainerLegacy):
        *,
        val_dataset: Optional[Dataset[T_co]] = None,
    ) -> None:
-        """Run the training loop using the configured strategy, store, and runner.
+        """Execute the full algorithm/runner training loop.
+
+        [`Trainer.fit`][agentlightning.Trainer.fit] packages the algorithm and runner bundles,
+        then hands them to the active [`ExecutionStrategy`][agentlightning.ExecutionStrategy].
+        The strategy rarely returns until:
+
+        * The algorithm exhausts the dataset(s) and stops enqueuing rollouts.
+        * `max_rollouts` causes individual runners to exit.
+        * An exception or interrupt cancels the shared [`ExecutionEvent`][agentlightning.ExecutionEvent].

        Args:
-            agent: The LitAgent instance to be trained on.
-            train_dataset: The dataset to train on.
-            val_dataset: The dataset to validate on.
+            agent: [`LitAgent`][agentlightning.LitAgent] implementation executed by runners.
+            train_dataset: Optional iterable of rollout inputs consumed by the algorithm.
+            val_dataset: Optional iterable consumed by validation passes.
        """
+        if isinstance(train_dataset, str):
+            logger.warning(
+                "Trainer.fit will no longer accepts a string URL in future version. "
+                "To continue using a string URL, please use Trainer.fit_v0 instead. "
+                "See documentation for how to migrate to latest version: https://microsoft.github.io/agent-lightning/stable/"
+            )
+            return self.fit_v0(  # type: ignore
+                agent,
+                train_dataset,
+                val_dataset,  # type: ignore
+            )
+
        agent.set_trainer(self)

        algorithm_bundle = functools.partial(
@@ -335,15 +430,22 @@ class Trainer(TrainerLegacy):
        *,
        val_dataset: Optional[Dataset[T_co]] = None,
    ) -> None:
-        """Dry run the training loop with a FastAlgorithm and the real runner.
+        """Exercise the infrastructure using a fast, synchronous algorithm.
+
+        [`Trainer.dev`][agentlightning.Trainer.dev] mirrors [`fit()`][agentlightning.Trainer.fit] but
+        insists on an [`Algorithm`][agentlightning.Algorithm] subtype that also derives from
+        [`FastAlgorithm`][agentlightning.FastAlgorithm]. This keeps the loop responsive for
+        debugging while still touching the same store, runners, hooks, and tracer plumbing.
+
+        If no algorithm is provided, a default [`Baseline`][agentlightning.Baseline] algorithm will be used.

        Args:
-            agent: The LitAgent instance to be trained on.
-            train_dataset: The dataset to train on.
-            val_dataset: The dataset to validate on.
+            agent: [`LitAgent`][agentlightning.LitAgent] implementation to execute.
+            train_dataset: Optional iterable passed to the algorithm.
+            val_dataset: Optional iterable passed to the algorithm.

        Raises:
-            TypeError: If the configured algorithm is not a :class:`FastAlgorithm`.
+            TypeError: If the configured algorithm does not inherit from `FastAlgorithm`.
        """
        agent.set_trainer(self)

@@ -374,8 +476,17 @@ class Trainer(TrainerLegacy):
        event: ExecutionEvent,
        train_dataset: Optional[Dataset[T_co]],
        val_dataset: Optional[Dataset[T_co]],
-        algorithm: Optional[BaseAlgorithm],
+        algorithm: Optional[Algorithm],
    ) -> None:
+        """Internal entry point executed by the strategy for the algorithm role.
+
+        This coroutine is scheduled inside the strategy's process/thread and is responsible
+        for binding algorithm dependencies (store, adapter, initial resources, proxy) before
+        invoking [`Algorithm.run`][agentlightning.Algorithm.run].
+        When `algorithm` is `None` the bundle simply waits for the
+        shared `event` to signal shutdown so runners can still execute (useful for manual queue
+        seeding or external algorithms).
+        """
        if algorithm is not None:
            algorithm.set_trainer(self)
            algorithm.set_store(store)
@@ -410,7 +521,14 @@ class Trainer(TrainerLegacy):
    async def _runner_bundle(
        self, store: LightningStore, worker_id: int, event: ExecutionEvent, agent: LitAgent[T_co]
    ) -> None:
-        runner_instance: BaseRunner[Any] | None = None
+        """Internal entry point executed by the strategy for each runner role.
+
+        The bundle materializes the configured runner, binds the agent and hooks, associates
+        the worker with the shared store, and then drives the runner's [`iter`][agentlightning.Runner.iter]
+        loop until the execution event is set or an exception occurs. Cleanup mirrors the initialization
+        sequence to keep tracer state, hooks, and agent resources consistent across restarts.
+        """
+        runner_instance: Runner[Any] | None = None
        runner_initialized = False
        worker_initialized = False
        try:
@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.

+"""Core data models shared across Agent Lightning components."""
+
 from __future__ import annotations

 from typing import (
@@ -25,8 +27,8 @@ from .tracer import Span

 if TYPE_CHECKING:
    from agentlightning.litagent import LitAgent
-    from agentlightning.runner.base import BaseRunner
-    from agentlightning.tracer.base import BaseTracer
+    from agentlightning.runner.base import Runner
+    from agentlightning.tracer.base import Tracer

 __all__ = [
    "Triplet",
@@ -53,7 +55,7 @@ T_co = TypeVar("T_co", covariant=True)


 class Triplet(BaseModel):
-    """A standard structure for a single turn in a trajectory."""
+    """Single interaction turn captured during reinforcement learning."""

    prompt: Any
    response: Any
@@ -62,7 +64,11 @@ class Triplet(BaseModel):


 class RolloutLegacy(BaseModel):
-    """The standard reporting object from client to server."""
+    """Legacy reporting payload exchanged with the deprecated HTTP server.
+
+    !!! warning "Deprecated"
+        Use [`Rollout`][agentlightning.Rollout] instead.
+    """

    rollout_id: str

@@ -97,6 +103,7 @@ RolloutStatus = Literal[
    "cancelled",  # cancelled by user (or watchdog)
    "requeuing",  # retrying
 ]
+"""The status of a rollout."""

 AttemptStatus = Literal[
    # A status is essentially a process.
@@ -108,66 +115,83 @@ AttemptStatus = Literal[
    "unresponsive",  # the worker has not reported results for a while
    "timeout",  # the worker has been emitting new logs, but have been working on the task for too long
 ]
+"""The status of an attempt."""

 RolloutMode = Literal["train", "val", "test"]
+"""Possible rollout modes."""


 class Attempt(BaseModel):
-    """An attempt to execute a rollout. A rollout can have multiple attempts if retries are needed."""
-
-    rollout_id: str  # the rollout this attempt belongs to
-    attempt_id: str  # the universal id for current attempt
-    sequence_id: int  # the sequence number of the attempt, starting from 1
-    start_time: float  # time when the attempt has started
-    end_time: Optional[float] = None  # time when the attempt has ended
+    """Execution attempt for a rollout, including metadata for retries."""

+    rollout_id: str
+    """The rollout which this attempt belongs to."""
+    attempt_id: str
+    """The universal id for current attempt."""
+    sequence_id: int
+    """The sequence number of the attempt, starting from 1."""
+    start_time: float
+    """The time when the attempt has started."""
+    end_time: Optional[float] = None
+    """The time when the attempt has ended."""
    status: AttemptStatus = "preparing"
-    # The rollout worker which is executing this attempt
+    """The status of the attempt."""
    worker_id: Optional[str] = None
+    """The rollout worker which is executing this attempt."""

-    last_heartbeat_time: Optional[float] = None  # last time when the worker has reported progress
+    last_heartbeat_time: Optional[float] = None
+    """The last time when the worker has reported progress (i.e., a span)."""

-    # A bucket for any other relevant information
    metadata: Optional[Dict[str, Any]] = None
+    """A bucket for any other relevant information."""


 class RolloutConfig(BaseModel):
-    """Configurations for rollout execution."""
+    """Configuration controlling rollout retries and timeouts."""

-    timeout_seconds: Optional[float] = None  # none indicates no timeout
-    unresponsive_seconds: Optional[float] = None  # none indicates no unresponsive timeout
-    max_attempts: int = Field(default=1, ge=1)  # including the first attempt
-    retry_condition: List[AttemptStatus] = Field(
-        default_factory=cast(Callable[[], List[AttemptStatus]], list)
-    )  # list of statuses that should trigger a retry
+    timeout_seconds: Optional[float] = None
+    """The timeout for the rollout, in seconds. None indicates no timeout."""
+    unresponsive_seconds: Optional[float] = None
+    """The unresponsive timeout for the rollout, in seconds. None indicates no unresponsive timeout."""
+    max_attempts: int = Field(default=1, ge=1)
+    """The maximum number of attempts for the rollout, including the first attempt."""
+    retry_condition: List[AttemptStatus] = Field(default_factory=cast(Callable[[], List[AttemptStatus]], list))
+    """The list of statuses that should trigger a retry."""


 class Rollout(BaseModel):
    rollout_id: str
+    """Unique identifier for the rollout."""

-    # Inputs
    input: TaskInput
+    """Task input used to generate the rollout."""

    # Time to track the lifecycle of the rollout
    start_time: float
+    """Timestamp when the rollout started."""
    end_time: Optional[float] = None
+    """Timestamp when the rollout ended."""

    mode: Optional[RolloutMode] = None
+    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][agentlightning.RolloutMode]."""
    resources_id: Optional[str] = None
+    """Identifier of the resources required to execute the rollout."""

-    # Overall scheduling/running information
    status: RolloutStatus = "queuing"
+    """Latest status emitted by the controller."""

    config: RolloutConfig = Field(default_factory=RolloutConfig)
+    """Retry and timeout configuration associated with the rollout."""

-    # A bucket for any other relevant information
    metadata: Optional[Dict[str, Any]] = None
+    """Additional metadata attached to the rollout."""


 class AttemptedRollout(Rollout):
-    """A rollout along with its active attempt."""
+    """Rollout paired with the currently active attempt."""

    attempt: Attempt
+    """The attempt that is currently processing the rollout."""

    @model_validator(mode="after")
    def check_consistency(self) -> AttemptedRollout:
@@ -177,11 +201,16 @@ class AttemptedRollout(Rollout):


 TaskInput = Any
-"""Task input type. Can be any type."""
+"""Task input type. Accepts arbitrary payloads."""


 class Task(BaseModel):
-    """A task (rollout request) to be processed by the client agent. Deprecated."""
+    """Rollout request served to client agents.
+
+    !!! warning "Deprecated"
+        The legacy HTTP client/server stack still uses this model. Prefer
+        [`LightningStore`][agentlightning.LightningStore] APIs for new workflows.
+    """

    rollout_id: str
    input: TaskInput
@@ -199,11 +228,23 @@ class Task(BaseModel):


 class TaskIfAny(BaseModel):
+    """A task or indication that no task is available.
+
+    !!! warning "Deprecated"
+        Use [`LightningStore`][agentlightning.LightningStore] APIs for new workflows.
+    """
+
    is_available: bool
+    """Indication that a task is available."""
    task: Optional[Task] = None


 RolloutRawResultLegacy = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], RolloutLegacy]
+"""Legacy rollout result type.
+
+!!! warning "Deprecated"
+    Use [`RolloutRawResult`][agentlightning.RolloutRawResult] instead.
+"""

 RolloutRawResult = Union[
    None,  # nothing (relies on tracer)
@@ -211,11 +252,23 @@ RolloutRawResult = Union[
    List[ReadableSpan],  # constructed OTEL spans by user
    List[Span],  # constructed Span objects by user
 ]
+"""Rollout result type.
+
+Possible return values of [`rollout`][agentlightning.LitAgent.rollout].
+"""


 class GenericResponse(BaseModel):
-    """
-    A generic response message that can be used for various purposes.
+    """Generic server response used by compatibility endpoints.
+
+    !!! warning "Deprecated"
+        This response is no longer used by the new
+        [`LightningStore`][agentlightning.LightningStore] APIs.
+
+    Attributes:
+        status: Status string describing the result of the request.
+        message: Optional human readable explanation.
+        data: Arbitrary payload serialized as JSON.
    """

    status: str = "success"
@@ -224,19 +277,18 @@ class GenericResponse(BaseModel):


 class ParallelWorkerBase:
-    """Base class for objects that can be parallelized across multiple worker processes.
+    """Base class for workloads executed across multiple worker processes.

-    This class defines the standard lifecycle for parallel processing:
+    The lifecycle is orchestrated by the main process:

-    Main Process:
-        1. init() - Initialize the object in the main process
-        2. spawn workers and call init_worker() in each worker
-        3. run() - Execute the main workload in parallel across workers
-        4. teardown_worker() - Clean up resources in each worker
-        5. teardown() - Final cleanup in the main process
+    * [`init()`][agentlightning.ParallelWorkerBase.init] prepares shared state.
+    * Each worker calls [`init_worker()`][agentlightning.ParallelWorkerBase.init_worker] during start-up.
+    * [`run()`][agentlightning.ParallelWorkerBase.run] performs the parallel workload.
+    * Workers call [`teardown_worker()`][agentlightning.ParallelWorkerBase.teardown_worker] before exiting.
+    * The main process finalizes through [`teardown()`][agentlightning.ParallelWorkerBase.teardown].

-    Subclasses should implement the run() method and optionally override
-    the lifecycle methods for custom initialization and cleanup behavior.
+    Subclasses must implement [`run()`][agentlightning.ParallelWorkerBase.run]
+    and can override other lifecycle hooks.
    """

    def __init__(self) -> None:
@@ -244,25 +296,30 @@ class ParallelWorkerBase:
        self.worker_id: Optional[int] = None

    def init(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize before spawning the workers. This method can be overridden by subclasses."""
        pass

    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Initialize the worker. This method can be overridden by subclasses."""
        self.worker_id = worker_id

    def run(self, *args: Any, **kwargs: Any) -> Any:
+        """Run the workload. This method can be overridden by subclasses."""
        pass

    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Teardown the worker. This method can be overridden by subclasses."""
        pass

    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Teardown after the workers have exited. This method can be overridden by subclasses."""
        pass


 class Dataset(Protocol, Generic[T_co]):
    """The general interface for a dataset.

-    It's currently implemented as a protocol, having a similar interface to torch.utils.data.Dataset.
+    It's currently implemented as a protocol, having a similar interface to `torch.utils.data.Dataset`.
    You don't have to inherit from this class; you can use a simple list if you want to.
    """

@@ -275,42 +332,42 @@ class Hook(ParallelWorkerBase):
    """Base class for defining hooks in the agent runner's lifecycle."""

    async def on_trace_start(
-        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: Rollout
+        self, *, agent: LitAgent[Any], runner: Runner[Any], tracer: Tracer, rollout: Rollout
    ) -> None:
        """Hook called immediately after the tracer enters the trace context but before the rollout begins.

        Args:
-            agent: The :class:`LitAgent` instance associated with the runner.
-            runner: The :class:`BaseRunner` managing the rollout.
-            tracer: The :class:`BaseTracer` instance associated with the runner.
-            rollout: The :class:`Rollout` object that will be processed.
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: The [`Tracer`][agentlightning.Tracer] instance associated with the runner.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that will be processed.

        Subclasses can override this method to implement custom logic such as logging,
        metric collection, or resource setup. By default, this is a no-op.
        """

    async def on_trace_end(
-        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: Rollout
+        self, *, agent: LitAgent[Any], runner: Runner[Any], tracer: Tracer, rollout: Rollout
    ) -> None:
        """Hook called immediately after the rollout completes but before the tracer exits the trace context.

        Args:
-            agent: The :class:`LitAgent` instance associated with the runner.
-            runner: The :class:`BaseRunner` managing the rollout.
-            tracer: The :class:`BaseTracer` instance associated with the runner.
-            rollout: The :class:`Rollout` object that has been processed.
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: The [`Tracer`][agentlightning.Tracer] instance associated with the runner.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that has been processed.

        Subclasses can override this method to implement custom logic such as logging,
        metric collection, or resource cleanup. By default, this is a no-op.
        """

-    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: BaseRunner[Any], rollout: Rollout) -> None:
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: Runner[Any], rollout: Rollout) -> None:
        """Hook called immediately before a rollout *attempt* begins.

        Args:
-            agent: The :class:`LitAgent` instance associated with the runner.
-            runner: The :class:`BaseRunner` managing the rollout.
-            rollout: The :class:`Rollout` object that will be processed.
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that will be processed.

        Subclasses can override this method to implement custom logic such as
        logging, metric collection, or resource setup. By default, this is a
@@ -321,16 +378,16 @@ class Hook(ParallelWorkerBase):
        self,
        *,
        agent: LitAgent[Any],
-        runner: BaseRunner[Any],
+        runner: Runner[Any],
        rollout: Rollout,
        spans: Union[List[ReadableSpan], List[Span]],
    ) -> None:
        """Hook called after a rollout *attempt* completes.

        Args:
-            agent: The :class:`LitAgent` instance associated with the runner.
-            runner: The :class:`BaseRunner` managing the rollout.
-            rollout: The :class:`Rollout` object that has been processed.
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that has been processed.
            spans: The spans that have been added to the store.

        Subclasses can override this method for cleanup or additional
@@ -2,6 +2,8 @@

 from __future__ import annotations

+"""Typed representations of tunable resources shared between Agent Lightning components."""
+
 import inspect
 import logging
 from typing import (
@@ -32,40 +34,40 @@ __all__ = [


 class Resource(BaseModel):
-    """
-    Base class for all tunable resources.
-    """
+    """Base class for tunable resources distributed to executors."""

    resource_type: Any
+    """Alias of the resource type."""


 class LLM(Resource):
-    """
-    Provide an LLM endpoint and model name as a resource.
-
-    Attributes:
-        endpoint (str): The URL of the LLM API endpoint.
-        model (str): The identifier for the model to be used (e.g., 'gpt-4o').
-        sampling_parameters (SamplingParameters): A dictionary of hyperparameters
-            for model inference, such as temperature, top_p, etc.
-    """
+    """Resource that identifies an LLM endpoint and its configuration."""

    resource_type: Literal["llm"] = "llm"
    endpoint: str
+    """The URL of the LLM API endpoint."""
    model: str
+    """The identifier for the model to be used (e.g., 'gpt-4o')."""
    api_key: Optional[str] = None
+    """Optional secret used to authenticate requests."""
    sampling_parameters: Dict[str, Any] = Field(default_factory=dict)
+    """A dictionary of hyperparameters for model inference, such as temperature, top_p, etc."""

    def get_base_url(self, *args: Any, **kwargs: Any) -> str:
-        """The base_url to put into openai.OpenAI.
+        """Return the base URL consumed by OpenAI-compatible clients.

-        Users are encouraged to use `base_url` to get the LLM endpoint instead of accessing `endpoint` directly.
+        Users are encouraged to use `get_base_url(rollout_id, attempt_id)` to get
+        the LLM endpoint instead of accessing `.endpoint` directly.
        """
        return self.endpoint


 class ProxyLLM(LLM):
-    """Proxy LLM resource that is tailored by `llm_proxy.LLMProxy`."""
+    """LLM resource that rewrites endpoints through [`LLMProxy`][agentlightning.LLMProxy].
+
+    The proxy injects rollout- and attempt-specific routing information into the
+    endpoint so that downstream services can attribute requests correctly.
+    """

    resource_type: Literal["proxy_llm"] = "proxy_llm"  # type: ignore
    _initialized: bool = False
@@ -76,7 +78,7 @@ class ProxyLLM(LLM):
        object.__setattr__(self, "_initialized", True)

    def __getattribute__(self, name: str) -> Any:
-        """Override to emit a warning when endpoint is accessed directly."""
+        """Emit a warning when `endpoint` is accessed directly after initialization."""
        # Check if we're accessing endpoint after initialization and not from base_url
        if name == "endpoint":
            try:
@@ -97,7 +99,7 @@ class ProxyLLM(LLM):
        return super().__getattribute__(name)

    def with_attempted_rollout(self, rollout: AttemptedRollout) -> LLM:
-        """Bake the rollout and attempt id into the endpoint."""
+        """Bake rollout metadata into a concrete [`LLM`][agentlightning.LLM] instance."""
        return LLM(
            endpoint=self.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
            model=self.model,
@@ -106,6 +108,18 @@ class ProxyLLM(LLM):
        )

    def get_base_url(self, rollout_id: Optional[str], attempt_id: Optional[str]) -> str:
+        """Return the routed endpoint for a specific rollout/attempt pair.
+
+        Args:
+            rollout_id: Identifier of the rollout making the request.
+            attempt_id: Identifier of the attempt within that rollout.
+
+        Returns:
+            Fully qualified endpoint including rollout metadata.
+
+        Raises:
+            ValueError: If exactly one of ``rollout_id`` or ``attempt_id`` is provided.
+        """
        if rollout_id is None and attempt_id is None:
            return self.endpoint

@@ -130,22 +144,20 @@ class ProxyLLM(LLM):


 class PromptTemplate(Resource):
-    """
-    A prompt template as a resource.
-
-    Attributes:
-        template (str): The template string. The format depends on the engine.
-        engine (Literal['jinja', 'f-string', 'poml']): The templating engine
-            to use for rendering the prompt. I imagine users can use their own
-            customized engines, but algos can only well operate on a subset of them.
-    """
+    """Resource describing a reusable prompt template."""

    resource_type: Literal["prompt_template"] = "prompt_template"
    template: str
+    """The template string. The format depends on the engine."""
    engine: Literal["jinja", "f-string", "poml"]
+    """The templating engine to use for rendering the prompt."""

    def format(self, **kwargs: Any) -> str:
-        """Format the prompt template with the given kwargs."""
+        """Format the prompt using keyword arguments.
+
+        !!! warning
+            Only the `f-string` engine is supported for now.
+        """
        if self.engine == "f-string":
            return self.template.format(**kwargs)
        else:
@@ -158,32 +170,29 @@ class PromptTemplate(Resource):
 # TODO: migrate to use a registry
 ResourceUnion = Annotated[Union[LLM, ProxyLLM, PromptTemplate], Field(discriminator="resource_type")]
 NamedResources = Dict[str, ResourceUnion]
-"""
-A dictionary-like class to hold named resources.
+"""Mapping from resource names to their configured instances.

-Example:
+Examples:
+    ```python
    resources: NamedResources = {
-        'main_llm': LLM(
+        "main_llm": LLM(
            endpoint="http://localhost:8080",
            model="llama3",
-            sampling_parameters={'temperature': 0.7, 'max_tokens': 100}
+            sampling_parameters={"temperature": 0.7, "max_tokens": 100},
        ),
-        'system_prompt': PromptTemplate(
+        "system_prompt": PromptTemplate(
            template="You are a helpful assistant.",
-            engine='f-string'
-        )
+            engine="f-string",
+        ),
    }
+    ```
 """


 class ResourcesUpdate(BaseModel):
-    """
-    A resource update message to be sent from the server to clients.
-
-    This message contains a dictionary of resources that clients should use
-    for subsequent tasks. It is used to update the resources available to
-    clients dynamically.
-    """
+    """Update payload broadcast to clients when resources change."""

    resources_id: str
+    """Identifier used to version the resources."""
    resources: NamedResources
+    """Mapping of resource names to their definitions."""
@@ -2,17 +2,19 @@

 from __future__ import annotations

+"""Data models that mirror OpenTelemetry spans for Agent Lightning."""
+
 import json
 from enum import Enum
 from typing import Any, Dict, List, Optional, Sequence, Union

 from opentelemetry import trace as trace_api
-from opentelemetry.sdk.resources import Resource as OtelResource
+from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import Event as OtelEvent
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 from opentelemetry.trace.status import Status as OtelStatus
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict

 __all__ = [
    "AttributeValue",
@@ -22,7 +24,7 @@ __all__ = [
    "TraceStatus",
    "Event",
    "Link",
-    "Resource",
+    "OtelResource",
    "Span",
    "SpanNames",
    "SpanAttributeNames",
@@ -31,9 +33,13 @@ __all__ = [


 def convert_timestamp(timestamp: Optional[int]) -> Optional[float]:
-    """Convert timestamp from nanoseconds to seconds if needed.
+    """Normalize OpenTelemetry timestamps to seconds.

-    Auto-detects format: if > 1e12, assumes nanoseconds; otherwise seconds.
+    Args:
+        timestamp: Timestamp expressed either in seconds or nanoseconds.
+
+    Returns:
+        Timestamp in seconds when `timestamp` is provided; otherwise `None`.
    """
    if not timestamp:
        return None
@@ -41,7 +47,15 @@ def convert_timestamp(timestamp: Optional[int]) -> Optional[float]:


 def extract_extra_fields(src: Any, excluded_fields: List[str]) -> Dict[str, Any]:
-    """Extract extra fields from source object, excluding specified fields and private fields."""
+    """Capture custom attributes from an OpenTelemetry object.
+
+    Args:
+        src: Object that exposes a `__dict__` of potential attributes.
+        excluded_fields: Attribute names that should be removed from the output.
+
+    Returns:
+        Dictionary containing JSON-serializable representations of the remaining fields.
+    """
    excluded_fields_set = set(excluded_fields) | set(["_" + k for k in excluded_fields])
    # Exclude the function fields
    excluded_fields_set |= set(src.__class__.__dict__.keys())
@@ -62,23 +76,31 @@ AttributeValue = Union[
    Sequence[int],
    Sequence[float],
 ]
+"""Possible values for OpenTelemetry attributes."""
 Attributes = Dict[str, AttributeValue]
+"""Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
 TraceState = Dict[str, str]
+"""Mapping from trace state key to its value. Same as OpenTelemetry `TraceState` type."""


 class SpanContext(BaseModel):
-    """Corresponding to opentelemetry.trace.SpanContext"""
+    """Pydantic representation of `opentelemetry.trace.SpanContext` values."""

    trace_id: str
+    """The trace ID of the span."""
    span_id: str
+    """The span ID of the span."""
    is_remote: bool
+    """Whether the span is remote."""
    trace_state: TraceState
+    """Mapping from trace state key to its value."""

-    class Config:
-        allow_extra = True
+    model_config = ConfigDict(extra="allow")

    @classmethod
    def from_opentelemetry(cls, src: trace_api.SpanContext) -> "SpanContext":
+        """Construct a [`SpanContext`][agentlightning.SpanContext] from OpenTelemetry data."""
+
        return cls(
            trace_id=trace_api.format_trace_id(src.trace_id),
            span_id=trace_api.format_span_id(src.span_id),
@@ -89,16 +111,19 @@ class SpanContext(BaseModel):


 class TraceStatus(BaseModel):
-    """Corresponding to opentelemetry.trace.Status"""
+    """Serializable variant of `opentelemetry.trace.Status`."""

    status_code: str
+    """The status code of the span. Same as OpenTelemetry `Status.status_code` type."""
    description: Optional[str] = None
+    """The description of the span. Same as OpenTelemetry `Status.description` type."""

-    class Config:
-        allow_extra = True
+    model_config = ConfigDict(extra="allow")

    @classmethod
    def from_opentelemetry(cls, src: OtelStatus) -> "TraceStatus":
+        """Create a [`TraceStatus`][agentlightning.TraceStatus] from OpenTelemetry metadata."""
+
        return cls(
            status_code=src.status_code.name,
            description=src.description,
@@ -107,17 +132,21 @@ class TraceStatus(BaseModel):


 class Event(BaseModel):
-    """Corresponding to opentelemetry.trace.Event"""
+    """Serializable representation of OpenTelemetry `Event` values."""

    name: str
+    """The name of the event."""
    attributes: Attributes
+    """Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
    timestamp: Optional[float] = None
+    """The timestamp of the event. Same as OpenTelemetry `Event.timestamp` type."""

-    class Config:
-        allow_extra = True
+    model_config = ConfigDict(extra="allow")

    @classmethod
    def from_opentelemetry(cls, src: OtelEvent) -> "Event":
+        """Create an [`Event`][agentlightning.Event] from an OpenTelemetry event."""
+
        return cls(
            name=src.name,
            attributes=dict(src.attributes) if src.attributes else {},
@@ -127,16 +156,19 @@ class Event(BaseModel):


 class Link(BaseModel):
-    """Corresponding to opentelemetry.trace.Link"""
+    """Serializable representation of OpenTelemetry `Link` values."""

    context: SpanContext
+    """The context of the link."""
    attributes: Optional[Attributes] = None
+    """Optional attributes."""

-    class Config:
-        allow_extra = True
+    model_config = ConfigDict(extra="allow")

    @classmethod
    def from_opentelemetry(cls, src: trace_api.Link) -> "Link":
+        """Create a [`Link`][agentlightning.Link] from an OpenTelemetry link."""
+
        return cls(
            context=SpanContext.from_opentelemetry(src.context),
            attributes=dict(src.attributes) if src.attributes else None,
@@ -144,14 +176,24 @@ class Link(BaseModel):
        )


-class Resource(BaseModel):
-    """Corresponding to opentelemetry.sdk.resources.Resource"""
+class OtelResource(BaseModel):
+    """Serializable representation of OpenTelemetry `Resource` values.
+
+    Named as `OtelResource` to avoid confusion with the [`Resource`][agentlightning.Resource] class.
+    Users will very rarely need to construct this class directly. Most of the times,
+    they deal with the [`Resource`][agentlightning.Resource] class instead, which describes
+    a very different concept.
+    """

    attributes: Attributes
+    """Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
    schema_url: str
+    """The schema URL of the resource."""

    @classmethod
-    def from_opentelemetry(cls, src: OtelResource) -> "Resource":
+    def from_opentelemetry(cls, src: Resource) -> "OtelResource":
+        """Create a [`Resource`][agentlightning.Resource] from an OpenTelemetry resource."""
+
        return cls(
            attributes=dict(src.attributes) if src.attributes else {},
            schema_url=src.schema_url if src.schema_url else "",
@@ -160,35 +202,58 @@ class Resource(BaseModel):


 class Span(BaseModel):
+    """Agent Lightning's canonical span model used for persistence and analytics.

-    class Config:
-        allow_extra = True  # allow extra fields if needed
+    The model captures the most relevant fields from
+    `opentelemetry.sdk.trace.ReadableSpan` instances while preserving unmodeled
+    attributes in Pydantic `BaseModel`'s extra storage. This keeps the serialized format
+    stable even as upstream OpenTelemetry types evolve.
+    """
+
+    model_config = ConfigDict(extra="allow")

    rollout_id: str
+    """The rollout which this span belongs to."""
    attempt_id: str
-    # The ID to make spans ordered within a single attempt
+    """The attempt which this span belongs to."""
    sequence_id: int
+    """The ID to make spans ordered within a single attempt."""

    # Current ID (in hex, formatted via trace_api.format_*)
    trace_id: str  # one rollout can have traces coming from multiple places
+    """The trace ID of the span. One rollout/attempt can have multiple traces.
+    This ID comes from the OpenTelemetry trace ID generator.
+    """
    span_id: str
+    """The span ID of the span. This ID comes from the OpenTelemetry span ID generator."""
    parent_id: Optional[str]
+    """The parent span ID of the span."""

    # Core ReadableSpan fields
    name: str
+    """The name of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
    status: TraceStatus
+    """The status of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
    attributes: Attributes
+    """The attributes of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
    events: List[Event]
+    """The events of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
    links: List[Link]
+    """The links of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""

    # Timestamps
    start_time: Optional[float]
+    """The start time of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
    end_time: Optional[float]
+    """The end time of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""

    # Other parsable fields
    context: Optional[SpanContext]
+    """The context of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
    parent: Optional[SpanContext]
-    resource: Resource
+    """The parent context of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    resource: OtelResource
+    """The resource of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""

    # Preserve other fields in the readable span as extra fields
    # Make sure that are json serializable (so no bytes, complex objects, ...)
@@ -201,6 +266,17 @@ class Span(BaseModel):
        attempt_id: str,
        sequence_id: int,
    ) -> "Span":
+        """Convert an OpenTelemetry span into the Agent Lightning data model.
+
+        Args:
+            src: Span captured by OpenTelemetry.
+            rollout_id: Identifier for the rollout that produced the span.
+            attempt_id: Identifier of the attempt within the rollout.
+            sequence_id: Monotonically increasing identifier assigned to the span.
+
+        Returns:
+            Parsed [`Span`][agentlightning.Span] instance suitable for persistence.
+        """
        context = src.get_span_context()
        if context is None:
            trace_id = span_id = 0
@@ -223,7 +299,7 @@ class Span(BaseModel):
            end_time=convert_timestamp(src.end_time),
            context=SpanContext.from_opentelemetry(context) if context else None,
            parent=(SpanContext.from_opentelemetry(src.parent) if src.parent else None),
-            resource=Resource.from_opentelemetry(src.resource),
+            resource=OtelResource.from_opentelemetry(src.resource),
            **extract_extra_fields(
                src,
                [
@@ -261,8 +337,28 @@ class Span(BaseModel):
        parent_id: Optional[str] = None,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
-        resource: Optional[Resource] = None,
+        resource: Optional[OtelResource] = None,
    ) -> "Span":
+        """Build a synthetic span from raw attributes.
+        Different from the [`from_opentelemetry`][agentlightning.Span.from_opentelemetry] method,
+        all parameters other than `attributes` are optional and will be generated if not provided.
+
+        Args:
+            attributes: Span attributes to persist.
+            rollout_id: Optional rollout identifier associated with the span.
+            attempt_id: Optional attempt identifier associated with the span.
+            sequence_id: Optional sequence number to preserve ordering.
+            name: Optional human-readable span name.
+            trace_id: Custom trace identifier. When omitted, a random identifier is generated.
+            span_id: Custom span identifier. When omitted, a random identifier is generated.
+            parent_id: Optional parent span identifier.
+            start_time: Span start timestamp in seconds.
+            end_time: Span end timestamp in seconds.
+            resource: Explicit resource information to attach to the span.
+
+        Returns:
+            [`Span`][agentlightning.Span] populated with the provided attributes.
+        """

        id_generator = RandomIdGenerator()
        trace_id = trace_id or trace_api.format_trace_id(id_generator.generate_trace_id())
@@ -284,7 +380,7 @@ class Span(BaseModel):
                trace_state={},
            ),
            name=name or SpanNames.VIRTUAL.value,
-            resource=resource or Resource(attributes={}, schema_url=""),
+            resource=resource or OtelResource(attributes={}, schema_url=""),
            attributes=attributes,
            status=TraceStatus(status_code="OK"),
            events=[],
@@ -303,24 +399,28 @@ class Span(BaseModel):


 class SpanNames(str, Enum):
-    """Standard span name values for AgentLightning.
-
-    Currently reward, message, object and exception spans are supported.
-    We will add more spans related to error handling in the future.
-    """
+    """Enumerated span names recognised by Agent-lightning."""

    REWARD = "agentlightning.reward"
+    """The name of the reward span."""
    MESSAGE = "agentlightning.message"
+    """The name of the message span."""
    OBJECT = "agentlightning.object"
+    """The name of the object span."""
    EXCEPTION = "agentlightning.exception"
+    """The name of the exception span."""
    VIRTUAL = "agentlightning.virtual"
+    """The name of the virtual span. It represents derived spans without concrete operations."""


 class SpanAttributeNames(str, Enum):
-    """Standard attribute names for AgentLightning spans."""
+    """Canonical attribute names written by Agent Lightning emitters."""

    MESSAGE = "message"
+    """The name of the message attribute."""
    OBJECT = "object"
+    """The name of the object attribute."""


 SpanLike = Union[ReadableSpan, Span]
+"""Union type of OpenTelemetry `ReadableSpan` and Agent-lightning [`Span`][agentlightning.Span]."""
@@ -32,7 +32,7 @@ class PatchedvLLMServer(_unwrap_ray_remote(AsyncvLLMServer)):
    async def chat_completion(self, raw_request: Request):
        """OpenAI-compatible HTTP endpoint.

-        API reference: https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
+        API reference: [OpenAI-compatible server documentation](https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html)
        """
        request_json = await raw_request.json()
        request = ChatCompletionRequest(**request_json)
@@ -7,6 +7,7 @@ import socket
 import threading
 import time
 import uuid
+from collections import defaultdict
 from collections.abc import Mapping
 from typing import Any, Dict, List, Literal, Optional, Tuple

@@ -558,13 +559,28 @@ class AgentModeDaemon:
        assert len(self._completed_rollouts_v0) == self._total_tasks_queued

        sample_stat_list: List[Dict[str, Any]] = []
-        for _, rollout in self._completed_rollouts_v0.items():
+        sample_stat_list_by_source: Dict[str, List[Dict[str, Any]]] = defaultdict(
+            list
+        )  # FIXME: Evaluate whether grouping stats by source is actually needed.
+
+        for rollout_id, rollout in self._completed_rollouts_v0.items():
            final_reward = self._fillna_reward(rollout)
            if not rollout.triplets:
                print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
                sample_stat_list.append({"reward": final_reward})
                continue
            response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
+            if "data_source" in self._task_id_to_original_sample[rollout_id]:
+                # When a test sample includes a 'data_source' field, record per-source statistics for test results.
+                data_source = self._task_id_to_original_sample[rollout_id]["data_source"]
+                sample_stat_list_by_source[data_source].append(
+                    {
+                        "sum_response_length": np.sum(response_length_list),
+                        "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
+                        "turn_count": len(rollout.triplets),
+                        "reward": final_reward,
+                    }
+                )
            sample_stat_list.append(
                {
                    "sum_response_length": np.sum(response_length_list),
@@ -573,18 +589,45 @@ class AgentModeDaemon:
                    "reward": final_reward,
                }
            )
+        metric_dict: Dict[str, Any] = {}

        stats_w_trace = [stat for stat in sample_stat_list if "sum_response_length" in stat]
-        return {
-            "val/n_rollouts": len(sample_stat_list),
-            "val/n_rollouts_w_trace": len(stats_w_trace),
-            "val/reward": np.mean(
-                [stat["reward"] for stat in sample_stat_list]
-            ),  # each rollout must have a reward (fillna if missing)
-            "val/mean_response_length": np.mean([stat["mean_response_length"] for stat in stats_w_trace]),
-            "val/sum_response_length": np.mean([stat["sum_response_length"] for stat in stats_w_trace]),
-            "val/turn_count": np.mean([stat["turn_count"] for stat in stats_w_trace]),
+        stats_w_trace_by_source = {
+            data_source: [stat for stat in sample_stats if "sum_response_length" in stat]
+            for data_source, sample_stats in sample_stat_list_by_source.items()
        }
+        for data_source, sample_stats in sample_stat_list_by_source.items():
+            metric_dict.update(
+                {
+                    f"val/{data_source}/n_rollouts": len(sample_stats),
+                    f"val/{data_source}/n_rollouts_w_trace": len(stats_w_trace_by_source[data_source]),
+                    f"val/{data_source}/reward": np.mean(
+                        [stat["reward"] for stat in sample_stats]
+                    ),  # each rollout must have a reward (fillna if missing)
+                    f"val/{data_source}/mean_response_length": np.mean(
+                        [stat["mean_response_length"] for stat in stats_w_trace_by_source[data_source]]
+                    ),
+                    f"val/{data_source}/sum_response_length": np.mean(
+                        [stat["sum_response_length"] for stat in stats_w_trace_by_source[data_source]]
+                    ),
+                    f"val/{data_source}/turn_count": np.mean(
+                        [stat["turn_count"] for stat in stats_w_trace_by_source[data_source]]
+                    ),
+                }
+            )
+        metric_dict.update(
+            {
+                "val/n_rollouts": len(sample_stat_list),
+                "val/n_rollouts_w_trace": len(stats_w_trace),
+                "val/reward": np.mean(
+                    [stat["reward"] for stat in sample_stat_list]
+                ),  # each rollout must have a reward (fillna if missing)
+                "val/mean_response_length": np.mean([stat["mean_response_length"] for stat in stats_w_trace]),
+                "val/sum_response_length": np.mean([stat["sum_response_length"] for stat in stats_w_trace]),
+                "val/turn_count": np.mean([stat["turn_count"] for stat in stats_w_trace]),
+            }
+        )
+        return metric_dict

    def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, device: torch.device):
        """
@@ -63,6 +63,7 @@ class AgentLightningTrainer(RayPPOTrainer):
    RayPPOTrainer and focusing on the agent mode workflow.

    Key differences from RayPPOTrainer:
+
    1. Uses AgentModeDaemon for server communication
    2. Simplified data flow without pop/union operations
    3. Direct batch processing through agent daemon
@@ -16,9 +16,39 @@
 pip install agentlightning[apo]
 ```

+## Scope of Current Implementation
+
+APO is currently scoped to optimize a single prompt template. Optimizing multiple prompt templates is not supported yet.
+
+There is however no restriction on the number of variable placeholders in the prompt template (can range from zero to many). It's possible that invalid prompts are created during the optimization process. It is up to the agent developer to ensure that the prompt template is valid for the agent's task.
+
+## Initial Prompt
+
+APO expects the initial prompt to be provided in the `initial_resources` dictionary. This can be done in two approaches:
+
+1. Pass to the [Trainer][agentlightning.Trainer] constructor:
+
+```python
+trainer = agl.Trainer(
+    algorithm=agl.APO(...),
+    initial_resources={"main_prompt": agl.PromptTemplate(template="You are a helpful assistant.", engine="f-string")},
+)
+```
+
+2. Pass to the `[APO][agentlightning.algorithm.apo.APO].set_initial_resources()` method:
+
+```python
+algo = agl.APO(...)
+algo.set_initial_resources(
+    {"this_is_also_valid_key": agl.PromptTemplate(template="You are a helpful assistant.", engine="f-string")}
+)
+```
+
+The resource key can be arbitrary, which is used to identify the prompt template in [class-based implementations](../tutorials/write-agents.md) when you have multiple resources. When the key changes, the agent developer needs to update the key in the `rollout` method.
+
 ## Tutorials Using APO

-TBD
+- [Train the First Agent with APO](../how-to/train-first-agent.md) - A step-by-step guide to training your first agent using APO.

 ## References

@@ -6,5 +6,5 @@ For customizing algorithms, see [Algorithm-side References](../reference/algorit

 | Algorithm | Optimizing Resources | Description |
 | --------- | ------------------- | ----------- |
-| [APO](./apo.md) | [PromptTemplate][agentlightning.PromptTemplate] | Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search. |
-| [VERL](./verl.md) | [LLM][agentlightning.LLM] | Reinforcement Learning with [VERL framework](https://github.com/volcengine/verl). |
+| [APO](./apo.md) | `{<initial_prompt_key>: [PromptTemplate][agentlightning.PromptTemplate]}` | Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search. |
+| [VERL](./verl.md) | `{"main_llm": [LLM][agentlightning.LLM]}` | Reinforcement Learning with [VERL framework](https://github.com/volcengine/verl). |
@@ -10,10 +10,6 @@
    agl.VERL(...)
    ```

-!!! warning "Customization note"
-
-    Customization of VERL is not supported as of current version. We recommend copying the source code from VERL and modifying it as needed to suit your requirements.
-
 ## Installation

 ```bash
@@ -22,11 +18,36 @@ pip install agentlightning[verl]

 !!! warning

-    For best results, follow the steps in the [installation guide](../quickstart/installation.md) to set up VERL and its dependencies. Installing VERL directly with `pip install agentlightning[verl]` can cause issues unless you already have a compatible version of PyTorch installed.
+    To avoid various compatibility issues, follow the steps in the [installation guide](../tutorials/installation.md) to set up VERL and its dependencies. Installing VERL directly with `pip install agentlightning[verl]` can cause issues unless you already have a compatible version of PyTorch installed.
+
+!!! note "Notes for Readers"
+
+    [VERL][agentlightning.algorithm.verl.VERL] in this article refers to a wrapper, provided by Agent-lightning, of the [VERL framework](https://github.com/volcengine/verl). It's a subclass of [agentlightning.Algorithm][]. To differentiate it from the VERL framework, all references to the VERL framework shall use the term "VERL framework", and all references to the Agent-lightning wrapper shall be highlighted with a link.
+
+## Resources
+
+[VERL][agentlightning.algorithm.verl.VERL] expects no initial resources. The first LLM endpoint is directly deployed from the VERL configuration (`.actor_rollout_ref.model.path`). The resource key is always `main_llm`.
+
+[VERL][agentlightning.algorithm.verl.VERL] currently does not support optimizing multiple [LLM][agentlightning.LLM]s together.
+
+!!! note
+
+    The resource type created by [VERL][agentlightning.algorithm.verl.VERL] is actually a [ProxyLLM][agentlightning.ProxyLLM], a subclass of the [LLM][agentlightning.LLM] type. This object contains a **URL template** provided by [VERL][agentlightning.algorithm.verl.VERL], with placeholders for rollout and attempt IDs. When a rollout begins on the agent side, the framework uses the current `rollout_id` and `attempt_id` to format this template, generating a final, unique endpoint URL. This URL points to [VERL][agentlightning.algorithm.verl.VERL]'s internal proxy, allowing it to intercept and log all traffic for that specific attempt, for tracing and load balancing purposes. For agents created with the `@rollout` decorator, this resolution of the template is handled automatically ("auto-stripped"). Class-based agents will need to manually resolve the `ProxyLLM` using the rollout context.
+
+    ```python
+    proxy_llm = resources["main_llm"]
+    proxy_llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id)
+    ```
+
+## Customization
+
+Internally, [VERL][agentlightning.algorithm.verl.VERL] decomposes each agent execution into prompt–response pairs via the [Adapter][agentlightning.Adapter] and associates them with their corresponding reward signals as [Triplet][agentlightning.Triplet] objects. The final scalar reward, derived from the last triplet in the trajectory, is propagated to all preceding triplets following the [identical assignment strategy](https://arxiv.org/abs/2508.03680). This ensures that each triplet receives an identical reward signal and can be independently optimized as a valid RLHF trajectory within the VERL framework.
+
+At present, [VERL][agentlightning.algorithm.verl.VERL] does not expose fine-grained control over its reward propagation or credit assignment mechanisms. Users requiring customized reward shaping or trajectory decomposition are advised to clone and modify the [VERL][agentlightning.algorithm.verl.VERL] source implementation directly.

 ## Tutorials Using VERL

-TBD
+- [Train SQL Agent with RL](../how-to/train-sql-agent.md) - A practical example of training a SQL agent using VERL.

 ## References - Entrypoint

@@ -0,0 +1,106 @@
+# Changelog
+
+## Agent-lightning v0.2.1 (10/30/2025)
+
+Agent-lightning v0.2.1 is a stabilization release for v0.2.0. It introduces several bug fixes and new features, plus a number of unlisted CI improvements.
+
+### Bug fixes
+
+* Fix LiteLLM issues when restarting the proxy multiple times in the same process (#174 #206)
+* Fix LiteLLM model name selection when multiple servers use the same model (#197)
+* Fix store port conflict handling (#227)
+
+### New Features
+
+* Add trainer port option for client-server strategies (#198)
+
+### Documentation
+
+* Add tutorial for launching workers on separate machines (#213)
+* Add link to VERL framework (#210)
+* Add link to vLLM blog (#215)
+* Fix a couple of typos and avoid emacs backup files (#237)
+
+### New Contributors
+
+A warm welcome to our first-time contributors: @scott-vsi, @ddsfda99, @jeis4wpi 🎉
+
+**Full Changelog**: https://github.com/microsoft/agent-lightning/compare/v0.2.0...v0.2.1
+
+---
+
+## Agent-lightning v0.2.0 (10/22/2025)
+
+Agent-Lightning v0.2.0 introduces major framework improvements, new execution strategies, expanded documentation, and enhanced reliability across the agent training and deployment workflow. This release includes **78 pull requests** since v0.1.2.
+
+### Core Enhancements
+
+* **Lightning Store**: Added unified interface and implementation for Agent-lightning's core storage.
+* **Emitter**: Emitting any objects as spans to the store.
+* **Adapter** and **Tracer**: Adapting to OpenAI-like messages, and OpenTelemetry dummy tracer.
+* **LLM Proxy**: Added LLM Proxy as the first-class citizen in Agent-lightning.
+* **Agent Runner**: New version providing a more modular and robust runner design.
+* **Embedded Algorithms**: Algorithms are now embedded directly into trainers for simplicity.
+* **New Execution Strategies**: Introduced *Client-Server* and *Shared Memory* execution models.
+* **Trainer Updates**: Integrated v0.2 interfaces and FastAlgorithm validation.
+
+### Documentation & Examples
+
+* Revamped documentation with new guides for **agent creation**, **training**, **debugging**, and **store concepts**.
+* Improved quickstart tutorials, clarified installation and new deep-dive articles.
+* Added and updated examples: *SQL Agent*, *Calc-X*, *Local SFT*, *Search-R1*, and *APO algorithm*.
+
+### Developer Experience
+
+* Migrated build and CI pipelines to **1ES**, split workflows and aggregate badges for clarity.
+* Adopted **uv** as the dependency manager.
+* Added GPU-based pytest workflows for full test coverage.
+* Enhanced debugging UX, pre-commit configs, and linting (Pyright fixes, import sorting).
+
+### Ecosystem & Integrations
+
+* Added support for agents built with [**Agent-framework**](https://github.com/microsoft/agent-framework).
+* Added new community listings: [*DeepWerewolf*](https://github.com/af-74413592/DeepWerewolf) and [*AgentFlow*](https://agentflow.stanford.edu/).
+
+### New Contributors
+
+A warm welcome to our first-time contributors:
+@hzy46, @lunaqiu, @syeehyn, @linhx1999, @SiyunZhao, and @acured 🎉
+
+**Full changelog:** [v0.1.2 → v0.2.0](https://github.com/microsoft/agent-lightning/compare/v0.1.2...v0.2.0)
+
+---
+
+## Agent-lightning v0.1.2 (08/12/2025)
+
+### What's Changed
+* Add basic documentation in https://github.com/microsoft/agent-lightning/pull/33
+* RAG example by @wizardlancet in https://github.com/microsoft/agent-lightning/pull/21
+
+### New Contributors
+* @wizardlancet made their first contribution in https://github.com/microsoft/agent-lightning/pull/21
+
+**Full Changelog**: https://github.com/microsoft/agent-lightning/compare/v0.1.1...v0.1.2
+
+---
+
+## Agent-lightning v0.1.1 (08/06/2025)
+
+### What's Changed
+* Disable HTTP tracer tests and bump to 0.1.1 in https://github.com/microsoft/agent-lightning/pull/26
+* Fix trainer bugs in v0.1 in https://github.com/microsoft/agent-lightning/pull/24
+
+**Full Changelog**: https://github.com/microsoft/agent-lightning/compare/v0.1...v0.1.1
+
+---
+
+## Agent-lightning v0.1.0 (08/04/2025)
+
+The first release of Agent-lightning!
+
+- 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! 🤖
+- **Selectively** optimize one or more agents in a multi-agent system. 🎯
+- Embraces Reinforcement Learning, Automatic Prompt Optimization and more **algorithms**. 🤗
+
+Install via `pip install agentlightning`.
@@ -1,16 +1,21 @@
 # The Bird's Eye View of Agent-lightning

-This article summarizes how Agent-lightning (as of v0.2) wires algorithms, runners, and stores together and shows where auxiliary components (tracer, adapters, proxies) plug into the loop. Each section provides a diagram for a different perspective of the system.
+!!! warning "High Volume of Information Ahead"
+
+    This article provides an in-depth exploration of the Agent-lightning architecture.
+    It is not intended as a beginner’s guide or usage tutorial.
+
+This article summarizes how Agent-lightning (as of v0.2) wires the [Algorithm][agentlightning.Algorithm], [Runner][agentlightning.Runner], and [LightningStore][agentlightning.LightningStore] loop together and shows where auxiliary components (the [Tracer][agentlightning.Tracer], [Adapter][agentlightning.Adapter], and [LLM Proxy][agentlightning.LLMProxy]) plug into the loop. Each section provides a diagram for a different perspective of the system.

 ## Algorithm ↔ Runner ↔ Store data flow

 At its heart, Agent-lightning is built on three main components that work in a coordinated loop:

-* **Algorithm:** The "brain" of the system. It decides what tasks to run, learns from the results, and updates resources (like AI models or prompts).
-* **Runner:** The "worker" of the system. It executes tasks assigned by the algorithm, runs the agent, and records the results.
-* **LightningStore:** The central "database" and message queue. It acts as the single source of truth, storing tasks, results, and resources, and enabling communication between the Algorithm and Runner.
+* **[Algorithm][agentlightning.Algorithm]:** The "brain" of the system. It decides what tasks to run, learns from the results, and updates resources (like AI models or prompts).
+* **[Runner][agentlightning.Runner]:** The "worker" of the system. It executes tasks assigned by the algorithm, runs the agent, and records the results.
+* **[LightningStore][agentlightning.LightningStore]:** The central "database" and message queue. It acts as the single source of truth, storing tasks, results, and resources, and enabling communication between the Algorithm and Runner.

-The typical data flow in a training loop is as follows: The **Algorithm** enqueues tasks (called **Rollouts**) into the **Store**. A **Runner** then dequeues a task, executes it, and streams the results (called **Spans**) back to the Store. Once the task is complete, the Algorithm can query the new data from the Store to learn and update its resources.
+The typical data flow in a training loop is as follows: The **[Algorithm][agentlightning.Algorithm]** enqueues tasks (called **[Rollouts][agentlightning.Rollout]**) into the **[LightningStore][agentlightning.LightningStore]**. A **[Runner][agentlightning.Runner]** then dequeues a task, executes it, and streams the results (called **[Spans][agentlightning.Span]**) back to the store. Once the task is complete, the algorithm can query the new data from the store to learn and update its resources.

 The diagram below shows this fundamental interaction in a simple, non-parallel setup.

@@ -42,16 +47,16 @@ Solid lines represent direct calls, while dashed lines are asynchronous or long-

 We define the following terms, which may be helpful for understanding the diagram above.

-* **Resources:** A collection of assets to be tuned or trained. Agents perform rollouts against resources and collect span data. Algorithms use those data to update the resources. In RL training, the resources are a tunable model. In prompt tuning, the resources are prompt templates.
-* **Rollout:** A unit of work that an agent performs against a resource. A rollout (noun) can be incomplete, in which case it is also known as a **task**, **sample**, or **job** (these terms are used interchangeably). The agent executes its own defined workflow against the rollout — the process is also called "to rollout" (verb). After execution, the rollout (noun) is considered *complete*.
-* **Attempt:** A single execution of a rollout. One rollout can have multiple attempts in case of failures or timeouts.
-* **Span:** During the rollout, the agent can generate multiple spans (also known as "traces" or "events"). The recorded spans are collected in the store, which is crucial for understanding agent behavior and optimizing agents.
-* **Reward:** A special span that is defined as a number judging the quality of the rollout during some period of the rollout.
-* **Dataset:** A collection of incomplete rollouts (i.e., tasks) for the agent to process. The dual datasets (train, val) serve as the initial input for the algorithm to enqueue the first batch of rollouts.
+* **[Resources][agentlightning.Resource]:** A collection of assets to be tuned or trained. Agents perform rollouts against resources and collect span data. Algorithms use those data to update the resources. In RL training, the resources are a tunable model. In prompt tuning, the resources are prompt templates.
+* **[Rollout][agentlightning.Rollout]:** A unit of work that an agent performs against a resource. A rollout (noun) can be incomplete, in which case it is also known as a **task**, **sample**, or **job** (these terms are used interchangeably). The agent executes its own defined workflow against the rollout — the process is also called "to rollout" (verb). After execution, the rollout (noun) is considered *complete*.
+* **[Attempt][agentlightning.Attempt]:** A single execution of a rollout. One rollout can have multiple attempts in case of failures or timeouts.
+* **[Span][agentlightning.Span]:** During the rollout, the agent can generate multiple spans (also known as "traces" or "events"). The recorded spans are collected in the store, which is crucial for understanding agent behavior and optimizing agents.
+* **[Reward][agentlightning.emit_reward]:** A special span that is defined as a number judging the quality of the rollout during some period of the rollout.
+* **[Dataset][agentlightning.Dataset]:** A collection of incomplete rollouts (i.e., tasks) for the agent to process. The dual datasets (train, val) serve as the initial input for the algorithm to enqueue the first batch of rollouts.

 ### Store

-As discussed previously, the store is the central hub for all data in Agent-lightning. The store exposes a set of APIs for algorithms and runners to interact with the data; the most important ones are:
+As discussed previously, the [LightningStore][agentlightning.LightningStore] is the central hub for all data in Agent-lightning. The store exposes a set of APIs for algorithms and runners to interact with the data; the most important ones are:

 ```python
 from agentlightning.types import AttemptedRollout, ResourcesUpdate, Span, TaskInput
@@ -75,9 +80,11 @@ class LightningStore:
    ...
 ```

+These interfaces operate on [`AttemptedRollout`][agentlightning.AttemptedRollout], [`ResourcesUpdate`][agentlightning.ResourcesUpdate], [`Span`][agentlightning.Span], and [`TaskInput`][agentlightning.TaskInput] instances from `agentlightning.types`.
+
 As the APIs show, the store essentially provides a queue for rollouts and storage for resources, spans, and attempts. Developers should implement the store carefully to ensure data integrity and consistency, especially when multiple runners work in parallel across multiple attempts.

-The store is designed to be extensible. Users can implement their own store by inheriting from `LightningStore` and overriding methods. Agent-lightning provides a few reference implementations, such as `InMemoryLightningStore` (default) and `SqliteLightningStore` (under construction). When parallelized, the store may need special wrappers to ensure thread/process safety or delegate computation to a store in another process or machine.
+The store is designed to be extensible. Users can implement their own store by inheriting from [`LightningStore`][agentlightning.LightningStore] and overriding methods. Agent-lightning provides a few reference implementations, such as [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] (default) and `SqliteLightningStore` (under construction). When parallelized, the store may need special wrappers to ensure thread/process safety or delegate computation to a store in another process or machine.

 ## Supporting Components in the Loop

@@ -85,7 +92,7 @@ While the core loop is simple, Agent-lightning provides several components to ma

 ### Tracer

-The tracer is a component within the Runner that records detailed spans (events) during an agent's execution and sends them to the Store. Instead of requiring the agent to manually log every span, the tracer automatically instruments key methods (e.g., LLM calls) and captures their inputs, outputs, and metadata. This provides a detailed log of the agent's behavior with minimal effort.
+The [`Tracer`][agentlightning.Tracer] is a component within the [`Runner`][agentlightning.Runner] that records detailed spans (events) during an agent's execution and sends them to the [`LightningStore`][agentlightning.LightningStore]. Instead of requiring the agent to manually log every span, the tracer automatically instruments key methods (e.g., LLM calls) and captures their inputs, outputs, and metadata. This provides a detailed log of the agent's behavior with minimal effort.

 ```mermaid
 sequenceDiagram
@@ -114,11 +121,11 @@ sequenceDiagram
    Tracer->>Agent: Unapply instrumentation
 ```

-The above diagram shows the overall data flow between store, tracer and agent. In realistic, it's a bit more complicated than that. Spans are not emitted actively by the agent; they are intercepted by the tracer by hooking and instrumenting key methods used in the agents.  The tracer uses a callback (called exporter) to monitor events and log to the store. Before a rollout starts, the runner enters a `trace_context` before invoking the agent, wiring store identifiers into the tracer. Each span completion streams back to the store through `LightningSpanProcessor.on_end`, so the agent’s instrumentation lands in `add_otel_span`. If the agent’s rollout method returns a numeric reward, the runner emits one more OpenTelemetry span before finalizing the attempt.
+The above diagram shows the overall data flow between store, tracer and agent. In realistic, it's a bit more complicated than that. Spans are not emitted actively by the agent; they are intercepted by the tracer by hooking and instrumenting key methods used in the agents.  The tracer uses a callback (called exporter) to monitor events and log to the store. Before a rollout starts, the runner enters a [`trace_context`][agentlightning.Tracer.trace_context] before invoking the agent, wiring store identifiers into the tracer. Each span completion streams back to the store through `LightningSpanProcessor`, so the agent’s instrumentation lands in [`add_otel_span`][agentlightning.LightningStore.add_otel_span]. If the agent’s rollout method returns a numeric reward, the runner emits one more OpenTelemetry span before finalizing the attempt.

 ### Hooks

-Hooks are user-defined callback functions that allow you to augment a Runner's behavior at specific points in its lifecycle. You can use hooks to add custom logging, set up resources before a rollout begins, or tear them down after it ends. Hooks can be triggered at four key moments: `on_rollout_start`, `on_trace_start`, `on_trace_end`, and `on_rollout_end`.
+[`Hook`][agentlightning.Hook] implementations are user-defined callback functions that allow you to augment a [`Runner`][agentlightning.Runner]'s behavior at specific points in its lifecycle. You can use hooks to add custom logging, set up resources before a rollout begins, or tear them down after it ends. Hooks can be triggered at four key moments: `on_rollout_start`, `on_trace_start`, `on_trace_end`, and `on_rollout_end`.

 Users should pay special attention to the difference between `on_trace_end` and `on_rollout_end`. The former is called right before the tracer exits the trace context, while the latter is called after the runner processes the final leftover rewards and spans, and finalizes the attempt in the store.

@@ -154,9 +161,9 @@ sequenceDiagram

 ### Adapter

-The Adapter is a component used by the Algorithm to transform raw data from the Store into a format suitable for learning. Runners stream raw spans into the Store during execution. Later, the Algorithm queries these spans and uses an Adapter to convert them into structured data, like training examples for a reinforcement learning model.
+The [`Adapter`][agentlightning.Adapter] is a component used by the [`Algorithm`][agentlightning.Algorithm] to transform raw data from the [`LightningStore`][agentlightning.LightningStore] into a format suitable for learning. Runners stream raw spans into the store during execution. Later, the algorithm queries these spans and uses an adapter to convert them into structured data, like training examples for a reinforcement learning model.

-For instance, the `TracerTraceToTriplet` processes OpenTelemetry spans to create `(prompt, response, reward)` triplets, which are the fundamental data structure for many RL fine-tuning algorithms.
+For instance, the [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] processes OpenTelemetry spans to create `(prompt, response, reward)` triplets, which are the fundamental data structure for many RL fine-tuning algorithms.

 ```mermaid
 flowchart LR
@@ -168,13 +175,13 @@ flowchart LR

 ### LLM Proxy

-The LLM Proxy is an optional bridge component that sits between an agent and the algorithms' resources. It acts as a centralized endpoint for all LLM calls. Usually the proxy URL is added to the store as a special resource, so that the runner can fetch it along with other resources when dequeuing a rollout. During rollouts, the runner invokes the proxy's HTTP endpoint instead of calling a model backend directly.
+The [`LLMProxy`][agentlightning.LLMProxy] is an optional bridge component that sits between an agent and the algorithms' resources. It acts as a centralized endpoint for all LLM calls. Usually the proxy URL is added to the store as a special resource, so that the [`Runner`][agentlightning.Runner] can fetch it along with other resources when dequeuing a rollout. During rollouts, the runner invokes the proxy's HTTP endpoint instead of calling a model backend directly.

 This design offers several benefits:

-1. **Instrumentation:** It automatically captures detailed traces of LLM interactions (prompts, responses, metadata) and sends them to the Store, complementing the Tracer, especially when the agent's code is hard to instrument directly.
+1. **Instrumentation:** It automatically captures detailed traces of LLM interactions (prompts, responses, metadata) and sends them to the store, complementing the tracer, especially when the agent's code is hard to instrument directly.
 2. **Backend Abstraction:** It provides a unified interface for various LLM backends (OpenAI, Anthropic, local models) and can add features like retry logic, rate limiting, and caching.
-3. **Resource Management:** The Algorithm can dynamically update which LLM the agent uses (e.g., swapping to a newly fine-tuned model) by simply swapping the backend model the proxy is using, without interrupting the agent's code.
+3. **Resource Management:** The algorithm can dynamically update which LLM the agent uses (e.g., swapping to a newly fine-tuned model) by simply swapping the backend model the proxy is using, without interrupting the agent's code.

 The benefits above seem to be all discussed within the context of model fine-tuning. As a matter of fact, the proxy can be useful for prompt tuning as well. The algorithm can register one of the following two types of endpoints into the proxy:

@@ -221,7 +228,7 @@ In this diagram, the store receives spans from both the proxy and the runner. We

 ### Trainer

-The Trainer is the high-level orchestrator that initializes and connects all major components -- algorithm, runner, store, tracer, adapter, LLM proxy, and hooks. The components can have a lifecycle as long as the trainer. The trainer manages their lifecycles and handles dependency injection, ensuring that every part of the system operates within a consistent and shared environment.
+The [Trainer][agentlightning.Trainer] is the high-level orchestrator that initializes and connects all major components -- [Algorithm][agentlightning.Algorithm], [Runner][agentlightning.Runner], [LightningStore][agentlightning.LightningStore], [Tracer][agentlightning.Tracer], [Adapter][agentlightning.Adapter], [LLM Proxy][agentlightning.LLMProxy], and [Hook][agentlightning.Hook]. The components can have a lifecycle as long as the trainer. The trainer manages their lifecycles and handles dependency injection, ensuring that every part of the system operates within a consistent and shared environment.

 Below, we demonstrate how the components relate to each other and their roles. We first clarify the roles and relationships shown in the diagram:

@@ -230,7 +237,7 @@ Below, we demonstrate how the components relate to each other and their roles. W
 3. **References:** weak links for coordination without ownership.
 4. **Uses:** components that are temporarily interacted with.

-For example, the store is injected into the algorithm and runner. The tracer and agent are injected into the runner. The adapter and LLM proxy are injected into the algorithm. The store is further injected into the tracer, adapter and LLM proxy by the runner and algorithm respectively.
+For example, the [LightningStore][agentlightning.LightningStore] is injected into the [Algorithm][agentlightning.Algorithm] and [Runner][agentlightning.Runner]. The [Tracer][agentlightning.Tracer] and [LitAgent][agentlightning.LitAgent] are injected into the runner. The [Adapter][agentlightning.Adapter] and [LLM Proxy][agentlightning.LLMProxy] are injected into the algorithm. The store is further injected into the tracer, adapter and LLM proxy by the runner and algorithm respectively.

 ```mermaid
 flowchart TD
@@ -294,6 +301,8 @@ flowchart TD

 ## Putting It All Together: A Reinforcement Learning Example (VERL)

+[](){ #birds-eye-view-verl-example }
+
 VERL shows how an algorithm consumes the shared infrastructure. For historical reasons, code lives in `agentlightning.algorithm.verl` and `agentlightning.verl`. The latter is legacy and reuses terms like `Trainer` in confusing ways. The former is a thin wrapper that conforms to the new algorithm interface. Future versions will merge the two.

 Reinforcement learning aims to learn a policy that takes actions in states to maximize expected reward. For agents, the policy is usually a language model. Inputs are prompts (state). Outputs are generated text (action). A numeric score judges quality (reward). The `(state, action, reward)` **triplet** is the basic learning unit.
@@ -358,9 +367,9 @@ sequenceDiagram

 ## Execution Strategies and Parallelism

-Readers might have observed from the diagram above that there is absolutely no communication between (1) runner and agents and (2) algorithm. The only overlap of them is the trainer and store. This observation is very clear with the diagram within the trainer section. This design allows us to flexibly scale the runner and algorithm independently, which is crucial for large-scale training.
+Readers might have observed from the diagram above that there is absolutely no communication between (1) runner and agents and (2) algorithm. The only overlap of them is the [Trainer][agentlightning.Trainer] and [LightningStore][agentlightning.LightningStore]. This observation is very clear with the diagram within the trainer section. This design allows us to flexibly scale the runner and algorithm independently, which is crucial for large-scale training.

-Agent-lightning packages two executable bundles: a runner bundle (runner, tracer, hooks, agent) and an algorithm bundle (algorithm, adapter, LLM proxy). Both share the store. The trainer initializes and connects the bundles.
+Agent-lightning packages two executable bundles: a runner bundle ([Runner][agentlightning.Runner], [Tracer][agentlightning.Tracer], [Hook][agentlightning.Hook], [LitAgent][agentlightning.LitAgent]) and an algorithm bundle ([Algorithm][agentlightning.Algorithm], [Adapter][agentlightning.Adapter], [LLM Proxy][agentlightning.LLMProxy]). Both share the [LightningStore][agentlightning.LightningStore]. The trainer initializes and connects the bundles.

 ```mermaid
 graph TD
@@ -385,7 +394,7 @@ graph TD
    linkStyle 0,1,2,3,4 opacity:0;
 ```

-An **execution strategy**, defined and owned by the trainer, governs how algorithm and runner bundles are placed, connected, scaled, and aborted. It serves four primary purposes.
+An [execution strategy][agentlightning.ExecutionStrategy], defined and owned by the trainer, governs how algorithm and runner bundles are placed, connected, scaled, and aborted. It serves four primary purposes.

 Execution strategies first determine **bundle placement** — whether the two bundles run in the same thread, process, machine, or across separate machines. They also define **store management**, wrapping the store and specifying how data is shared between bundles.

@@ -397,7 +406,7 @@ Agent-lightning currently provides two execution strategies: **shared-memory** a

 ### Shared-memory Strategy

-`SharedMemoryExecutionStrategy` runs algorithm and runner bundles as threads in one process. The strategy wraps the store with `LightningStoreThreaded`, which guards calls with a lock for safe concurrency.
+[`SharedMemoryExecutionStrategy`][agentlightning.SharedMemoryExecutionStrategy] runs algorithm and runner bundles as threads in one process. The strategy wraps the store with [`LightningStoreThreaded`][agentlightning.LightningStoreThreaded], which guards calls with a lock for safe concurrency.

 This is good for lightweight debugging because components share one Python heap and avoid serialization. It is not suitable for heavy RL training or compute-intensive agents.

@@ -431,7 +440,9 @@ You can configure which role runs on the main thread. If the main thread runs th

 ### Client-server Strategy

-`ClientServerExecutionStrategy` splits concerns across processes. The algorithm bundle starts a `LightningStoreServer` (HTTP API) that wraps the underlying store. Runners connect via `LightningStoreClient` to call the same interface over REST. The server embeds a client to support algorithm-launched subprocesses (e.g., an LLM proxy worker) that need to talk back to the algorithm’s process through the same API.
+[](){ #birds-eye-view-client-server-strategy }
+
+[`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy] splits concerns across processes. The algorithm bundle starts a [`LightningStoreServer`][agentlightning.LightningStoreServer] (HTTP API) that wraps the underlying store. Runners connect via [`LightningStoreClient`][agentlightning.LightningStoreClient] to call the same interface over REST. The server embeds a client to support algorithm-launched subprocesses (e.g., an LLM proxy worker) that need to talk back to the algorithm’s process through the same API.

 Currently this design introduces an extra wrapper in the Server side (as shown in the diagram), which helps debugging and improves fault tolerance. We might revisit this design in the future and enforce the client to be the only way to communicate with the store.

@@ -485,11 +496,11 @@ flowchart TD
 Continuous learning keeps the algorithm loop running while runners report tasks and spans opportunistically. Key differences from batch mode:

 1. The algorithm does not enqueue rollouts from a fixed dataset. Runners report tasks/rollouts and spans spontaneously.
-2. The algorithm can wait for rollouts with a expected set of rollout IDs, but more oftenly polls for new rollouts and spans or waits for a count to arrive.
-3. The runner processes one rollout at a time via `step(task)` instead of exhausting a task queue. It notifies the store when starting a rollout so the store records it.
+2. The algorithm can wait for rollouts with a expected set of rollout IDs, but more often polls for new rollouts and spans or waits for a count to arrive.
+3. The [`Runner`][agentlightning.Runner] processes one rollout at a time via [`step(task)`][agentlightning.Runner.step] instead of exhausting a task queue. It notifies the store when starting a rollout so the store records it.
 4. A user or higher-level loop controls which resources the next step uses and when to retry.

-Spans, adapters, and LLM proxies work the same way.
+[Spans][agentlightning.Span], [Adapter][agentlightning.Adapter] implementations, and the [LLM Proxy][agentlightning.LLMProxy] work the same way.

 ```mermaid
 sequenceDiagram
@@ -1,31 +0,0 @@
-# Server-client Architecture
-
-Article to be written.
-
-```mermaid
-sequenceDiagram
-    participant RL as RL Framework
-    participant TS as Training Server
-    participant AC as Agent Client
-    participant AG as Agent
-
-    AC->>TS: Upload Dataset (1)
-    RL->>TS: Start RL Server (2)
-    TS->>RL: Latest Model (3)
-
-    loop for each batch of tasks
-        loop for each task in the batch
-            AC->>TS: Request Task (4)
-            TS->>AC: Send Task & Model API (5)
-            AC->>AG: Run Agent with Task & Model API (6)
-            loop for each LLM call
-                AC->>AG: Prompt (7)
-                AG->>AC: Response (8)
-            end
-            AG->>AC: Rewarded Trace (9)
-            AC->>TS: Send Rewarded Trace (10)
-        end
-        TS->>RL: Send Batch of Traces (11)
-        RL->>TS: Return Updated Model (12)
-    end
-```
@@ -0,0 +1,135 @@
+# Serving LLMs under Agent-lightning
+
+Agent-lightning focuses on data, learning signals, and control flow — **not** on running model inference. This deep dive explains how to **serve** a model alongside Agent-lightning so runners can call it reliably, how the **LLM Proxy** fits into the loop, and why **token IDs** matter if you care about correctness in training and evaluation.
+
+## General background on LLM serving
+
+[](){ #general-llm-serving-background }
+
+Serving a model is essential if you want to train it, especially when you use the model’s own generations as training data. We’ll briefly review the general background to ensure all readers are aligned.
+
+Modern LLM servers solve a difficult scheduling problem: keeping GPUs fully utilized while handling prompts of different lengths, streaming tokens as they arrive, and fitting large KV caches into limited memory. Techniques like [**continuous batching**](https://www.anyscale.com/blog/continuous-batching-llm-inference) and [**paged attention**](https://arxiv.org/abs/2309.06180) address these challenges. Continuous batching interleaves decoding across requests to reuse weights efficiently; with careful memory planning, it achieves major throughput gains without increasing latency. PagedAttention reduces KV-cache fragmentation so batching remains effective as sequences grow. See [vLLM’s PagedAttention paper](https://arxiv.org/abs/2309.06180) and [industry analyses](https://www.baseten.co/blog/continuous-vs-dynamic-batching-for-ai-inference/) for details. Balancing inference correctness and efficiency is difficult — a [recent blog](https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/) from Thinking Machines Labs highlights how inference nondeterminism ultimately affects training.
+
+Beyond scheduling, servers expose an HTTP API, often **OpenAI-compatible** (`/v1/chat/completions` and `/v1/responses`), which is itself a complex stack. In addition to text prompts and chat messages, the API defines many parameters and response fields such as [tool calls](https://platform.openai.com/docs/guides/function-calling), [structured output](https://platform.openai.com/docs/guides/structured-outputs), and [multimodal support](https://platform.openai.com/docs/guides/images-vision). Much effort has been put into implementing all these parameters for many frameworks. Popular engines like **vLLM** and [**SGLang**](https://github.com/sgl-project/sglang) ship with OpenAI-compatible frontends so you can reuse existing client code. [Ollama](https://ollama.com/blog/openai-compatibility) and [llama.cpp](https://llama-cpp-python.readthedocs.io/en/latest/server/) provide similar capabilities. However, because models differ internally, each framework interprets and implements the API slightly differently. Even with identical requests, the tokens passed to the model can vary substantially across frameworks.
+
+## What Agent-lightning expects from a served LLM
+
+Most of the issues above either have workarounds or remain open research problems. Keep them in mind, but the key question is: what does Agent-lightning expect from a served LLM? The answer includes at least two things:
+
+* An OpenAI-compatible **Chat Completions** or **Responses** endpoint the agent can call during rollouts.
+* Optional training and debugging signals: **logprobs**, **usage**, and ideally **token IDs**. (OpenAI’s public API exposes usage and logprobs, but **not** token IDs — more on [why IDs matter][token-ids-matter] later.)
+
+## Launching a serving framework
+
+For many algorithms, you’ll start an engine (e.g., **vLLM** or **SGLang**) before rollouts, then shut it down afterward to free GPU memory. Most frameworks provide a one-line “serve” command to launch the OpenAI-compatible server. You can use those to bring up `/v1/chat/completions` with your checkpoint, ensuring streaming and any required tool-calling features are enabled. A working example is shown in [Unsloth SFT](../how-to/unsloth-sft.md).
+
+Weight updates — which occur after each training step — are trickier. Some frameworks like [vLLM](https://vllm.ai/) support hot-updating model weights, but it’s usually simpler and more reliable to restart the engine to load new weights. For medium-sized tasks (hundreds of rollouts taking 10+ minutes), the restart overhead (under 30 seconds) is typically negligible.
+
+If you’re using Agent-lightning’s [**VERL**][agentlightning.algorithm.verl.VERL] integration, the algorithm can **manage the server automatically**. The [VERL framework](https://github.com/volcengine/verl) intelligently allocates compute resources and wraps vLLM/SGLang behind an `AsyncLLMServer` abstraction. You can directly use this as the LLM endpoint for agents. Since VERL can spawn multiple vLLM replicas, using [`LLMProxy`][agentlightning.LLMProxy] to manage them adds an additional safety layer.
+
+A full sequence diagram of how [VERL][agentlightning.algorithm.verl.VERL] interacts with the LLM server and proxy is available [here][birds-eye-view-verl-example].
+
+## LLM Proxy
+
+The **LLM Proxy** is a utility class in Agent-lightning, built on [LiteLLM](https://docs.litellm.ai/), that sits between runners and your backend engine(s) or server(s). In Agent-lightning it acts as a single URL registered as a [`Resource`][agentlightning.Resource] in the store, offering three key benefits:
+
+1. **Unified endpoint & hot-swaps.** You can redirect traffic between OpenAI, Anthropic, local vLLM/SGLang, or canary checkpoints without modifying agent code — simply repoint the proxy.
+2. **First-class tracing.** The proxy emits **OpenTelemetry** spans for every call and sends them to the [`LightningStore`][agentlightning.LightningStore]. It includes rollout and attempt identifiers in request headers so spans are correctly attributed. Sequence numbers are allocated monotonically via the store to [prevent clock-skew issues][distributed-tracing] and allow reliable reconstruction of execution trees.
+3. **Token IDs.** The proxy can return prompt and response token IDs along with the model output. More details are available in the [next section][token-ids-matter].
+
+Operationally, running the proxy alongside the algorithm works best: the algorithm registers the backend (e.g., the vLLM URL) via [`LLMProxy.update_model_list`][agentlightning.LLMProxy.update_model_list], publishes the proxy URL as a resource via [`LightningStore.add_resources`][agentlightning.LightningStore.add_resources], and runners simply use that URL during rollouts. This mirrors many production client–server setups.
+
+## Token IDs and why they matter
+
+[](){ #token-ids-matter }
+
+This section explains how Agent-lightning handles and uses token IDs — a subtle but important detail for training stability and accuracy.
+
+Most agents interact with LLMs via **Chat Completion APIs**, exchanging chat messages. There are two main approaches to collecting training data from such agents.
+
+!!! note
+
+    Tokenization here refers to the process of converting **Chat Messages** into **Token IDs**. Detokenization is the reverse process of converting **Token IDs** back to **Chat Messages**. Normally, the tokenizer is published along with the pretrained model, which includes a vocabulary, special tokens, and a chat template to dealing with chat messages.
+
+**1. Retokenizing chat messages.**
+In this approach, you store chat messages as text and let training algorithms **retokenize** them later, as done in many SFT workflows (e.g., [HuggingFace SFT](https://huggingface.co/docs/trl/sft_trainer)).
+In practice, we’ve found this method unstable and less accurate. The chart below compares training results. The retokenization approach is run twice. All settings are the same except for the retokenization approach.
+
+<div style="height:400px">
+<canvas data-chart='{
+  "type": "line",
+  "data": {
+    "labels": [0.0, 32.0, 64.0, 96.0, 128.0, 160.0, 192.0, 224.0, 256.0, 288.0, 320.0, 352.0, 384.0, 416.0, 448.0, 480.0],
+    "datasets": [
+      {
+        "label": "With Token IDs from Retokenization",
+        "data": [0.49, 0.512, 0.54, 0.532, 0.54, 0.466, 0.328, 0.358, 0.348, 0.35, 0.346, 0.372, 0.346, 0.33, 0.346, 0.332],
+        "spanGaps": true
+      },
+      {
+        "label": "Retokenization (Second Run)",
+        "data": [0.494, 0.526, 0.536, 0.554, 0.544, 0.556, 0.568, 0.552, 0.45, 0.466, 0.474, 0.47, 0.464, 0.476, 0.488, 0.432],
+        "spanGaps": true
+      },
+      {
+        "label": "With Token IDs from Engine",
+        "data": [0.494, 0.522, 0.514, 0.538, 0.53, 0.564, 0.564, 0.586, 0.594, 0.604, 0.618, 0.584, 0.606, 0.558, 0.612, 0.588],
+        "spanGaps": true
+      }
+    ]
+  },
+  "options": {
+    "interaction": {
+      "mode": "nearest",
+      "intersect": false
+    },
+    "plugins": {
+      "legend": {
+        "display": true,
+        "position": "top"
+      },
+      "title": {
+        "display": true,
+        "text": "Agent Training Results Comparison"
+      }
+    },
+    "scales": {
+      "x": {
+        "title": {
+          "display": true,
+          "text": "Step"
+        }
+      },
+      "y": {
+        "title": {
+          "display": true,
+          "text": "Reward"
+        }
+      }
+    }
+  }
+}'></canvas>
+</div>
+
+This instability has three causes. Firstly, chat template used in different frameworks could be slightly different. For example, one single LLaMA model can work with multiple chat templates (multiple in [vLLM](https://github.com/vllm-project/vllm/tree/1d165d6d859d3c50720f0c07209db2363c4fd33b/examples) and one in [HuggingFace](https://huggingface.co/meta-llama)). It's possible that the chat template used in detokenization is different from the one used in tokenization (this is actually an implementation bug).
+
+Secondly, a word might be generated as two tokens (e.g., `H + AVING`) but later retokenized as `HAV + ING`. The text looks identical, but the token IDs differ from what the model originally produced.
+
+Thirdly, a generated tool call text like `<tool_call>{ "name": ... }</tool_call>` is parsed by tool call parser into an object that is required by chat completion API. Later, the object is rendered back to `<tool_call>{ "name": ... }</tool_call>` and retokenized again, tool call parsing and re-rendering might cause changes in whitespace and formatting. In some situations, JSON errors may even be auto-corrected by the tool call parser — masking the model’s true generation errors and preventing them from being trained away.
+
+----
+
+**2. Saving token IDs directly.**
+The alternative is to save the token IDs generated by the model, as done in RL setups like [Tinker](https://thinkingmachines.ai/tinker/). This requires a training pipeline that treats tokens as first-class entities, meaning agents must communicate with the inference engine at the token level.
+
+However, most agents — especially those built with frameworks like LangChain — rely on OpenAI-compatible APIs and can’t tokenize or detokenize themselves. As mentioned [earlier][general-llm-serving-background], implementing this layer manually is complex and error-prone. Some frameworks implement custom solutions (e.g., [VERL Agent Loop](https://github.com/volcengine/verl/blob/4da0d3d3188072772cb2ec817b3d6cf4a463821f/recipe/langgraph_agent/chat_model.py), [Tinker Renderer](https://github.com/thinking-machines-lab/tinker-cookbook/blob/34a6588d7055040c259985d98e71c0140b389ba7/tinker_cookbook/renderers.py)), while others leave it to users (e.g., [SkyRL Search-R1](https://novasky-ai.notion.site/skyrl-searchr1)).
+
+----
+
+A better solution is to use an **OpenAI-compatible API that returns token IDs directly.** This lets agents continue using familiar APIs while capturing token IDs via [tracing](../tutorials/traces.md) for training. The limitation, of course, is that the serving framework must actually support this capability.
+
+When Agent-lightning was first released, we implemented an [instrumented vLLM server](https://github.com/microsoft/agent-lightning/blob/v0.1/agentlightning/instrumentation/vllm.py) that monkey-patched vLLM’s OpenAI server to return token IDs. Since then, the Agent-lightning and vLLM teams have collaborated to add this feature directly to [vLLM core](https://github.com/vllm-project/vllm/pull/22587). Starting with **vLLM v0.10.2**, the OpenAI-compatible API includes a [`return_token_ids` parameter](https://docs.vllm.ai/en/v0.10.2/serving/openai_compatible_server.html#api-reference), allowing token IDs to be requested alongside chat messages. SGLang has tracked [similar feature requests](https://github.com/sgl-project/sglang/issues/2634), though its OpenAI-compatible layer doesn’t yet support them.
+
+In short, when using vLLM v0.10.2 or newer, [`LLMProxy`][agentlightning.LLMProxy] automatically adds `return_token_ids` to each request so the engine includes token IDs in its response. For older vLLM versions, you still need the instrumented version (via `agl vllm` CLI command).
+
+Finally, if you only save token IDs in spans, it will have its own limitations — if you train one model using spans from another model with a different tokenizer, incompatibilities can arise. In practice, though, spans in Agent-lightning always store both chat messages and token IDs (actually the full request and response objects), allowing you to fall back to retokenization when necessary.
@@ -0,0 +1,204 @@
+# Understanding Store
+
+The **[`LightningStore`][agentlightning.LightningStore]** is the central coordination point for Agent-lightning. It holds the task queue, rollouts, attempts, spans, and versioned resources, and exposes a small API both Runners and Algorithms use to communicate. This document explains what’s in the store, how statuses transition, how spans are recorded, and the concurrency model (threads & processes).
+
+## What’s in the Store?
+
+![Store Architecture](../assets/store-api-visualized.svg){ .center }
+
+At a high level:
+
+* **Task Queue** – [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout] adds work; workers poll with [`dequeue_rollout`][agentlightning.LightningStore.dequeue_rollout]. When a rollout is dequeued, it automatically creates a new attempt associated with itself.
+* **Rollouts** – A rollout is one unit of work. It has input, metadata, links to resources, and a lifecycle (`queuing → preparing → running → ...`). Valid [RolloutStatus][agentlightning.RolloutStatus] are **`queuing`**, `preparing`, `running`, `succeeded`, `failed`, **`requeuing`**, **`cancelled`**. For algorithms and runners, the rollout can be seen as a whole, without worrying about the internal attempts.
+* **Attempts** – Each rollout can have multiple executions (retries). Attempts track [`status`][agentlightning.Attempt.status], [`start_time`][agentlightning.Attempt.start_time], [`end_time`][agentlightning.Attempt.end_time], [`last_heartbeat_time`][agentlightning.Attempt.last_heartbeat_time] and link to spans. Valid [AttemptStatus][agentlightning.AttemptStatus] are `preparing`, `running`, `succeeded`, `failed`, `requeuing`, `cancelled`.
+* **Spans** – Structured trace events produced by the Tracer during an attempt. Spans are ordered by a **monotonic sequence id** per `(rollout_id, attempt_id)`.
+* **Resources** – Versioned, named bundles (e.g., prompt templates) referenced by rollouts.
+
+Rollout and Task share the same surface in practice: [`Rollout.input`][agentlightning.types.Rollout] is the task input. The queue stores rollouts that are not yet running; [Runners][agentlightning.Runner] dequeue them and update the same rollout’s status as work progresses.
+
+All [`LightningStore`][agentlightning.LightningStore] implementations must inherit from [`LightningStore`][agentlightning.LightningStore] and override the methods to implement the storage logic.
+
+Before we look at status transitions, it helps to keep in mind that rollouts are the “outside view,” while attempts are the “inside view.” Attempts are what actually run; rollouts summarize the latest attempt plus a small set of control actions like queueing and cancellation.
+
+## Attempt Status Transitions
+
+The status model is intentionally small and operationally clear.
+
+```mermaid
+stateDiagram-v2
+direction LR
+
+[*] --> preparing: <b>Runner calls</b> dequeue_rollout()<br>or start_rollout()<br>or start_attempt()
+preparing --> running: <b>Runner calls</b><br>add_[otel_]span()<br>for the first time
+
+state c_runner <<choice>>
+state c_watch  <<choice>>
+
+preparing --> c_runner: <b>Runner calls</b><br>update_attempt(...)</b>
+running --> c_runner: <b>Runner calls</b><br>update_attempt(...)
+running --> c_watch: <b>Watchdog checks</b>
+preparing --> c_watch: <b>Watchdog checks</b>
+
+state "Client-set outcome" as Client {
+  direction TB
+  succeeded
+  failed
+}
+state "Watchdog / policy" as Watch {
+  direction TB
+  timeout
+  unresponsive
+}
+
+c_runner --> succeeded: update_attempt(status=succeeded)
+c_runner --> failed: update_attempt(status=failed)
+
+c_watch --> timeout: now - start_time > timeout_seconds
+c_watch --> unresponsive: now - last_heartbeat > unresponsive_seconds
+
+unresponsive --> running: <b>Runner calls</b><br>add_[otel_]span()
+```
+
+Each attempt begins in **preparing**, created either when a rollout is dequeued or explicitly started. It transitions to **running** the first time a span is recorded. From there, a few clear rules govern how it can change:
+
+* When the runner explicitly marks completion, the attempt becomes **succeeded** or **failed** (when the runner catches exception thrown out by the agent).
+* When the watchdog detects that the total elapsed time since start exceeds the configured limit, it marks the attempt as **timeout**.
+* If heartbeats stop arriving for too long, the watchdog marks it **unresponsive**.
+* A new span from the runner can immediately revive an **unresponsive** attempt back to **running**.
+
+!!! info "What's a Watchdog?"
+
+    The watchdog enforces timing and liveness rules defined by each rollout’s [`RolloutConfig`][agentlightning.RolloutConfig]. It’s not a separate thread or service, but a function periodically invoked (e.g., before store mutations) to keep attempts healthy and consistent.
+
+This simple model allows the system to distinguish between normal termination, abnormal stalling, and recoverable interruption without additional state flags.
+
+
+## Rollout Transition Map
+
+Rollout status is an **aggregated view** of its latest attempt’s status, with additional transitions for queueing and explicit cancellation.
+
+A rollout’s retry behavior is controlled by [`Rollout.config`][agentlightning.types.Rollout] (a [`RolloutConfig`][agentlightning.types.RolloutConfig]). The key fields are:
+
+* [`timeout_seconds`][agentlightning.RolloutConfig.timeout_seconds] – maximum wall-clock time for an attempt before it is marked `timeout`.
+* [`unresponsive_seconds`][agentlightning.RolloutConfig.unresponsive_seconds] – maximum silence between heartbeats before an attempt is marked `unresponsive`.
+* [`max_attempts`][agentlightning.RolloutConfig.max_attempts] – total number of attempts allowed for the rollout (including the first).
+* [`retry_condition`][agentlightning.RolloutConfig.retry_condition] – which attempt terminal statuses should trigger a retry (e.g., `["failed", "timeout", "unresponsive"]`).
+
+**How it plays out:** The runner works on attempt `k`. If the attempt ends in a status that is listed in `retry_condition`, and `k < max_attempts`, the rollout moves to **requeuing** and the store creates attempt `k+1`. Otherwise, the rollout becomes **failed** (or **succeeded** if the runner marked it so). `timeout_seconds` and `unresponsive_seconds` are enforced by the watchdog and feed into the same decision flow.
+
+A minimal example of how to use `RolloutConfig`:
+
+```python
+from agentlightning import RolloutConfig
+
+# Retry on explicit failures or timeouts, up to 3 attempts in total.
+cfg = RolloutConfig(
+    timeout_seconds=600,
+    unresponsive_seconds=120,
+    max_attempts=3,
+    retry_condition=["failed", "timeout"]
+)
+
+# When creating/enqueuing a rollout, attach this config.
+# The store will propagate attempt outcomes according to cfg.
+rollout = await store.enqueue_rollout(input, config=cfg)
+```
+
+| Latest attempt status                 | Rollout transition                        | Notes / guards                                                                                    |
+| ------------------------------------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------- |
+| N/A                                   | `queuing`                                 | Created by `enqueue_rollout()`.                                                                   |
+| `preparing`                           | `queuing/requeuing` → `preparing`         | Typically `dequeue_rollout()` or `start_rollout()`/`start_attempt()` creates a new attempt.       |
+| `running`                             | `preparing/queuing/requeuing` → `running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `propagate_status`. |
+| `succeeded`                           | `*` → `succeeded`                         | Terminal. Rollout `end_time` set.                                                                 |
+| `failed` / `timeout` / `unresponsive` | `*` → `requeuing`                         | **Only if** `status ∈ retry_condition ∧ sequence_id < max_attempts`.                              |
+| `failed` / `timeout` / `unresponsive` | `*` → `failed`                            | Otherwise (no retries left or retries disabled).                                                  |
+| `*`                                   | `*` → `cancelled`                         | Explicitly set by `update_rollout(status=cancelled)`.                                             |
+
+!!! note "Why aggregation?"
+
+    In code, we use `propagate_status()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollout’s transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
+
+## Spans
+
+Every traceable operation in a rollout is stored as a [Span][agentlightning.Span]. Spans not only capture fine-grained instrumentation but also act as periodic heartbeats that demonstrate liveness. The first span marks activation; each subsequent one both extends the trace and refreshes the attempt’s [`last_heartbeat_time`][agentlightning.Attempt.last_heartbeat_time]. If no span arrives within the configured [`unresponsive_seconds`][agentlightning.RolloutConfig.unresponsive_seconds], the watchdog downgrades the attempt to **unresponsive** until activity resumes.
+
+Spans are indexed by `(rollout_id, attempt_id, sequence_id)` where the sequence is monotonic. Tracing analysis tools like [Adapter][agentlightning.Adapter] usually rely on "time order" to reconstruct the trace. However, in a distributed system, the recorded start time and end time of a span are not necessarily in the right order when they aggregated into a central store. Therefore, we enforce every span creation to retrieve a monotonically increasing [`sequence_id`][agentlightning.Span.sequence_id] from the store before adding the span.
+
+!!! note
+
+    In practice, one `sequence_id` can be used to create multiple spans. In that case, the orders between the multiple spans are determined by the order of `start_time` and `end_time` of the spans.
+
+### OpenTelemetry conversion
+
+Runners often produce [OpenTelemetry `ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/#attributes) objects directly. The store normalizes them into [`Span`][agentlightning.Span] as follows:
+
+1. The runner first requests [`get_next_span_sequence_id`][agentlightning.LightningStore.get_next_span_sequence_id] via `sequence_id = await store.get_next_span_sequence_id(rollout_id, attempt_id)`. This guarantees ordering within the attempt regardless of clock skew.
+2. `trace_id`, `span_id`, `parent_id`, `name`, `status`, timestamps, attributes, events, links, and resource are copied from the OTEL span. Timestamps are auto-normalized to seconds (nanoseconds are converted).
+3. OTEL `SpanContext` and parent context are preserved so downstream tools can correlate traces across systems.
+4. Any additional serializable fields present on the `ReadableSpan` are retained in the stored span (after safe JSON serialization), which keeps the representation forward-compatible.
+
+Programmatically this is encapsulated by [`Span.from_opentelemetry(readable_span, rollout_id, attempt_id, sequence_id)`][agentlightning.Span.from_opentelemetry]; [`store.add_otel_span(...)`][agentlightning.LightningStore.add_otel_span] simply wraps the fetch-then-add flow. The end result is a store span that is stable to sort, merge, and query, while still preserving the richness of the original OTEL payload.
+
+!!! tip
+
+    [`add_span`][agentlightning.LightningStore.add_span] or [`add_otel_span`][agentlightning.LightningStore.add_otel_span] both appends a span *and* acts as a heartbeat that can revive `unresponsive` → `running`.
+
+## Store Implementations
+
+Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore]:
+
+- Fast startup, zero external dependencies, and ideal for local development, CI, and unit tests.
+- Fully asyncio-safe for writes; most reader operations can iterate without locks, except those that need to perform multiple queries.
+- Includes a best-effort span eviction policy once memory crosses a configured watermark; querying evicted spans raises a clear error so callers can fall back.
+
+For production you will likely want persistence. We’re actively building a SQLite-backed store that keeps the same API surface while adding durability, crash recovery, and better historical span queries. If you need something sooner, implement your own store by subclassing [`LightningStore`][agentlightning.LightningStore] and providing concrete storage for the small set of abstract methods (`enqueue_rollout`, `dequeue_rollout`, `update_attempt`, `add_span`, etc.). This document plus the tests in `tests/store/` illustrate the expected behavior.
+
+## Thread Safety
+
+**[`LightningStoreThreaded`][agentlightning.LightningStoreThreaded]** is a subclass of [`LightningStore`][agentlightning.LightningStore] that wraps another underlying store to make a store instance safe for multi-threaded callers. It wraps every state-mutating call in a mutex. Specifically:
+
+* Methods like [`start_rollout`][agentlightning.LightningStore.start_rollout], [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout], [`update_attempt`][agentlightning.LightningStore.update_attempt], [`add_span`][agentlightning.LightningStore.add_span], etc. are guarded by a lock.
+* Non-mutating, potentially blocking calls remain pass-through by design (e.g., [`wait_for_rollouts`][agentlightning.LightningStore.wait_for_rollouts]), as they don’t modify shared state and should not hold the lock for long periods.
+
+## Process Safety and Client-server Store
+
+**[`LightningStoreServer`][agentlightning.LightningStoreServer]** wraps another underlying store and runs a FastAPI app to expose the store API over HTTP. [`LightningStoreClient`][agentlightning.LightningStoreClient] is a small [`LightningStore`][agentlightning.LightningStore] implementation that talks to the HTTP API.
+
+!!! warning
+
+    The server HTTP API is not considered a stable API at this moment. Users are encouraged to use the [`LightningStoreClient`][agentlightning.LightningStoreClient] to communicate with the server as a stable interface.
+
+The server tracks the creator PID. In the owner process it delegates directly to the in-memory store; in other processes it lazily constructs a [`LightningStoreClient`][agentlightning.LightningStoreClient] to talk to the HTTP API. This prevents accidental cross-process mutation of the wrong memory image. When the server is pickled (e.g., via `multiprocessing`), only the minimal fields are serialized, but **NOT** the FastAPI/uvicorn objects. Subprocesses won’t accidentally carry live server state. Forked subprocess should also use [`LightningStoreClient`][agentlightning.LightningStoreClient] to communicate with the server in the main process.
+
+On the client side, the client retries network/5xx failures using a small backoff, and probes `/health` between attempts. Application exceptions inside the server are wrapped as HTTP 400 with a traceback—these are **not retried**. The client also maintains a **per-event-loop** `aiohttp.ClientSession` map so that tracer callbacks (often on separate loops/threads) don’t hang by reusing a session from another loop.
+
+Minimal lifecycle:
+
+```python
+import agentlightning as agl
+
+# Server (owner process)
+in_memory_store = agl.InMemoryLightningStore()
+server = agl.LightningStoreServer(store=in_memory_store, host="0.0.0.0", port=4747)
+await server.start()      # starts uvicorn in a daemon thread and waits for /health
+# or keep your own event loop and stop via await server.stop()
+# await server.run_forever()
+
+# Client (same or different process)
+client = agl.LightningStoreClient("http://localhost:4747")
+
+print(await client.query_rollouts(status=["queuing"]))
+
+await client.close()
+await server.stop()
+```
+
+Another approach is to use a dedicated command line to start a long running server process, possibly sharable across multiple processes. In the main process, you can always use [`LightningStoreClient`][agentlightning.LightningStoreClient] to communicate with the server.
+
+```bash
+agl store --port 4747
+```
+
+!!! note
+
+    [`LightningStoreClient.wait_for_rollouts`][agentlightning.LightningStoreClient.wait_for_rollouts] intentionally enforces a tiny timeout (≤ 0.1s) to avoid blocking event loops. Poll with short timeouts or compose with `asyncio.wait_for` at a higher layer.
@@ -0,0 +1,229 @@
+# Train the First Agent with Agent-lightning
+
+Welcome! This tutorial is your first step into making AI agents smarter using the **Agent-lightning** framework. We'll show you how to take a simple agent and automatically improve its performance through a process called [**Automatic Prompt Optimization (APO)**](../algorithm-zoo/apo.md).
+
+The main goal of Agent-lightning is to provide a structured way to **train your agents**. Just like you train a machine learning model on data, you can train an agent on a task dataset. This could involve using Reinforcement Learning (RL) to teach it new behaviors or, as we'll do today, optimizing its prompts to make it more accurate and reliable.
+
+!!! tip
+
+    You can open the sample code [room_selector_apo.py]({{ src("examples/apo/room_selector_apo.py") }}) and [room_selector.py]({{ src("examples/apo/room_selector.py") }}) as you go through this tutorial.
+
+## Our Example: The Room Selector Agent
+
+Today, we'll work with an agent whose job is to book a meeting room. It's a common but tricky task with multiple constraints.
+
+Here's how the agent works:
+
+- **Input:** It receives a task with specific requirements, like "`Find a room for 4 people at 10:00 AM with a whiteboard.`"
+- **Action:** The agent uses a Large Language Model (LLM) to understand the request. It can also use tools, which are pre-defined functions it can call, to get more information, such as checking room availability in an external database.
+- **Output:** Its final decision is the ID of the best room it found, like "`A103`".
+- **Reward:** After the agent makes its choice, a separate "grader" function scores its performance on a scale of 0 to 1. This score is called its **reward**. A perfect choice gets a 1.0, while a wrong one gets a 0.0.
+
+The agent's logic is sound, but its performance heavily depends on its initial prompt. A poorly worded prompt will confuse the LLM, leading to bad decisions. Our goal is to use Agent-lightning to find the best possible prompt automatically.
+
+!!! tip "A Closer Look at the Agent's Logic"
+
+    Modern LLMs can do more than just generate text; they can decide to call functions you provide. This is often called tool use or function calling. Our agent uses this capability to make informed decisions. If you're new to this concept, you can read more about it in [OpenAI's documentation](https://platform.openai.com/docs/guides/function-calling).
+
+    Here is a sketch of the agent's logic, adhering closely to the OpenAI API:
+
+    ```python
+    # Pseudo-code for the Room Selector agent
+
+    import openai
+    import json
+
+    def room_selector_agent(task, prompt):
+        client = openai.OpenAI()
+        messages = [{"role": "user", "content": prompt.format(**task)}]
+        tools = [ ... ] # Tool definition for the LLM
+
+        # 1. First LLM call to decide if a tool is needed.
+        response = client.chat.completions.create(
+            model="gpt-5-mini",
+            messages=messages,
+            tools=tools,
+            tool_choice="auto",
+        )
+        response_message = response.choices[0].message
+        tool_calls = response_message.tool_calls
+
+        # 2. Check if the LLM wants to use a tool.
+        if tool_calls:
+            messages.append(response_message) # Append assistant's reply
+
+            # 3. Execute the tool and get the real-world data.
+            for tool_call in tool_calls:
+                function_name = tool_call.function.name
+                if function_name == "get_rooms_and_availability":
+                    function_args = json.loads(tool_call.function.arguments)
+                    # Query the local room database
+                    function_response = get_rooms_and_availability(
+                        date=function_args.get("date"),
+                        time_str=function_args.get("time"),
+                        duration_min=function_args.get("duration_min"),
+                    )
+                    messages.append({
+                        "tool_call_id": tool_call.id,
+                        "role": "tool",
+                        "name": function_name,
+                        "content": json.dumps(function_response),
+                    })
+
+            # 4. Second LLM call with the tool's output to get a final choice.
+            second_response = client.chat.completions.create(
+                model="gpt-5-mini",
+                messages=messages,
+            )
+            final_choice = second_response.choices[0].message.content
+        else:
+            final_choice = response_message.content
+
+        # 5. Grade the final choice to get a reward.
+        reward = grade_the_choice(final_choice, task["expected_choice"])
+        return reward
+    ```
+
+In Agent-lightning, you wrap this logic in a Python function marked with the [`@rollout`][agentlightning.rollout] decorator, so that the agent can be managed and tuned by Agent-lightning's runner and trainer. The `prompt_template` that the APO algorithm tunes is passed in as an argument:
+
+```python
+import agentlightning as agl
+
+@agl.rollout
+def room_selector(task: RoomSelectionTask, prompt_template: agl.PromptTemplate) -> float:
+    # ... agent logic using the prompt_template ...
+
+    # The final reward is determined by a grader function
+    reward = room_selection_grader(client, final_message, task["expected_choice"])
+    return reward
+```
+
+## Core Concepts: Tasks, Rollouts, Spans, and Prompt Templates
+
+To understand how Agent-lightning works, you need to know these key terms.
+
+### Task
+
+A task is a specific input or problem statement given to the agent. It defines what the agent needs to accomplish.
+
+!!! example "Analogy: Task"
+
+    If the agent is a chef, a task is the recipe request: "Bake a chocolate cake."
+
+### Rollout
+
+A rollout is a single, complete execution of an agent attempting to solve a given **task**. It's the entire story from receiving the task to producing a final result and receiving a reward. A rollout captures a full trace of the agent's execution.
+
+!!! example "Analogy: Rollout"
+
+    A rollout is one full attempt by the chef to bake the chocolate cake, from gathering ingredients to the final taste test.
+
+### Span
+
+A span represents a single unit of work or an operation within a **rollout**. Spans are the building blocks of a trace. They have a start and end time and contain details about the specific operation, like an LLM call, a tool execution, or a reward calculation. For a more precise definition, see the [OpenTelemetry documentation](https://opentelemetry.io/docs/concepts/signals/traces/).
+
+!!! example "Analogy: Span"
+
+    If the rollout is "baking a cake," a span could be "preheating the oven," "mixing flour and sugar," or "adding frosting." Each is a distinct step or unit of work.
+
+The picture below from [ADK](https://google.github.io/adk-docs/observability/cloud-trace/) shows a typical rollout, where each rectangle in the waterfall visualizes a span. As can be seen in the visualization, spans can be sequential, parallel or nested among each other. In other frameworks, the terminilogy might be slightly different. Agent-lightning follows the terminalogies used by OpenTelemetry to avoid confusion.
+
+![AgentOps Waterfall Visualization](../assets/agentops-waterfall-visualization.jpg)
+
+### Prompt Template
+
+A prompt template is a reusable instruction for the agent, often containing placeholders that can be filled in with specific details from a task. It is a key **"resource"** that the algorithm learns and improves over time.
+
+!!! example "Analogy: Resource (Prompt Template)"
+
+    If the task is the recipe request, the prompt template is the master recipe card that the chef follows. The algorithm's job is to edit this recipe card to make the instructions clearer and the final dish better.
+
+## The Training Loop: How the Magic Happens
+
+Training in Agent-lightning revolves around a clear, managed loop, orchestrated by the **Trainer**. The diagram below illustrates this core interaction:
+
+![Loop of Tasks and Spans](../assets/tasks-spans-loop.svg){ .center }
+
+**The Loop Explained:**
+
+- **Algorithm to Agent (via Trainer):** The **Algorithm** (the "brain") creates an improved **Prompt Template** and selects **Tasks**. The Trainer then sends both to the Trainer.
+- **Agent to Algorithm (via Trainer):** For each task it receives, the Agent uses the provided prompt template to perform a Rollout, executing its logic and potentially using tools. During this rollout, the runner that runs the agent captures Spans that detail every step. The agent also calculates a Reward for its performance on the task. These spans and rewards are then sent back to the Algorithm via the Trainer.
+- **Algorithm Learning:** The Algorithm then analyzes these spans and rewards to learn how to improve the agent's behavior, for example, by generating a better prompt. This improved prompt is then used in the next iteration of tasks.
+
+This cycle continues, allowing the agent to continuously learn and get better at solving tasks.
+
+!!! note
+
+    In the next tutorial, we will see that the "via Trainer" here is not accurate. It's actually via the runner and store.
+
+### The Algorithm
+
+The algorithm is the smart part of the system that drives the improvement. In this tutorial, we use [**APO**][agentlightning.algorithm.apo.APO] (Automatic Prompt Optimization). It works in a few steps:
+
+1. **Evaluate:** The algorithm first asks for rollouts to be run using the current prompt template to see how well it performs.
+2. **Critique:** It then looks at the detailed spans from those rollouts. Using a powerful LLM (`gpt-5-mini`), it generates a "textual gradient", which is a natural language critique of the prompt. For example: "The prompt is ambiguous about how to handle tie-breakers for equally good rooms."
+3. **Rewrite:** Finally, it gives the critique and the original prompt to another LLM (`gpt-4.1-mini`) and asks it to apply the edits, generating a new, improved prompt template.
+
+This cycle repeats, with each round producing a slightly better prompt. To use it, you simply initialize the APO class with your desired hyperparameters.
+
+```python
+# In the main training script: run_apo.py
+from openai import AsyncOpenAI
+
+openai = AsyncOpenAI()
+algo = agl.APO(openai)
+```
+
+!!! tip
+
+    Make sure you have `OPENAI_API_KEY` set in your environment variables.
+
+### The Trainer
+
+The Trainer is the central component you'll interact with. It connects everything and manages the entire workflow by running the loop described above. You configure the Trainer, providing the algorithm, the number of parallel runners, and the initial prompt. A single call to [`trainer.fit()`][agentlightning.Trainer.fit] kicks off the entire process!
+
+```python
+# 1. Configure the Trainer with the algorithm and initial prompt
+trainer = agl.Trainer(
+    algorithm=algo,
+    n_runners=8, # Run 8 agents in parallel to try out the prompts
+    initial_resources={
+        # The initial prompt template to be tuned
+        "prompt_template": prompt_template_baseline()
+    },
+    # This is used to convert the span data into a message format consumable by APO algorithm
+    adapter=agl.TraceToMessages(),
+)
+
+# 2. Load datasets: They can be list of task objects consumable by `room_selector`.
+dataset_train, dataset_val = ...
+
+# 3. Start the training process!
+trainer.fit(
+    agent=room_selector,
+    train_dataset=dataset_train,
+    val_dataset=dataset_val
+)
+```
+
+!!! tip
+
+    [`TraceToMessages`][agentlightning.TraceToMessages] is a convenience adapter that converts spans into OpenAI chat messages. It requires `openai >= 1.100.0` to be installed.
+
+## Training Results
+
+The APO algorithm successfully improved the agent's performance. We ran the example with the following hyper-parameters:
+
+* `val_batch_size` = 10
+* `gradient_batch_size` = 4
+* `beam_width` = 2
+* `branch_factor` = 2
+* `beam_rounds` = 2
+
+The validation accuracy on the 29 samples of datasets steadily increase from 0.569 (baseline) to **0.721** (after round 2). The tuning takes around 10 minutes with 8 runners. We ran twice, and the results are shown in the chart below.
+
+<div style="height:400px">
+<canvas data-chart='{ "type": "line", "data": { "labels": ["Baseline", "After round 1", "After round 2"], "datasets": [ { "label": "Run #1", "data": [0.569, 0.638, 0.721], "spanGaps": true }, { "label": "Run #2", "data": [0.534, 0.628, 0.645], "spanGaps": true } ] }, "options": { "interaction": { "mode": "nearest", "intersect": false }, "plugins": { "legend": { "display": true, "position": "top" }, "title": { "display": true, "text": "Validation Accuracy Across Rounds" } }, "scales": { "x": { "title": { "display": true, "text": "Round" } }, "y": { "title": { "display": true, "text": "Accuracy" } } } } }'></canvas>
+</div>
+
+This demonstrates how Agent-lightning can efficiently and automatically enhance your agent's capabilities with just a few lines of code.
@@ -1,14 +1,14 @@
-# SQL Agent with Agent Lightning
+# Train SQL Agent with Agent-lightning and VERL

-> This tutorial is tested with `verl==0.5.0` and `vllm==0.10.0`.
+This walkthrough builds upon the **Agent-lightning v0.2 SQL Agent** example and explains how the system components integrate: a **LangGraph-based SQL agent** wrapped as a [`LitAgent`][agentlightning.LitAgent], the **[`VERL`][agentlightning.algorithm.verl.VERL] reinforcement learning (RL) algorithm**, and the **[`Trainer`][agentlightning.Trainer]**, which coordinates both training and debugging.

-This example demonstrates how to build and train a self-correcting SQL agent. It leverages [Agent Lightning]({{ config.repo_url }}) and the `verl` framework for Reinforcement Learning (RL) based training, and LangGraph to define the agent's complex, cyclical reasoning workflow. The goal is to fine-tune a Large Language Model (LLM) to accurately convert natural language questions into executable SQL queries.
+The command-line interface in [`examples/spider/train_sql_agent.py`]({{ src("examples/spider/train_sql_agent.py") }}) provides a complete runnable example. However, this document focuses on understanding the underlying architecture so you can effectively adapt the workflow to your own agents.

-## SQL Agent Implementation
+## SQL Agent Architecture

-The design of Agent-lightning **allows flexible integration with various agent frameworks**, including AutoGen, CrewAI, OpenAI Agent SDK, LangGraph, and more. It can also work without agent frameworks, allowing you to train an agent built from scratch with Python code. See [our example gallery]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/examples) for more details.
+Agent-lightning integrates seamlessly with various orchestration frameworks, including [Agent Framework](https://github.com/microsoft/agent-framework), [AutoGen](https://github.com/microsoft/autogen), [CrewAI](https://www.crewai.com/), [LangGraph](https://github.com/langchain-ai/langgraph), and the [OpenAI Agents SDK](https://github.com/openai/openai-agents-python). It can also interoperate with custom Python logic.

-The core of the agent is a state machine built with LangGraph, which allows for a robust and transparent workflow. The agent's logic, as visualized below, starts by writing a query, executes it, and then enters a refinement loop where it checks and rewrites the query until it is deemed correct or a turn limit is reached.
+In this example, **LangGraph** defines a cyclic workflow that mirrors an analyst’s iterative SQL development process. The following graph (rendered directly from [`sql_agent.py`]({{ src("examples/spider/sql_agent.py") }})) illustrates how the agent drafts, executes, critiques, and refines queries until a satisfactory result is achieved.

 ```mermaid
 ---
@@ -34,151 +34,323 @@ graph LR;
        classDef last fill:#cccccc
 ```

-This workflow is implemented in the `SQLAgent` class within `sql_agent.py`. It consists of the following key steps:
+!!! note

-1. **write_query**: Given a user's question and database schema, the agent makes an initial attempt to write a SQL query.
-2. **execute_query**: The generated query is run against the target database.
-3. **check_query**: The agent analyzes the original query and its execution result (or error) to check for mistakes. It uses a specific prompt (`CHECK_QUERY_PROMPT`) to determine if the query is correct.
-4. **rewrite_query**: If the `check_query` step finds errors, the agent enters this step. It uses the feedback from the previous step to generate a corrected SQL query. The process then loops back to `check_query` for re-evaluation.
-5. **END**: The loop terminates when `check_query` confirms the query is correct or the maximum number of turns (`max_turns`) is exceeded. One turn corresponds to a complete cycle of `write_query` (if first round), `execute_query`, `check_query`, and potentially `rewrite_query`.
+    The workflow proceeds through the following stages:

-We aim to train **write_query** and **rewrite_query** step in the setup of this example. The **check_query** step is not trained but will share the same LLM weights as the other steps.
+    1. **write_query** – Generates an initial SQL query from the user’s question and the database schema.
+    2. **execute_query** – Executes the generated query against the target database.
+    3. **check_query** – Evaluates the query and its results (or errors) using a specialized prompt (`CHECK_QUERY_PROMPT`) to detect issues.
+    4. **rewrite_query** – If issues are identified, the agent rewrites the query using feedback from the previous step and re-enters the loop.
+    5. **END** – The cycle terminates when the query is validated or the maximum iteration count (`max_turns`) is reached. Each *turn* consists of one full loop through the `write_query`, `execute_query`, `check_query`, and (if applicable) `rewrite_query` stages.

-## Client-Server Training with Agent Lightning
+In this tutorial, **reinforcement learning (RL)** is used to optimize the `write_query` and `rewrite_query` stages. While the `check_query` step shares the same underlying LLM weights, its trace data is not used for learning.

-The training process uses a distributed client-server architecture designed by Agent Lightning to efficiently fine-tune the underlying LLM. This separation allows for scalable data generation across multiple clients while centralizing the computationally intensive model training on a dedicated server with GPUs, and also provides opportunities for customizing algorithms and training strategies (like [prompt optimization]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/examples/apo)) with minimal code changes.
+To keep the design modular and maintainable, it is recommended to define the LangGraph-based SQL Agent in a separate file and expose it via a builder function such as:

-* **Training Server (`agentlightning.verl`)**: The server, launched with the first command below, manages the core training loop. It runs an RL algorithm (with `verl` of course) and hosts an OpenAI-compatible LLM endpoint (with `verl`'s async server). The server's sole purpose is to receive interaction data from clients and update the LLM's weights to improve its performance. [This link]({{ config.repo_url }}/tree/{{ config.extra.source_commit }}/agentlightning/verl) points to the implementation of the server, which is built upon `verl`.
-* **Agent Clients (`sql_agent.py`)**: The clients run the LangGraph agent logic described above. They connect to the server to fetch tasks (natural language questions) and use the server's **OpenAI-compatible endpoint** for all generation steps (`write_query`, `check_query`, `rewrite_query`). After completing a task, the client exports its interaction traces (traced by [AgentOps](https://www.agentops.ai/) and filtered by trace hierarchy), evaluates its correctness to calculate a reward, and sends the entire interaction history (the "trajectory") back to the server for training. To adapt any agent to an "agent client", you do not need to change the agent logic, but only need to invoke the client's `run` method with `agentlightning.trainer`.
+```python
+def build_langgraph_sql_agent(
+    database_path: str,
+    openai_base_url: str,
+    model: str,
+    sampling_parameters: Dict[str, Any],
+    max_turns: int,
+    truncate_length: int
+):
+    builder = StateGraph(State)
+    builder.add_node(write_query)
+    ...

-![Difference between the original agent and modified agent client](../assets/sql-agent-diff.png)
+    builder.add_edge(START, "write_query")
+    ...

-## Running the Example
-
-1. Prepare the dataset: download from [here](https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view) and unzip it to the `data` folder. It's basically a [Spider V1](https://yale-lily.github.io/spider) dataset converted to Parquet format. The dataset contains about 8000 training samples and about 2000 test samples, from which we sampled 500 samples for evaluation.
-   ```bash
-   pip install gdown
-   gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
-   unzip -q spider-data.zip -d data
-   rm spider-data.zip
-   ```
-
-2. Install the required dependencies:
-   ```bash
-   pip install -r requirements.txt
-   ```
-
-3. Launch the training server:
-   ```bash
-   python -m agentlightning.verl \
-       agentlightning.port=9997 \
-       algorithm.adv_estimator=grpo \
-       data.train_files=data/train_spider.parquet \
-       data.val_files=data/test_dev_500.parquet \
-       actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
-       trainer.n_gpus_per_node=1 \
-       data.train_batch_size=32 \
-       actor_rollout_ref.rollout.n=4 \
-       actor_rollout_ref.actor.ppo_mini_batch_size=32 \
-       actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
-       actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
-       actor_rollout_ref.rollout.multi_turn.format=hermes \
-       actor_rollout_ref.model.path=meta-llama/Llama-3.2-3B-Instruct \
-       data.max_prompt_length=4096 \
-       data.max_response_length=2048 \
-       data.truncation='error' \
-       trainer.val_before_train=True \
-       actor_rollout_ref.actor.optim.lr=1e-6 \
-       actor_rollout_ref.model.use_remove_padding=True \
-       actor_rollout_ref.actor.use_kl_loss=False \
-       actor_rollout_ref.actor.kl_loss_coef=0.000 \
-       actor_rollout_ref.actor.entropy_coeff=0 \
-       actor_rollout_ref.actor.clip_ratio_low=0.2 \
-       actor_rollout_ref.actor.clip_ratio_high=0.3 \
-       actor_rollout_ref.model.enable_gradient_checkpointing=True \
-       actor_rollout_ref.actor.fsdp_config.param_offload=True \
-       actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
-       actor_rollout_ref.rollout.name=vllm \
-       actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
-       actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
-       actor_rollout_ref.ref.fsdp_config.param_offload=True \
-       algorithm.use_kl_in_reward=False \
-       trainer.critic_warmup=0 \
-       trainer.logger=['console','wandb'] \
-       trainer.project_name=AgentLightning \
-       trainer.experiment_name=train_sql_agent \
-       trainer.nnodes=1 \
-       trainer.save_freq=256 \
-       trainer.test_freq=32 \
-       trainer.total_epochs=2
-    ```
-
-4. Launch agent clients that connect with the server:
-   ```bash
-   export VERL_API_BASE=http://localhost:9997/  # Same as the server port. This is used for receiving tasks and sending results.
-   python sql_agent.py \
-       --litsqlagent.trained-agents write \  # Will only train the write and rewrite agent.
-       --trainer.n-workers 16 \
-       --litsqlagent.val-temperature 0
-   ```
-
-There is no hard requirement in the launching order of the server and clients. But remember to kill the long-running agent clients after the training is done.
-
-## Debug the Agent without verl
-
-You can run the agent client alone without the `verl` server. This is useful for debugging the agent logic and SQL execution.
-
-1. Copy `.env.example` to `.env` and fill in your OpenAI API key. `VERL_API_BASE` does not really matter here because you are not connecting to the server end.
-
-2. Run the agent client:
-   ```bash
-   dotenv run python sql_agent.py \
-       --litsqlagent.trained-agents write \  # Will only select the trajectories related to write and rewrite.
-       --trainer.n-workers 1 \  # For debug, use single process.
-       --trainer.dev true  # Enable the dev debug mode.
-   ```
-
-## Evaluation
-
-The example is evaluated using Llama-3.2-Instruct models. The models are trained on the Spider dataset for 2 epochs, with evaluation performed on a randomly selected subset of 500 test samples to compute held-out accuracy. The default setup for running agent clients during evaluation is as follows:
-
-```bash
-python sql_agent.py \
-   --litsqlagent.trained-agents write \
-   --trainer.n-workers 16 \
-   --trainer.daemon true \
-   --litsqlagent.val-temperature 0 \
-   --litsqlagent.max-turns 3 \
-   --litsqlagent.table-info-truncate 2048 \
-   --litsqlagent.execution-truncate 2048
+    return builder.compile().graph()
 ```

-The setup of training server is the same as the command above.
+This approach isolates your LangGraph logic from Agent-lightning version changes, improving both readability and debuggability.

-### W&B Report
+## Bridging LangGraph and Agent-lightning

-[link](https://api.wandb.ai/links/ultmaster/4cid500g)
+!!! tip

-### Performance Metrics
+    Keep [`sql_agent.py`]({{ src("examples/spider/sql_agent.py") }}) open on the side while reading this section. This will help you understand how the code snippets shown here work in practice.

-![](../assets/sql-agent-val-reward-curve.png)
+The **`LitSQLAgent`** class defined in [`sql_agent.py`]({{ src("examples/spider/sql_agent.py") }}) acts as the bridge. It subclasses [`agl.LitAgent`][agentlightning.LitAgent], allowing the runner to provision shared resources (e.g., [LLMs][agentlightning.LLM]) for each rollout.

-| Model         | Size   |   Context |   Max Turns | Agents                        |   Acc (Initial) |   Acc (Final) | Transitions   |   Prompt Length | Response Length   |
-|---------------|--------|-----------|-------------|-------------------------------|-----------------|---------------|---------------|-----------------|-------------------|
-| Llama3.2      | 1B     |      2048 |           3 | write&#124;rewrite            |            21   |          49.6 | 2.87 → 3.08   |           821.2 | 319.2 → 249.4     |
-| Llama3.2      | 3B     |      2048 |           3 | write&#124;rewrite            |            51.8 |          66.4 | 2.20 → 2.72   |           865.6 | 116.2 → 314.3     |
+Below is a simplified illustration of the key logic (note: this is conceptual pseudocode; the actual implementation includes dataset-specific details):

-**Notes:**
+```python
+class LitSQLAgent(agl.LitAgent[Dict[str, Any]]):

-1. **Context Length**: Controlled via `--litsqlagent.table-info-truncate <context-length>` and `--litsqlagent.execution-truncate <context-length>`
-2. **Max Turns**: Set using `--litsqlagent.max-turns <max-turns>`
-3. **Agents**: Specified with `--litsqlagent.agents <regex>` (defaults to `write`, which matches both write and rewrite agents)
-4. **Transitions**: Represents the number of prompt-response pairs traced (collected) during each rollout. Note that this differs from the turn count in the SQL agent workflow, where one turn may encompass 2-3 transitions in the check-rewrite cycle. The number of transitions is also related to which *agents* get involved in the training.
-5. **Prompt/Response Length**: Average token count per **traced** prompt/transition response.
+    def __init__(self, max_turns: int, truncate_length: int):
+        # Every turn here refers to a full cycle of write/exe/check/rewrite
+        self.max_turns = max_turns
+        self.truncate_length = truncate_length

-### Efficiency Metrics
+    def rollout(
+        self,
+        task: Dict[str, Any],
+        resources: agl.NamedResources,
+        rollout: agl.Rollout
+    ) -> float | None:
+        llm: agl.LLM = resources["main_llm"]
+        agent = build_langgraph_sql_agent(
+            database_path="sqlite:///" + task["db_id"],
+            max_turns=self.max_turns,
+            truncate_length=self.truncate_length,
+            openai_base_url=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+            model=llm.model,
+            sampling_parameters=llm.sampling_parameters,
+        )
+        result = agent.invoke({"question": question}, {
+            "callbacks": [self.tracer.get_langchain_handler()],
+            "recursion_limit": 100,
+        })
+        reward = evaluate_query(result["query"], ground_truth, db_path, raise_on_error=False)
+        return reward
+```

-| Model         | Size   |   Context |   Max Turns | Agents                        |   # GPUs |   # Steps |   Time (h) |   Time/Step (s) |   Rollout Time (%) |   Update Actor Time (%) |
-|---------------|--------|-----------|-------------|-------------------------------|----------|-----------|------------|-----------------|--------------------|-------------------------|
-| Llama3.2      | 1B     |      2048 |           3 | write&#124;rewrite            |        1 |       436 |      13.06 |            98.9 |               66.7 |                    25.2 |
-| Llama3.2      | 3B     |      2048 |           3 | write&#124;rewrite            |        2 |       436 |      10.3  |           181.3 |               63.9 |                    27.9 |
+The `LitSQLAgent` serves as a lightweight wrapper around the LangGraph agent, providing the correct interface for the [`rollout`][agentlightning.LitAgent.rollout] method. It constructs the LangGraph agent, invokes it, and returns the evaluation result as a reward signal.
+
+The `"main_llm"` resource key is a convention between the agent and [VERL][agentlightning.algorithm.verl.VERL]. It is used to inject an OpenAI-compatible endpoint from the [VERL][agentlightning.algorithm.verl.VERL] algorithm during rollout. Two approaches are supported to use this [agentlightning.LLM][] resource:
+
+1. **Direct access** – Use [`llm.endpoint`][agentlightning.LLM.endpoint] for a simple integration (identical to the v0.1 example).
+2. **Context-aware access** – Use [`get_base_url`][agentlightning.ProxyLLM.get_base_url] with [`rollout.rollout_id`][agentlightning.Rollout.rollout_id] and [`rollout.attempt.attempt_id`][agentlightning.Attempt.attempt_id].
+   This approach enables per-caller trace attribution, improving trace collection per rollout or attempt when runner-side tracers are unavailable. For details, see [Working with Traces](../tutorials/traces.md).
+
+## Reward Signal and Evaluation
+
+The `evaluate_query` function provides the reward mechanism for RL training. In agent training, obtaining a consistent and meaningful reward signal is often challenging. Fortunately, this is simplified when using the [**Spider dataset**](https://yale-lily.github.io/spider). The dataset includes ~8k samples containing natural-language questions, database schemas, and ground-truth SQL queries.
+
+Using the [**Spider evaluator**](https://github.com/taoyds/test-suite-sql-eval), the agent's generated query is executed and compared to the ground-truth query on the target database. The two queries are considered equivalent if they produce identical execution results.
+
+!!! attention
+
+    The ground-truth queries must **never** be exposed to the agent during training to prevent data leakage.
+
+In this setup, the reward is returned directly from the [`rollout`][agentlightning.LitAgent.rollout] method, enabling the runner to forward it back to the RL algorithm.
+
+!!! warning
+
+    Avoid using [`emit_reward`][agentlightning.emit_reward] in conjunction with returning a reward value. Doing both will cause the algorithm to receive duplicate reward signals, leading to inconsistent training behavior.
+
+## Configuring VERL for Reinforcement Learning
+
+View [`examples/spider/train_sql_agent.py`]({{ src("examples/spider/train_sql_agent.py") }}) for a full reinforcement learning configuration, which is a plain Python dictionary. It mirrors (and actually *is*) the [shell arguments](https://verl.readthedocs.io/en/latest/index.html) used to launch training in the VERL framework but is easier to tweak programmatically:
+
+```python
+verl_config: Dict[str, Any] = {
+    "algorithm": {"adv_estimator": "grpo", "use_kl_in_reward": False},
+    "data": {
+        # train_files and val_files are no longer needed here
+        # because data are read in agl.Trainer
+        ...,
+        # Controls how many tasks are pooled per step
+        # (multiplied by actor_rollout_ref.rollout.n)
+        "train_batch_size": 32,
+        # Prompt and responses larger than these lengths are truncated
+        "max_prompt_length": 4096,
+        "max_response_length": 2048,
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            # Only vLLM is supported currently
+            "name": "vllm",
+            # Equals to group size of GRPO
+            "n": 4,
+            # Used to enable tool call parser in vLLM
+            "multi_turn": {"format": "hermes"},
+            ...
+        },
+        "actor": {"ppo_mini_batch_size": 32, "optim": {"lr": 1e-6}, ...},
+        "model": {
+            # Config your preferred LLM here
+            "path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+            ...
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        # Validation once before training starts
+        "val_before_train": True,
+        # Validation every N training steps
+        "test_freq": 32,
+        # Save checkpoints every N training steps
+        "save_freq": 64,
+        # Go through the train dataset this many times
+        "total_epochs": 2
+    },
+}
+```
+
+This is equivalent to the following CLI invocation:
+
+```bash
+python3 -m verl.trainer.main_ppo \
+    algorithm.adv_estimator=grpo \
+    algorithm.use_kl_in_reward=False \
+    data.train_batch_size=32 \
+    data.max_prompt_length=4096 \
+    data.max_response_length=2048 \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.n=4 \
+    actor_rollout_ref.rollout.multi_turn.format=hermes \
+    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.path=Qwen/Qwen2.5-Coder-1.5B-Instruct \
+    trainer.n_gpus_per_node=1 \
+    trainer.val_before_train=True \
+    trainer.test_freq=32 \
+    trainer.save_freq=64 \
+    trainer.total_epochs=2
+```
+
+!!! warning
+    We used to provide a CLI called `python -m agentlightning.verl` to launch training in v0.1. This is no longer the recommended approach. Instead, use [`agl.Trainer`][agentlightning.Trainer] to run VERL and agent runners together, or follow the [debugging tutorial](../tutorials/debug.md) if you want an isolated experience similar to v0.1.
+## Orchestrating Training with [`Trainer`][agentlightning.Trainer]
+
+[`Trainer`][agentlightning.Trainer] is the high-level orchestrator that integrates the agent, algorithm, dataset, and distributed runners. The key benefits of using the [`Trainer`][agentlightning.Trainer] are:
+
+1. It allows you to launch everything with a single line of code: `trainer.fit(...)`.
+2. It exposes configuration options such as `n_runners` to control parallelism and `adapter` to define how algorithms interpret the trace data produced by the agent.
+
+An example usage is shown below:
+
+```python
+import agentlightning as agl
+
+agent = LitSQLAgent()
+algorithm = agl.VERL(verl_config)
+trainer = agl.Trainer(
+    n_runners=10,
+    algorithm=algorithm,
+    adapter={"agent_match": active_agent},
+)
+train_data = pd.read_parquet("data/train_spider.parquet").to_dict("records")
+val_data = pd.read_parquet("data/test_dev_500.parquet").to_dict("records")
+trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)
+```
+
+First, `agl.VERL(verl_config)` launches the [`VERL`][agentlightning.algorithm.verl.VERL] algorithm and its OpenAI-compatible proxy. The `train_data` and `val_data` are passed into [`VERL`][agentlightning.algorithm.verl.VERL], which enqueues tasks to a centralized task queue managed by the [`LightningStore`][agentlightning.LightningStore], accessible to all runners.
+
+When [`Trainer.fit`][agentlightning.Trainer.fit] is called, it launches 10 concurrent runners (as specified by `n_runners=10`). Each runner pulls tasks from the centralized task queue, executes the agent’s [`rollout`][agentlightning.LitAgent.rollout] method, collects traces, and returns rewards to VERL for training.
+
+The [`Adapter`][agentlightning.Adapter], as discussed earlier, is used at the algorithm side, and receives the traces emitted by the agent and runners. The `agent_match` parameter ensures [`VERL`][agentlightning.algorithm.verl.VERL] only ingests spans from the specific agent you want to optimize.
+In the example above, there are at least three agents—`write_query`, `rewrite_query`, and `check_query`. By setting `agent_match` to a regex like `"write"`, both `write_query` and `rewrite_query` agents are optimized simultaneously. You can also set it to `"write|check"` or `None` to include all agents if desired.
+
+## Dry-Run the Pipeline with [`Trainer.dev`][agentlightning.Trainer.dev]
+
+Before committing hours of GPU time, you can **dry-run** the agent with [`Trainer.dev()`][agentlightning.Trainer.dev]. This method swaps in the lightweight [`Baseline`][agentlightning.Baseline] algorithm, enqueues up to ten tasks, and prints every span emitted by the agent. Because it uses the same runner stack as full training, it’s ideal for verifying database connections and LangGraph control flow.
+
+To begin, the agent needs a valid OpenAI-compatible endpoint since VERL is not active in this mode. You can use OpenAI’s official API or your own local LLM endpoint. Wrap it as follows:
+
+```python
+trainer = agl.Trainer(
+    n_workers=1,
+    initial_resources={
+        "main_llm": agl.LLM(
+            endpoint=os.environ["OPENAI_API_BASE"],
+            model="gpt-4.1-nano",
+            sampling_parameters={"temperature": 0.7},
+        )
+    },
+)
+```
+
+Then, call [`trainer.dev(...)`][agentlightning.Trainer.dev] with a small number of tasks:
+
+```python
+dev_data = pd.read_parquet("data/test_dev_500.parquet").to_dict("records")[:10]
+trainer.dev(agent, dev_dataset=dev_data)
+```
+
+Run this in a Python session or adapt your script to include a `--dev` flag. Once the spans appear healthy and the rewards are non-zero, switch back to [`trainer.fit(...)`][agentlightning.Trainer.fit] for full RL training.
+
+## Running the Sample Code
+
+The following tutorial explains how to run the complete example in [`examples/spider`]({{ src("examples/spider") }}).
+
+### Dataset
+
+The trainer expects three Parquet files inside `examples/spider/data`:
+`train_spider.parquet`, `test_dev_500.parquet`, and `test_dev.parquet`.
+
+Download the curated dataset bundle provided with the repository:
+
+```bash
+cd examples/spider
+pip install gdown  # included in the 'experiment' optional dependency
+gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
+unzip -q spider-data.zip -d data
+rm spider-data.zip
+```
+
+If you prefer to generate the files yourself, download [Spider 1.0](https://yale-lily.github.io/spider) and run:
+
+```bash
+python spider_eval/convert_dataset.py
+```
+
+Set `VERL_SPIDER_DATA_DIR` if you store the dataset outside the default `data` directory.
+
+### Dependencies
+
+Create a clean virtual environment, activate it, and install Agent-lightning with the VERL extras required by [this tutorial](../tutorials/installation.md). Install LangChain-related dependencies as needed.
+
+For full training profiles, plan to use a GPU with at least **40 GB** of memory.
+
+### Launch Training
+
+From [`examples/spider`]({{ src("examples/spider") }}), run one of the helper scripts depending on your model preference:
+
+```bash
+python train_sql_agent.py qwen   # Default Qwen-2.5-Coder-1.5B run
+python train_sql_agent.py llama  # LLaMA-3.2-1B with llama3_json tool parser
+```
+
+The script instantiates `LitSQLAgent` and launches [`trainer.fit`][agentlightning.Trainer.fit].
+Provide `--active-agent my_agent_variant` if you only want to train one of the agents in the graph.
+
+For the LLaMA profile, export an `HF_TOKEN` before running so VERL can download the model weights.
+
+!!! tip "Troubleshooting"
+
+    If you have got some Ray worker errors on either `WANDB_API_KEY` not set, or `HF_TOKEN` not set, or data not found, please try to restart the Ray cluster with the helper script: [scripts/restart_ray.sh]({{ src("scripts/restart_ray.sh") }}), which essentially stops the ray cluster if any, and starts a new one:
+
+    ```bash
+    env RAY_DEBUG=legacy HYDRA_FULL_ERROR=1 VLLM_USE_V1=1 ray start --head --dashboard-host=0.0.0.0
+    ```
+
+### Debugging the Agent without VERL
+
+[`sql_agent.py`]({{ src("examples/spider/sql_agent.py") }}) also provides a `debug_sql_agent()` helper to run the LangGraph workflow directly against a local or hosted OpenAI-compatible endpoint before using VERL.
+
+Set the following environment variables, then execute the file:
+
+```bash
+export OPENAI_API_BASE=<your_api_base>
+export OPENAI_API_KEY=<your_api_key>
+cd examples/spider
+python sql_agent.py
+```
+
+This allows you to verify that the workflow and prompts behave as expected before reinforcement learning is introduced.
+
+### Evaluation
+
+The following results were obtained by running `python train_sql_agent.py qwen` on a single 80 GB GPU.
+Training completes in approximately **12 hours**.
+The training curves below are smoothed by aggregating every 16 steps for better visualization.
+
+Additional evaluation results were collected with a legacy version — Agent-lightning v0.1.1, `verl==0.5.0`, and `vllm==0.10.0`.
+You can find them in this write-up:
+[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)
+
+
+<div style="height:400px">
+<canvas data-chart='{"type": "line", "data": {"labels": [0.0, 16.0, 32.0, 48.0, 64.0, 80.0, 96.0, 112.0, 128.0, 144.0, 160.0, 176.0, 192.0, 208.0, 224.0, 240.0, 256.0, 272.0, 288.0, 304.0, 320.0, 336.0, 352.0, 368.0, 384.0, 400.0, 416.0, 432.0], "datasets": [{"label": "Training", "data": [0.4609375, 0.5041666666666667, 0.5790441176470589, 0.6015625, 0.6070772058823529, 0.6208333333333333, 0.6668198529411765, 0.66875, 0.6709558823529411, 0.6708333333333333, 0.6847426470588235, 0.6791666666666667, 0.6819852941176471, 0.690625, 0.7008272058823529, 0.7453125, 0.7398897058823529, 0.7119791666666667, 0.7224264705882353, 0.7114583333333333, 0.7431066176470589, 0.7427083333333333, 0.75, 0.7302083333333333, 0.7247242647058824, 0.7390625, 0.7463235294117647, 0.7376302083333334], "spanGaps": true}, {"label": "Validation", "data": [0.342, null, 0.594, null, 0.642, null, 0.66, null, 0.676, null, 0.676, null, 0.694, null, 0.712, null, 0.702, null, 0.678, null, 0.702, null, 0.702, null, 0.674, null, 0.734, 0.722], "spanGaps": true}]}, "options": {"interaction": {"mode": "nearest", "intersect": false}, "plugins": {"legend": {"display": true, "position": "top"}, "title": {"display": true, "text": "SQL Agent Training Result (agent_match = write)"}}, "scales": {"x": {"title": {"display": true, "text": "Step (aggregated)"}}, "y": {"title": {"display": true, "text": "Accuracy"}}}}}'></canvas>
+</div>
+
+<div style="height:400px">
+<canvas data-chart='{"type": "line", "data": {"labels": [0.0, 16.0, 32.0, 48.0, 64.0, 80.0, 96.0, 112.0, 128.0, 144.0, 160.0, 176.0, 192.0, 208.0, 224.0, 240.0, 256.0, 272.0, 288.0, 304.0, 320.0, 336.0, 352.0, 368.0, 384.0, 400.0, 416.0, 432.0], "datasets": [{"label": "Training", "data": [0.4560546875, 0.578125, 0.6167279411764706, 0.6401041666666667, 0.6461397058823529, 0.6598958333333333, 0.6838235294117647, 0.69375, 0.6916360294117647, 0.6833333333333333, 0.6893382352941176, 0.6921875, 0.6838235294117647, 0.70625, 0.7045036764705882, 0.7442708333333333, 0.7288602941176471, 0.7317708333333334, 0.7311580882352942, 0.7286458333333333, 0.7316176470588235, 0.7359375, 0.7366727941176471, 0.7208333333333333, 0.7118566176470589, 0.7296875, 0.7389705882352942, 0.7350260416666666], "spanGaps": true}, {"label": "Validation", "data": [0.33, null, 0.62, null, 0.662, null, 0.682, null, 0.696, null, 0.7, null, 0.708, null, 0.692, null, 0.72, null, 0.7, null, 0.7, null, 0.702, null, 0.694, null, 0.702, 0.682], "spanGaps": true}]}, "options": {"interaction": {"mode": "nearest", "intersect": false}, "plugins": {"legend": {"display": true, "position": "top"}, "title": {"display": true, "text": "SQL Agent Training Result (agent_match = null)"}}, "scales": {"x": {"title": {"display": true, "text": "Step (aggregated)"}}, "y": {"title": {"display": true, "text": "Value"}}}}}'></canvas>
+</div>
@@ -0,0 +1,405 @@
+# Fine-tune with Unsloth SFT
+
+!!! note "Prerequisites"
+
+    Please make sure you have read [Write the First Algorithm](./write-first-algorithm.md). Although that recipe is based on a simple prompt tuning algorithm, it introduces the core concepts of Agent-lightning and you should be familiar with them before proceeding.
+
+This recipe builds on [Write the First Algorithm](./write-first-algorithm.md). Instead of iterating on a prompt, we will fine-tune a large language model with [Unsloth](https://docs.unsloth.ai/)'s SFT Trainer and keep the whole loop inside Agent-lightning. The new pieces you will meet are the **LLM proxy**, the **trace-to-triplet adapter**, a [vLLM](https://github.com/vllm-project/vllm) inference endpoint, and an agent implemented with the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/). The full sample code is available in the [`examples/unsloth`]({{ src("examples/unsloth") }}) folder.
+
+!!! warning
+
+    You need a GPU that can host the Unsloth base model and run vLLM. The sample defaults to `unsloth/Qwen3-4B-Instruct-2507`, which requires at least 16GB of GPU memory under 4-bit quantization.
+
+## The Data and Serving Loop
+
+To tune a large language model in Supervised Fine-Tuning (SFT), we commonly need a dataset with input/output samples. For example, the [TRL SFT Trainer](https://huggingface.co/docs/trl/sft_trainer) expects a dataset with samples like the following:
+
+```json
+{"messages": [{"role": "user", "content": "What color is the sky?"},
+              {"role": "assistant", "content": "It is blue."}]}
+```
+
+With supervised fine-tuning, the LLM learns to generate the "assistant" response as close as possible to the completion in the dataset.
+
+Typically, the dataset used in SFT should be a curated set of samples. The samples can be either hand-written by humans, or generated by a more powerful model, which is known as [data distillation](https://docs.nvidia.com/nemo-framework/user-guide/24.12/modelalignment/knowledge-distillation.html). However, in this recipe, we use a different setup that relies on samples generated by the model itself. We use the reward emitted by the agent to select the top-performing samples.
+
+Overall, the flow of the algorithm is an iteration of the following steps:
+
+1. Serve the current checkpoint (with vLLM).
+2. Publish the vLLM endpoint through the LLM proxy and let runners roll out some tasks with the current model.
+3. Collect the traces from the rollouts and transform the highest-rewarded ones into a dataset that is acceptable for Unsloth SFT Trainer.
+4. Launch Unsloth to fine-tune on the dataset and save a new checkpoint.
+
+You will find the full source code of this iteration in `sft_one_iter` in [sft_algorithm.py]({{ src("examples/unsloth/sft_algorithm.py") }}). We will elaborate on each part below.
+
+### Serving the Model with vLLM and Proxy
+
+Most modern agents do not use the model directly; instead, they use an API like the OpenAI chat completions API to interact with the model. Therefore, we need a vLLM-based inference server launched before rollouts. The serving code looks like the following. See the `vllm_server` function in [sft_algorithm.py]({{ src("examples/unsloth/sft_algorithm.py") }}) if you want to see a more robust version.
+
+```python
+from openai import OpenAI
+
+vllm_process = subprocess.Popen([
+    "vllm", "serve", model_path, "--port", str(port),
+    "--enable-auto-tool-choice", "--tool-call-parser", "hermes"
+])
+
+# Wait for the server to be ready
+url = f"http://localhost:{port}/health"
+start = time.time()
+client = httpx.Client()
+
+while True:
+    if client.get(url).status_code == 200:
+        break
+
+server_address = f"http://localhost:{port}/v1"
+
+# Try using the vLLM server
+openai = OpenAI(base_url=server_address)
+...
+```
+
+In this recipe, we do not expose the server address directly to the agent runners, because we want to install a "middleware" to collect the prompts and responses of all the requests. In general, it's up to you to decide whether to hide the vLLM server behind a proxy or not.
+
+The "middleware" here is [`LLMProxy`][agentlightning.LLMProxy], which is an independent [LiteLLM](https://docs.litellm.ai/) server that forwards the requests to the vLLM server. It also exposes an OpenAI-compatible API that the runners can target without caring about where the model lives. The benefits of using the proxy are:
+
+1. **Traces:** The proxy automatically logs the prompts and responses of all the requests into the store.
+2. **Token IDs:** The proxy augments the requests so that the vLLM server can return the prompt and response token IDs (see more details in [Serving LLM](../deep-dive/serving-llm.md)).
+
+The [`LLMProxy`][agentlightning.LLMProxy] accepts a list of model configurations, in the same syntax as LiteLLM's [`model_list`](https://docs.litellm.ai/docs/proxy/configs). Include a `hosted_vllm/` prefix to the models to activate LiteLLM's [vLLM integration](https://docs.litellm.ai/docs/providers/vllm).
+
+```python
+import agentlightning as agl
+
+llm_proxy = agl.LLMProxy(port=port, store=store)
+model_list = [
+    {
+        "model_name": "Qwen3-4B-Instruct",
+        "litellm_params": {"model": f"hosted_vllm/{model_path}", "api_base": server_address},
+    }
+]
+llm_proxy.update_model_list(model_list)
+# If the proxy is not running, it will start automatically.
+llm_proxy.restart()
+# Add the proxy as a resource to the store so that the runners can access it via URL.
+resource_update = await store.add_resources({"main_llm": llm_proxy.as_resource()})
+```
+
+### Spawn Rollout and Collect Spans
+
+Once the proxy is registered as a resource, the algorithm schedules work for the rollout runners. Each problem from a training dataset becomes a rollout with the proxy baked into its resources:
+
+```python
+rollouts: list[Rollout] = []
+for sample in train_dataset:
+    rollouts.append(
+        await store.enqueue_rollout(
+            input=sample,
+            mode="train",
+            resources_id=resources_update.resources_id,
+        )
+    )
+```
+
+`resources_id` ties every rollout to the `main_llm` proxy resource we just uploaded. The runners on the other side poll the store ([`LitAgentRunner.iter()`][agentlightning.LitAgentRunner.iter]) and execute the agent for each rollout. On the algorithm side we wait for completions with a non-blocking polling loop:
+
+```python
+completed_rollouts: list[Rollout] = []
+while True:
+    completed_rollouts = await store.wait_for_rollouts(
+        rollout_ids=[r.rollout_id for r in rollouts],
+        timeout=0.0,
+    )
+    if len(completed_rollouts) == len(rollouts):
+        break
+    await asyncio.sleep(5.0)
+```
+
+!!! note
+
+    The `timeout=0.0` is needed here because this example uses a [`LightningStoreClient`][agentlightning.LightningStoreClient], and `wait_for_rollouts` establishes an HTTP connection to that store. Currently, only non-blocking wait requests are supported, which avoids holding the store connection open.
+
+Once the rollouts complete, we terminate the vLLM server to free up GPU memory.
+
+```python
+vllm_process.terminate()
+vllm_process.join(timeout=10.0)
+```
+
+### Adapt the Spans to HuggingFace Dataset
+
+[`LlmProxyTraceToTriplet`][agentlightning.LlmProxyTraceToTriplet] converts the proxy’s spans (which might be dozens to hundreds per rollout) into [`Triplet`][agentlightning.Triplet] objects that contain prompt/response token IDs plus an optional reward. The adapter may return multiple triplets per rollout (one per chat-completion call). To bias training toward successful reasoning chains the algorithm walks the triplets in reverse order, keeps the most recent reward, and turns each prompt/response pair into Hugging Face dataset rows:
+
+```python
+all_triplets = []
+data_adapter = agl.LlmProxyTraceToTriplet()
+
+for rollout in completed_rollouts:
+    spans = await store.query_spans(rollout.rollout_id, "latest")
+    triplets = data_adapter.adapt(spans)
+
+    recent_reward = None
+    for triplet in reversed(triplets):
+        if triplet.reward is not None:
+            recent_reward = triplet.reward
+        if recent_reward is None:
+            continue
+
+        input_ids = triplet.prompt["token_ids"] + triplet.response["token_ids"]
+        # We don't train on prompt tokens, so they are masked out by setting to -100.
+        labels = [-100] * len(triplet.prompt["token_ids"]) + triplet.response["token_ids"]
+        # This matches the dataset format required by the Unsloth SFT trainer.
+        all_triplets.append(
+            {
+                "input_ids": input_ids,
+                "attention_mask": [1] * len(input_ids),
+                "labels": labels,
+                "reward": recent_reward,
+            }
+        )
+```
+
+!!! note
+
+    You might notice that the dataset format used here differs from the format described in the **SFT Trainer** documentation. According to the documentation, dataset samples should be provided as plain text strings or message objects.
+
+    As a matter of fact, this example leverages some [undocumented behavior](https://github.com/huggingface/trl/blob/e0eec055b412c48ad754149c475a87a8fca34fb4/trl/trainer/sft_trainer.py#L887) in the SFT Trainer implementation. When the dataset already includes a `"input_ids"` column, the Trainer automatically marks it as `is_processed` and skips the internal tokenization step.
+
+    Since we already have spans with token IDs generated by the [`LLMProxy`][agentlightning.LLMProxy], providing them directly avoids unnecessary [**re-tokenization** and related complications](../deep-dive/serving-llm.md). This approach will both save processing time and increase consistency between training and inference.
+
+After aggregating every rollout we shuffle, sort by reward, and keep the top fraction (e.g., 50%) before shuffling again. The resulting list feeds directly into `datasets.Dataset.from_list`, which is the format Unsloth’s SFT trainer expects.
+
+```python
+from datasets import Dataset as HuggingFaceDataset
+
+random.shuffle(all_triplets)
+all_triplets.sort(key=lambda x: x["reward"], reverse=True)
+sliced_triplets = all_triplets[: max(1, int(len(all_triplets) * triplet_fraction))]
+# Shuffle the sliced triplets again
+random.shuffle(sliced_triplets)
+
+sft_dataset = HuggingFaceDataset.from_list(sliced_triplets)
+```
+
+### Launch Unsloth Training
+
+The heavy lifting happens in [`trl.SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer) (see [unsloth_helper.py]({{ src("examples/unsloth/unsloth_helper.py") }}) on how it's used). We launch it in a fresh process created with `multiprocessing.get_context("spawn")` so CUDA memory is reliably reclaimed when training ends. Launching it in the same process will also work for the first iteration, but we found that the memory won't be freed properly for subsequent vLLM serving.
+
+```python
+context = multiprocessing.get_context("spawn")
+unsloth_process = context.Process(
+    target=unsloth_training,
+    args=(model_path, sft_dataset, next_model_path),
+    daemon=True,
+)
+unsloth_process.start()
+unsloth_process.join(timeout=600.0)
+```
+
+Inside the `unsloth_training` subprocess, Unsloth loads the previous checkpoint in 4-bit, applies LoRA adapters, and forwards the Hugging Face dataset to [`trl.SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer) with the configuration defined in [`SFTConfig`](https://huggingface.co/docs/trl/sft_trainer#trl.SFTConfig) (batch size, accumulation steps, learning rate, etc.). The merged 16-bit weights are saved under `models/version_<iteration + 1>` so the next iteration can immediately serve them with vLLM.
+
+```python
+from unsloth import FastLanguageModel
+# TRL is patched by unsloth.
+from trl import SFTConfig, SFTTrainer
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=model_path,
+    load_in_4bit=True,  # 4 bit quantization to reduce memory
+)
+
+# Config the model to use LoRA
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=32,
+    ...
+)
+
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=sft_dataset,
+    ...
+)
+
+# This is the heaviest step.
+trainer_stats = trainer.train()
+
+# Save in 16-bit for vLLM inference later
+model.save_pretrained_merged(next_model_path, tokenizer, save_method="merged_16bit")
+```
+
+## Math Agent: OpenAI Agents SDK with MCP
+
+We build an agent with the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/) to wire a calculator MCP tool and an OpenAI-compatible chat completion model together. The agent aims to solve a math problem and returns a reward indicating whether the answer is correct or not. The runner injects the `LLM` resource supplied by the algorithm side:
+
+```python
+import os
+from typing import TypedDict
+
+import agentlightning as agl
+from agents import Agent, ModelSettings, OpenAIChatCompletionsModel, Runner as OpenAIRunner
+from agents.mcp import MCPServerStdio
+from openai import AsyncOpenAI
+
+class GsmProblem(TypedDict):
+    input: str
+    target: float
+
+def compute_reward(result: str, target: float) -> float:
+    ...
+
+@agl.rollout
+async def math_agent(task: GsmProblem, llm: agl.LLM) -> float:
+    async with MCPServerStdio(
+        name="Calculator via uvx",
+        params={"command": "uvx", "args": ["mcp-server-calculator"]},
+    ) as server:
+        agent = Agent(
+            name="Assistant",
+            instructions=(
+                "Use the calculator tool for every question. "
+                "Return only the numeric answer wrapped like ### <answer> ###."
+            ),
+            mcp_servers=[server],
+            model=OpenAIChatCompletionsModel(
+                model=llm.model,
+                openai_client=AsyncOpenAI(
+                    base_url=llm.endpoint,
+                    api_key=llm.api_key or "dummy",
+                ),
+            ),
+            model_settings=ModelSettings(
+                temperature=llm.sampling_parameters.get("temperature", 0.0),
+            ),
+        )
+        result = await OpenAIRunner.run(agent, task["input"])
+    return compute_reward(result.final_output, task["target"])
+```
+
+!!! tip
+
+    You can test the agent with a dry run:
+
+    ```python
+    import asyncio
+
+    llm = agl.LLM(
+        endpoint=os.environ["OPENAI_BASE_URL"],
+        api_key=os.environ["OPENAI_API_KEY"],
+        model="gpt-4.1-mini",
+    )
+    asyncio.run(math_agent({"input": "What is 1 + 1?", "target": 2.0}, llm))
+    ```
+
+## Run this Recipe
+
+The full runnable script for this recipe resides in [`examples/unsloth`]({{ src("examples/unsloth") }}) folder.
+
+Before running this example, install `unsloth`, `vllm`, and the other libraries used in the examples (the project uses CUDA tooling, TRL, rich, datasets, etc.). We tested with `unsloth==2025.10.1`. `unsloth==2025.10.2` and `2025.10.3` are not working because of an [issue](https://github.com/unslothai/unsloth/issues/3451) we have been investigating with the unsloth team.
+
+It's recommended to download the base model before running the example, such that the first iteration and subsequent iterations can both load from local checkpoints.
+
+```bash
+hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
+```
+
+The repository already contains `examples/unsloth/data_gsmhard.jsonl` (which is a very small subset of the [GSM-hard math dataset](https://huggingface.co/datasets/reasoning-machines/gsm-hard) for demonstration purposes).
+
+### Run Manually
+
+Similar to the [Write the First Algorithm](./write-first-algorithm.md) recipe, you can open three terminals and start each component in parallel.
+
+```bash
+agl store --port 4747
+python examples/unsloth/sft_rollout_runners.py
+python examples/unsloth/sft_algorithm.py
+```
+
+In this case, [`sft_rollout_runners.py`]({{ src("examples/unsloth/sft_rollout_runners.py") }}) is a simple spawner implemented in Python that spawns 4 runners in parallel. The runners all connect to the same store server executing in another terminal.
+
+```python
+import agentlightning as agl
+
+def run_rollout(store: agl.LightningStore, worker_id: int) -> None:
+    # Since the server side has already used LiteLLM proxy to collect traces,
+    # a simple OtelTracer to collect the rewards is enough.
+    tracer = agl.OtelTracer()
+
+    runner = agl.LitAgentRunner(tracer=tracer)
+
+    with runner.run_context(agent=math_agent, store=store, worker_id=worker_id):
+        asyncio.run(runner.iter())
+
+
+def spawn_runners(store: agl.LightningStore, n_runners: int) -> None:
+    runners = [
+        multiprocessing.Process(target=run_rollout, args=(store, worker_id))
+        for worker_id in range(n_runners)
+    ]
+    for runner in runners:
+        runner.start()
+
+    for runner in runners:
+        runner.join()
+
+
+store = agl.LightningStoreClient("http://localhost:4747")
+spawn_runners(store=store, n_runners=4)
+```
+
+!!! tip
+
+    Try to swap [`OtelTracer`][agentlightning.OtelTracer] in the runners with other tracers like [`AgentOpsTracer`][agentlightning.AgentOpsTracer]. Try to use a different adapter at the algorithm side such as [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] to see what happens.
+
+### Run Everything with Trainer
+
+We also show how to wrap everything into a single script using [`Trainer`][agentlightning.Trainer]. [`sft_allinone.py`]({{ src("examples/unsloth/sft_allinone.py") }}) wires the same components together, replacing the manual management of runners above.
+
+```python
+class UnslothSupervisedFinetuning(agl.Algorithm):
+
+    async def run(
+        self,
+        train_dataset: Optional[Dataset[GsmProblem]] = None,
+        val_dataset: Optional[Dataset[GsmProblem]] = None,
+    ):
+        # Use the store, llm_proxy, and adapter from the trainer
+        store = self.get_store()
+        llm_proxy = self.get_llm_proxy()
+        data_adapter = self.get_adapter()
+
+        for iteration in range(self.max_iterations):
+            ...  # Same logic as sft_algorithm.py
+
+algo = UnslothSupervisedFinetuning(
+    max_iterations=2,
+    vllm_port=12316,
+    train_triplet_fraction=0.5,
+    initial_model_path="models/version_0",
+)
+
+# The LLM proxy can be created before Trainer
+trainer = Trainer(
+    n_runners=4,
+    algorithm=algo,
+    llm_proxy=LLMProxy(port=12358),
+)
+
+trainer.fit(math_agent, load_math_dataset())
+```
+
+You might wonder where the initialization of [`Adapter`][agentlightning.Adapter] happens in this code. It turns out that [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] is the default adapter in [`Trainer`][agentlightning.Trainer], so we don't need to create one manually.
+
+Now you can run the example with:
+
+```bash
+python examples/unsloth/sft_allinone.py
+```
+
+It starts an [`InMemoryLighningStore`][agentlightning.InMemoryLightningStore] for you, launches four worker processes, iterates the SFT loop, and prints the final checkpoint path when done. Adjust `max_iterations`, `train_triplet_fraction`, `n_runners`, or the proxy port to match your hardware or training goals. If you already run an external store or proxy you can also pass those objects into [`Trainer`][agentlightning.Trainer] instead of relying on the [Trainer-managed defaults][debug-with-external-store].
+
+!!! info
+
+    As a future plan, we might graduate this example into a more powerful SFT algorithm bundled into [Algorithm Zoo](../algorithm-zoo/index.md). Currently, this `UnslothSupervisedFinetuning` is still for demo purposes.
@@ -0,0 +1,279 @@
+# Write the First Algorithm with Agent-lightning
+
+In the [first tutorial](./train-first-agent.md), "Train the First Agent," we introduced the [Trainer][agentlightning.Trainer] and showed how to use a pre-built algorithm like **Automatic Prompt Optimization (APO)** to improve an agent's performance. The [Trainer][agentlightning.Trainer] handled all the complex interactions, letting us focus on the agent's logic.
+
+Now, we'll go a step deeper. What if you have a unique training idea that doesn't fit a standard algorithm? This tutorial will show you how to write your own custom algorithm from scratch. We'll build a simple algorithm that systematically tests a list of prompt templates and identifies the one with the highest reward.
+
+By the end, you'll understand the core mechanics of how the **Algorithm**, **Runner**, and a new component, the **Store**, work together to create the powerful training loop at the heart of Agent-lightning.
+
+!!! tip
+    This tutorial helps you build a basic understanding of how to interact with Agent-lightning's core components. It's recommended that all users customizing algorithms should read this tutorial, even for those who are not planning to do prompt optimization.
+
+## Core Concepts for Training
+
+Before diving into the [LightningStore][agentlightning.LightningStore], let's define two key concepts that are central to any training process in Agent-lightning: **Resources** and the **Tracer**.
+
+### Resources: The Tunable Assets
+
+[](){ #introduction-to-resources }
+
+**Resources** are the assets your algorithm is trying to improve. Think of them as the "recipe" an agent uses to perform its task. This recipe can be:
+
+  * A **prompt template** that guides an LLM.
+  * The **weights** of a machine learning model.
+  * Any other configuration or data your agent needs.
+
+The algorithm's job is to run experiments and iteratively update these resources to find the best-performing version.
+
+### Tracer: The Data Collector
+
+How does the algorithm know if a change was an improvement? It needs data. This is where the **Tracer** comes in.
+
+The Tracer automatically **instruments** (aka modifies / patches) the agent's code. This means it watches for important events, like an LLM call, a tool being used, or **reward signals**, and records a detailed log of what happened. Each of these logs is called a **Span** (which has already been introduced in the [last tutorial](./train-first-agent.md)).
+
+A collection of spans from a single task execution gives the algorithm a complete, step-by-step trace of the agent's behavior, which is essential for learning and making improvements. Our default tracer is built on the AgentOps SDK to support instrumenting code written in various Agent/non-agent frameworks.
+
+## The Central Hub: The LightningStore
+
+Now, where do all these resources, tasks, and spans live? They are all managed by the **LightningStore**.
+
+The LightningStore acts as the central database and message queue for the entire system. It's the single source of truth that decouples the Algorithm from the Runners.
+
+!!! note
+
+    In the [last tutorial](./train-first-agent.md) we simplified the training loop, saying the Algorithm and Agent communicate "via the Trainer." That's true at a high level, but the component that makes it all possible is actually the **LightningStore**.
+
+* The **Algorithm** connects to the Store to `enqueue_rollout` (tasks) and `update_resources` (like new prompt templates). It also queries the Store to retrieve the resulting spans and rewards from completed rollouts.
+* The **Runners** connect to the Store to `dequeue_rollout` (polling for available tasks). After executing a task, they use the `Tracer` to write the resulting spans and status updates back to the Store.
+
+This architecture is key to Agent-lightning's scalability. Since the Algorithm and Runners only talk to the Store, they can run in different processes or even on different machines.
+
+![Store Architecture](../assets/store-api-visualized.svg){ .center }
+
+!!! tip "A Mental Model of What the Store Contains"
+
+    The [LightningStore][agentlightning.LightningStore] isn't just a simple database; it's an organized system for managing the entire training lifecycle. Here's what it keeps track of:
+
+    * **Task Queue**: A queue of pending **Rollouts** waiting for a Runner to pick them up, interactable via `enqueue_rollout` and `dequeue_rollout`.
+    * **Rollouts**: The record of a single task. A rollout contains metadata about the task and tracks all **Attempts** to complete it, interactable via `query_rollouts` and `wait_for_rollouts`.
+    * **Attempts**: A single execution of a rollout. If an attempt fails (e.g., due to a network error), the Store can automatically schedules a retry if it's configured. Each attempt is linked to its parent rollout and contains the status and timing information. The rollout status is [synced](../deep-dive/store.md) with its children's status. **For beginners, you can assume each rollout has only one attempt unless you have explicitly configure the retry.**
+    * **Spans**: The detailed, structured logs generated by the `Tracer` during an attempt. Each span is linked to its parent attempt and rollout.
+    * **Resources**: A versioned collection of the assets (like prompt templates) that the algorithm creates. Each rollout is linked to the specific version of the resources it should use.
+
+## Building a Custom Algorithm
+
+Let's build an algorithm that finds the best system prompt from a predefined list. The logic is straightforward:
+
+1.  Start with a list of candidate prompt templates.
+2.  For each template, create a "resource" bundle in the Store.
+3.  Enqueue a rollout (a task), telling the Runner to use this specific resource.
+4.  Wait for a Runner to pick up the task and complete it.
+5.  Query the Store to get the final reward from the rollout's spans.
+6.  After testing all templates, compare the rewards and declare the best one.
+
+We can implement this as a simple Python function that interacts directly with the [LightningStore][agentlightning.LightningStore].
+
+```python
+async def find_best_prompt(store, prompts_to_test, task_input):
+    """A simple algorithm to find the best prompt from a list."""
+    results = []
+
+    # Iterate through each prompt to test it
+    for prompt in prompts_to_test:
+        print(f"[Algo] Updating prompt template to: '{prompt}'")
+
+        # 1. Update the resources in the store with the new prompt
+        resources_update = await store.add_resources(
+            resources={"prompt_template": prompt}
+        )
+
+        # 2. Enqueue a rollout task for a runner to execute
+        print("[Algo] Queuing task for clients...")
+        rollout = await store.enqueue_rollout(
+            input=task_input,
+            resources_id=resources_update.resources_id,
+        )
+        print(f"[Algo] Task '{rollout.rollout_id}' is now available for clients.")
+
+        # 3. Wait for the rollout to be completed by a runner
+        await store.wait_for_rollouts([rollout.rollout_id])
+
+        # 4. Query the completed rollout and its spans
+        completed_rollout = (await store.query_rollouts([rollout.rollout_id]))[0]
+        print(f"[Algo] Received Result: {completed_rollout.model_dump_json(indent=None)}")
+
+        spans = await store.query_spans(rollout.rollout_id)
+        # We expect at least two spans: one for the LLM call and one for the final reward
+        print(f"[Algo] Queried Spans:\n  - " + "\n  - ".join(str(span) for span in spans))
+        # find_final_reward is a helper function to extract the reward span
+        final_reward = find_final_reward(spans)
+        print(f"[Algo] Final reward: {final_reward}\n")
+
+        results.append((prompt, final_reward))
+
+    # 5. Find and print the best prompt based on the collected rewards
+    print(f"[Algo] All prompts and their rewards: {results}")
+    best_prompt, best_reward = max(results, key=lambda item: item[1])
+    print(f"[Algo] Best prompt found: '{best_prompt}' with reward {best_reward}")
+```
+
+!!! note "Asynchronous Operations"
+
+    You'll notice the `async` and `await` keywords. Agent-lightning is built on asyncio to handle concurrent operations efficiently. All interactions with the store are asynchronous network calls, so they must be awaited.
+
+## The Agent and Runner
+
+Our algorithm needs an **agent** to execute the tasks and a **runner** to manage the process.
+
+The runner is a long-lived worker process. Its job is simple:
+
+1.  Connect to the [LightningStore][agentlightning.LightningStore] via a [LightningStoreClient][agentlightning.LightningStoreClient].
+2.  Enter a loop, constantly asking the [LightningStore][agentlightning.LightningStore] for new tasks (`dequeue_rollout`).
+3.  When it gets a task, it runs the `simple_agent` function.
+4.  Crucially, the runner wraps the agent execution with a **Tracer**. The tracer automatically captures all the important events (like the LLM call and the final reward) as spans and sends them back to the [LightningStore][agentlightning.LightningStore].
+
+```python
+# Connecting to Store
+store = agl.LightningStoreClient("http://localhost:4747")  # or some other address
+runner = LitAgentRunner[str](tracer=AgentOpsTracer())
+with runner.run_context(agent=simple_agent, store=store):  # <-- where the wrapping and instrumentation happens
+    await runner.iter()  # polling for new tasks forever
+```
+
+For this example, the agent's job is to take the prompt from the resources, use it to ask an LLM a question, and return a score.
+
+```python
+def simple_agent(task: str, prompt_template: PromptTemplate) -> float:
+    """An agent that answers a question and gets judged by an LLM."""
+    client = OpenAI()
+
+    # Generate a response using the provided prompt template
+    prompt = prompt_template.format(any_question=task)
+    response = client.chat.completions.create(
+        model="gpt-4.1-nano", messages=[{"role": "user", "content": prompt}]
+    )
+    llm_output = response.choices[0].message.content
+    print(f"[Rollout] LLM returned: {llm_output}")
+
+    # This llm_output and the final score are automatically logged as spans by the Tracer
+    score = random.uniform(0, 1)  # Replace with actual scoring logic if needed
+    return score
+```
+
+## Running the Example
+
+To see everything in action, you'll need three separate terminal windows.
+
+!!! tip
+
+    If you want to follow along, you can find the complete code for this example in the [apo_custom_algorithm.py]({{ src("examples/apo/apo_custom_algorithm.py") }}) file.
+
+**1. Start the Store:**
+In the first terminal, start the LightningStore server. This component will wait for connections from the algorithm and the runner. The store will be listening on port `4747` ⚡ by default.
+
+```bash
+agl store
+```
+
+**2. Start the Runner:**
+In the second terminal, start the runner process. It will connect to the store and wait for tasks.
+
+The code to start the runner looks like the following:
+
+```bash
+export OPENAI_API_KEY=sk-... # Your OpenAI API key
+python apo_custom_algorithm.py runner
+```
+
+You will see output indicating the runner has started and is waiting for rollouts.
+
+```text
+2025-10-14 22:23:41,339 [INFO] ... [Worker 0] Setting up tracer...
+2025-10-14 22:23:41,343 [INFO] ... [Worker 0] Instrumentation applied.
+2025-10-14 22:23:41,494 [INFO] ... [Worker 0] AgentOps client initialized.
+2025-10-14 22:23:41,494 [INFO] ... [Worker 0] Started async rollouts (max: unlimited).
+```
+
+**3. Start the Algorithm:**
+In the third terminal, run the algorithm. This will kick off the entire process.
+
+For example, we run the algorithm code shown above with the following parameters:
+
+```python
+prompts_to_test = [
+    "You are a helpful assistant. {any_question}",
+    "You are a knowledgeable AI. {any_question}",
+    "You are a friendly chatbot. {any_question}",
+]
+task_input = "Why is the sky blue?"
+store = agl.LightningStoreClient("http://localhost:4747")
+find_best_prompt(store, prompts_to_test, task_input)
+```
+
+Or you can simply use our pre-written script to try out:
+
+```bash
+python apo_custom_algorithm.py algo
+```
+
+### Understanding the Output
+
+As the algorithm runs, you'll see logs appear across all three terminals, showing the components interacting in real-time.
+
+**Algorithm Output:**
+The algorithm terminal shows the main control flow: updating prompts, queuing tasks, and receiving the final results. You can also see the raw span data it retrieves from the store.
+
+```text
+[Algo] Updating prompt template to: 'You are a helpful assistant. {any_question}'
+[Algo] Queuing task for clients...
+[Algo] Task 'ro-1d18988581cd' is now available for clients.
+[Algo] Received Result: rollout_id='ro-1d18988581cd' ... status='succeeded' ...
+[Algo] Queried Spans:
+  - Span(name='openai.chat.completion', attributes={'gen_ai.prompt.0.content': 'You are a helpful assistant...', 'gen_ai.completion.0.content': 'The sky appears blue...'})
+  - Span(name='reward', attributes={'value': 0.95})
+[Algo] Final reward: 0.95
+
+[Algo] Updating prompt template to: 'You are a knowledgeable AI. {any_question}'
+...
+[Algo] Final reward: 0.95
+
+[Algo] Updating prompt template to: 'You are a friendly chatbot. {any_question}'
+...
+[Algo] Final reward: 1.0
+
+[Algo] All prompts and their rewards: [('You are a helpful assistant. {any_question}', 0.95), ('You are a knowledgeable AI. {any_question}', 0.95), ('You are a friendly chatbot. {any_question}', 1.0)]
+[Algo] Best prompt found: 'You are a friendly chatbot. {any_question}' with reward 1.0
+```
+
+**Runner Output:**
+The runner terminal shows it picking up each task, executing the agent logic, and reporting the completion.
+
+```text
+[Rollout] LLM returned: The sky appears blue due to Rayleigh scattering...
+2025-10-14 22:25:50,803 [INFO] ... [Worker 0 | Rollout ro-a9f54ac19af5] Completed in 4.24s. ...
+
+[Rollout] LLM returned: The sky looks blue because of a process called Rayleigh scattering...
+2025-10-14 22:25:59,863 [INFO] ... [Worker 0 | Rollout ro-c67eaa9016b6] Completed in 4.06s. ...
+```
+
+**Store Server Output:**
+The store terminal shows a detailed log of every interaction, confirming its role as the central hub. You can see requests to enqueue and dequeue rollouts, add spans, and update statuses.
+
+```text
+... "POST /enqueue_rollout HTTP/1.1" 200 ...
+... "GET /dequeue_rollout HTTP/1.1" 200 ...
+... "POST /add_span HTTP/1.1" 200 ...
+... "POST /update_attempt HTTP/1.1" 200 ...
+... "POST /wait_for_rollouts HTTP/1.1" 200 ...
+... "GET /query_spans/ro-c67eaa9016b6 HTTP/1.1" 200 ...
+```
+
+!!! info "So Where is Trainer?"
+
+    You might be wondering why the [last tutorial](./train-first-agent.md) focused on the [Trainer][agentlightning.Trainer] class, but we haven't used it here.
+
+    Think of the [Trainer][agentlightning.Trainer] as a convenient wrapper that manages the entire training process for you. It's perfect when you want to apply a pre-built algorithm to your agent without worrying about the underlying mechanics. The [Trainer][agentlightning.Trainer] handles starting the [LightningStore][agentlightning.LightningStore], coordinating the [Runners][agentlightning.Runner], managing their lifecycles, and handling errors.
+
+    In this tutorial, however, our goal is to *build a new algorithm*. To do that, we need to interact directly with the core components: the [Store][agentlightning.LightningStore], the [Runner][agentlightning.Runner], and the algorithm logic itself. Running them separately gives you more control and clearer, isolated logs, which is ideal for development and debugging.
+
+    Once your custom algorithm is mature, you can package it to comply with our standard interface ([@algo][agentlightning.algo] or [Algorithm][agentlightning.Algorithm]). This allows you to use it with the [Trainer][agentlightning.Trainer] again, getting all the benefits of automated lifecycle management while using your own custom logic. A sample code doing this is available in [apo_custom_algorithm_trainer.py]({{ src("examples/apo/apo_custom_algorithm_trainer.py") }}).
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Yuge Zhang	49ae371714	.	2025-11-05 09:27:59 +08:00
Yuge Zhang	af26204934	.	2025-11-04 10:45:14 +08:00
Yuge Zhang	91f9e1be3e	.	2025-11-04 01:27:48 +08:00
Yuge Zhang	8c8d474b65	.	2025-11-04 01:23:21 +08:00
Yuge Zhang	d7749a0ad3	.	2025-11-04 00:47:24 +08:00
Yuge Zhang	b26e646d86	.	2025-11-03 16:42:51 +00:00
Yuge Zhang	d245218651	.	2025-11-04 00:30:39 +08:00
Yuge Zhang	261529cbe4	..	2025-11-03 23:33:53 +08:00
Yuge Zhang	fdc00cbce3	.	2025-11-03 21:00:47 +08:00
Yuge Zhang	9cb5767c46	add new metrics	2025-11-03 12:28:29 +00:00
Yuge Zhang	c7ff2b2a9e	.	2025-11-01 09:54:03 +08:00
Yuge Zhang	6177a4a97d	update q20 train verl	2025-11-01 09:45:15 +08:00
Yuge Zhang	4e8eb53deb	update hyperparameters	2025-11-01 01:12:28 +00:00
Yuge Zhang	772e3a7ebf	.	2025-11-01 01:06:32 +00:00
Yuge Zhang	728e3fefc8	.	2025-10-31 09:36:12 +00:00
Yuge Zhang	a0565e126f	update 5mini script	2025-10-30 23:48:27 +08:00
Yuge Zhang	ea410e34bf	.	2025-10-30 09:18:46 +00:00
Yuge Zhang	131c7034ef	.	2025-10-30 09:02:53 +00:00
Yuge Zhang	03b42b4d12	update grid search script	2025-10-30 09:01:45 +00:00
Yuge Zhang	6f7801a38d	update grid search	2025-10-30 15:58:03 +08:00
Yuge Zhang	8e7dda5416	.	2025-10-30 13:06:55 +08:00
Yuge Zhang	ddd030c8f8	revert test results	2025-10-30 13:05:52 +08:00
Yuge Zhang	ceb4c547ce	Merge branch 'main' of github.com:microsoft/agent-lightning into tinker	2025-10-30 13:01:22 +08:00
Yuge Zhang	496e793f0b	Tinker Integration (#245 )	2025-10-30 12:44:17 +08:00
Yuge Zhang	0d1cd44c98	resolve comments	2025-10-30 12:39:33 +08:00
Yuge Zhang	917f37d356	resolve comments	2025-10-30 12:33:22 +08:00
Yuge Zhang	4776134399	update Readme	2025-10-30 12:21:41 +08:00
Yuge Zhang	94e5d5d869	update readme	2025-10-30 12:16:26 +08:00
Yuge Zhang	14ecf0d9a9	.	2025-10-30 11:55:11 +08:00
Yuge Zhang	d1ed396f6f	minor fix	2025-10-30 11:32:37 +08:00
Yuge Zhang	2440ef0f3f	update docstrings	2025-10-30 11:12:22 +08:00
Yuge Zhang	5f01765001	checkpoint	2025-10-30 10:35:27 +08:00
Yuge Zhang	5241410c37	Preparing code for merge	2025-10-30 09:33:51 +08:00
Yuge Zhang	80d306ff54	[Release] v0.2.1 (#243 )	2025-10-30 08:31:45 +08:00
Yuge Zhang	003b8f7dd2	update q20 train	2025-10-30 07:36:07 +08:00
Yuge Zhang	af4c8d0f76	.	2025-10-30 00:49:45 +08:00
Yuge Zhang	efa2af1033	update search	2025-10-30 00:44:02 +08:00
Yuge Zhang	5f67bfe137	Normalize Store FastAPI (#241 )	2025-10-29 21:43:37 +08:00
Yuge Zhang	9ae7ae42df	.	2025-10-29 20:25:35 +08:00
Yuge Zhang	3633849d5d	.	2025-10-29 18:19:17 +08:00
Yuge Zhang	a0e59871c0	.	2025-10-29 16:25:07 +08:00
Yuge Zhang	2f29d04f8e	fix rollout status	2025-10-29 12:02:18 +08:00
Yuge Zhang	92a5a3f85f	.	2025-10-29 10:52:01 +08:00
Yuge Zhang	c6336343d1	fix error handling	2025-10-29 10:33:42 +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
Yuge Zhang	6b86d65383	.	2025-10-29 01:13:56 +08:00
Yuge Zhang	064e535664	update 32x4	2025-10-29 01:01:58 +08:00
Yuge Zhang	635788d0d6	.	2025-10-29 00:09:45 +08:00
Yuge Zhang	3d94def44a	.	2025-10-28 22:03:51 +08:00
Yuge Zhang	5bd9b61d85	update lr5e-4 16x8	2025-10-28 22:01:42 +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	2d7e6656c8	.	2025-10-28 14:17:26 +08:00
Yuge Zhang	91798bdeb0	.	2025-10-28 14:16:03 +08:00
Yuge Zhang	73e0140dae	.	2025-10-28 14:14:01 +08:00
Yuge Zhang	36e52c80c4	minor fix	2025-10-28 14:10:14 +08:00
Yuge Zhang	7a606f9726	fix q20 train	2025-10-28 14:02:53 +08:00
Yuge Zhang	f4cdc3dcbd	reorganize folder	2025-10-28 13:42:04 +08:00
Yuge Zhang	ebd9ce6664	Merge from main	2025-10-28 12:53:54 +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	860d6662c6	add tool results	2025-10-28 03:55:45 +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	12aadb0036	update q20 config	2025-10-27 21:16:22 +08:00
Yuge Zhang	38ea3857a6	add more results	2025-10-27 13:11:29 +00:00
Yuge Zhang	4f152856f3	update rollout logic	2025-10-27 21:10:32 +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	bf72877eac	Merge branch 'main' of github.com:microsoft/agent-lightning into tinker	2025-10-27 19:40:52 +08:00
Yuge Zhang	9a81fa6cc5	freeze split by category	2025-10-27 19:34:37 +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	12d6fcb530	update port usage in hello	2025-10-27 16:43:27 +08:00
Yuge Zhang	f6737b52ba	Merge branch 'main' of github.com:microsoft/agent-lightning into tinker	2025-10-27 16:36:41 +08:00
Yuge Zhang	30880cd45a	minor fix	2025-10-27 16:24:56 +08:00
Yuge Zhang	7c47e957db	minor fix	2025-10-27 14:53:23 +08:00
Yuge Zhang	106d6542fd	add qwen more results	2025-10-27 06:51:37 +00:00
Yuge Zhang	dd0c9488fc	minor fix	2025-10-27 14:40:53 +08:00
Yuge Zhang	1f35993626	fix argparse	2025-10-27 14:09:43 +08:00
Yuge Zhang	2b6987cc99	minor fix	2025-10-27 14:07:30 +08:00
Yuge Zhang	2a30f47df9	.	2025-10-27 13:56:33 +08:00
Yuge Zhang	852e4f80dd	.	2025-10-27 13:07:48 +08:00
Yuge Zhang	46e81dda2e	update	2025-10-27 11:48:28 +08:00
Yuge Zhang	c83a6977d4	update q20 training dry run	2025-10-27 10:45:57 +08:00
Yuge Zhang	306f591d3a	add partial qwen results	2025-10-27 01:05:05 +00:00
Yuge Zhang	4a617a480d	use 4.1 instead	2025-10-27 00:38:15 +08:00
Yuge Zhang	9a9b3bb74e	with search ver	2025-10-26 23:42:51 +08:00
Yuge Zhang	9e60077902	add partial results	2025-10-26 14:48:49 +00:00
Yuge Zhang	c3d430a9c2	.	2025-10-26 18:19:49 +08:00
Yuge Zhang	d6cee16f3d	split category	2025-10-26 17:54:09 +08:00
Yuge Zhang	3cfba40e1e	.	2025-10-26 17:02:31 +08:00
Yuge Zhang	a736fd92ff	update 20q	2025-10-26 17:00:36 +08:00
Yuge Zhang	5679cd21d2	update dataset	2025-10-26 16:31:54 +08:00
Yuge Zhang	c58be56ad4	continue to add qwen results	2025-10-26 08:12:26 +00:00
Yuge Zhang	00407764f9	add partial qwen result	2025-10-26 03:55:24 +00:00
Yuge Zhang	eb9a4d2932	.	2025-10-26 10:59:26 +08:00
Yuge Zhang	bc6d72985e	.	2025-10-26 10:37:08 +08:00
Yuge Zhang	891b92bfc5	.	2025-10-26 10:35:00 +08:00
Yuge Zhang	f0051ddb3a	try out qwen3 235b	2025-10-26 10:01:32 +08:00
Yuge Zhang	4d5e133a06	Add vLLM blog link to resources (#215 )	2025-10-26 09:57:07 +08:00
Yuge Zhang	373259a949	add first test result	2025-10-26 01:17:34 +00:00
Yuge Zhang	ce20340ec6	.	2025-10-26 09:13:32 +08:00
Yuge Zhang	b8197cca8d	update agent impl	2025-10-26 01:11:57 +08:00
Yuge Zhang	2459163c85	.	2025-10-26 00:15:26 +08:00
Yuge Zhang	fa51592e31	.	2025-10-25 21:49:15 +08:00
Yuge Zhang	7effbefd10	.	2025-10-25 18:54:28 +08:00
Yuge Zhang	9690c9a9fc	update impl	2025-10-25 18:41:58 +08:00
Yuge Zhang	c71c466c4c	.	2025-10-25 18:27:06 +08:00
Yuge Zhang	0a215353ad	.	2025-10-25 16:26:26 +08:00
Yuge Zhang	df2a159b00	Add tutorial for launching workers on separate machines (#213 )	2025-10-25 14:16:27 +08:00
Yuge Zhang	53b2d701c2	CrewAI baseline	2025-10-25 14:15:29 +08:00
Yuge Zhang	675fc86727	Issue comment responder (#214 )	2025-10-25 13:00:02 +08:00
Yuge Zhang	8ef2c82ea8	Merge branch 'main' of github.com:microsoft/agent-lightning into tinker	2025-10-25 12:36:55 +08:00
Yuge Zhang	991aa2be10	update uv lock from main	2025-10-25 12:36:40 +08:00
Yuge Zhang	c16b3a21b6	Add dependency groups from Tinker and CrewAI (#212 )	2025-10-25 12:27:37 +08:00
Yuge Zhang	f00b1f3523	update uv lock	2025-10-25 11:24:48 +08:00
Yuge Zhang	8c75644ec4	update twenty question nouns dataset	2025-10-25 10:43:49 +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	0e3244b782	Merge branch 'main' of github.com:microsoft/agent-lightning into tinker	2025-10-24 01:13:37 +08:00
Yuge Zhang	e28fb8cb6b	Fix LiteLLM logging worker reset on proxy restart (#174 )	2025-10-24 00:56:16 +08:00
Yuge Zhang	0a5f8e0470	hello example running	2025-10-24 00:00:30 +08:00
Yuge Zhang	fae0fba3d7	Fix trigger on label (#204 )	2025-10-23 23:05:43 +08:00
Yuge Zhang	c88e0cfc3d	.	2025-10-23 23:04:20 +08:00
Yuge Zhang	a9271bf8e4	.	2025-10-23 22:52:40 +08:00
Yuge Zhang	6d912b8a60	.	2025-10-23 16:52:06 +08:00
Yuge Zhang	99f1d7ab42	.	2025-10-23 16:33:25 +08:00
Yuge Zhang	46bdd2bef5	update test llm	2025-10-23 16:14:55 +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	faa46dbd24	.	2025-10-23 15:00:50 +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	5d8082570c	update test llm	2025-10-23 12:05:44 +08:00
Yuge Zhang	9b6ca72e35	move files	2025-10-23 11:19:15 +08:00
Yuge Zhang	580bc13229	Merge branch 'main' of github.com:microsoft/agent-lightning into tinker	2025-10-23 10:09:19 +08:00
Yuge Zhang	6bf00d80c6	test with hello 1234567	2025-10-23 00:21:34 +08:00
Yuge Zhang	a2c9b0a6ae	.	2025-10-22 21:51:31 +08:00
Yuge Zhang	8abe85ad91	Using Group Subscription for CI (#195 )	2025-10-22 21:39:58 +08:00
Yuge Zhang	2bdf7ff275	.	2025-10-22 21:28:56 +08:00
Yuge Zhang	38aa5fd6e8	.	2025-10-22 20:56:56 +08:00
Yuge Zhang	8a002f83e0	.	2025-10-22 20:55:48 +08:00
Yuge Zhang	b878de116a	minimize	2025-10-22 14:49:22 +08:00
Yuge Zhang	e7ef3264c9	copy rl_train	2025-10-22 14:12:24 +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	9007fb28db	.	2025-10-22 13:47:32 +08:00
Yuge Zhang	483c518d74	Adjust CI status placement and example catalog details (#193 )	2025-10-22 13:11:40 +08:00
Yuge Zhang	978d9f48b3	fix litellm proxy	2025-10-22 12:48:44 +08:00
Yuge Zhang	3ac0e1f239	.	2025-10-22 12:36:41 +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	05748c5311	.	2025-10-20 23:50:12 +08: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	20b13f18e9	.	2025-10-20 17:42:14 +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	87798cc12f	add tinker dependency	2025-10-20 14:43:12 +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