.

2025-11-24 00:10:28 +08:00 · 2025-11-24 00:09:05 +08:00 · 2025-11-24 00:06:45 +08:00 · 2025-11-23 23:52:39 +08:00 · 2025-11-23 23:28:53 +08:00 · 2025-11-23 11:30:23 +08:00
66 changed files with 4440 additions and 993 deletions
@@ -0,0 +1,14 @@
+.venv
+**/.venv
+__pycache__
+.git
+.gitignore
+**/node_modules
+dist
+build
+.env
+docker
+.pytest_cache
+.vscode
+**/*.log
+examples/**/data
@@ -0,0 +1,29 @@
+name: Badge - Azure
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Azure
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-azure.yml', label: 'azure', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -7,6 +7,8 @@ on:
      - Examples - Spider
      - Examples - APO
      - Examples - Unsloth
+      - Examples - Tinker
+      - Examples - Azure
    types: [completed]

  workflow_dispatch:
@@ -31,5 +33,7 @@ jobs:
              { workflow: 'examples-spider.yml', label: 'examples-spider.stable', variants: ['stable'] },
              { workflow: 'examples-apo.yml', label: 'examples-apo.stable', variants: ['stable'] },
              { workflow: 'examples-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
+              { workflow: 'examples-tinker.yml', label: 'examples-tinker.stable', variants: ['stable'] },
+              { workflow: 'examples-azure.yml', label: 'examples-azure.stable', variants: ['stable'] },
            ];
            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,29 @@
+name: Badge - Tinker
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Tinker
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-tinker.yml', label: 'tinker', variants: ['stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });
@@ -0,0 +1,19 @@
+# This workflow is used to benchmark the performance of the project.
+# It's kept as a placeholder for now.
+
+name: Benchmark
+permissions:
+  contents: read
+on:
+  workflow_dispatch:
+
+jobs:
+  benchmark:
+    name: Benchmark
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
@@ -0,0 +1,98 @@
+name: Examples - Azure
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: '0 20 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-azure, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Azure - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Azure - {0}', github.event_name) }}
+
+jobs:
+  azure:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-azure' ||
+      github.event.action == 'ci-all'
+    name: Azure (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
+    timeout-minutes: 400
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.12'
+          setup-script: 'stable'
+      fail-fast: false
+    steps:
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group core-stable
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-azure-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Azure Login
+        run: |
+          az login --identity
+        shell: bash
+
+      - name: Azure OpenAI Sanity Check
+        run: |
+          source .venv/bin/activate
+          cd examples/azure
+          python capital_agent.py
+        shell: bash
+        env:
+          AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+        id: azure_openai_sanity_check
+
+      - name: Azure OpenAI Supervised Fine-tuning
+        run: |
+          source .venv/bin/activate
+          cd examples/azure
+          python train_capital_agent.py --n-iterations 2 --cleanup
+        shell: bash
+        env:
+          AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
+          AZURE_OPENAI_API_VERSION: 2025-04-01-preview
+          AZURE_RESOURCE_GROUP: ${{ secrets.AZURE_RESOURCE_GROUP }}
+          AZURE_RESOURCE_NAME: ${{ secrets.AZURE_RESOURCE_NAME }}
+        id: azure_openai_finetune
@@ -0,0 +1,160 @@
+name: Examples - Tinker
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 3 AM UTC+8
+    - cron: '0 19 * * *'
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-tinker, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Tinker - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Tinker - {0}', github.event_name) }}
+
+jobs:
+  tinker:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-tinker' ||
+      github.event.action == 'ci-all'
+    name: Tinker (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
+    timeout-minutes: 150
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group torch-cpu --group core-stable --group tinker
+      - name: Freeze dependencies
+        run: |
+          set -euo pipefail
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-tinker-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Tinker LLM sanity check
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          # TODO: Currently only test the client tracer implementation.
+          python -m tests.test_tinker_llm
+        shell: bash
+        env:
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Hello
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          python hello.py oneclick --ci
+        shell: bash
+        env:
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Evaluate (GPT-4.1)
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          mkdir -p logs
+          python q20_evaluate.py --ci --model gpt-4.1 --output-file logs/q20_evaluate_gpt-4.1.jsonl
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Evaluate (Qwen3-30B-A3B-Instruct-2507)
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          python q20_evaluate.py --ci --model Qwen/Qwen3-30B-A3B-Instruct-2507 --output-file logs/q20_evaluate_qwen3-30b-a3b.jsonl
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Training Dry Run
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          python q20_train.py dryrun --model qwen4b
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
+
+      - name: Tinker Q20 Training
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          cd examples/tinker
+          agl store --port 4747 &
+          sleep 5
+          python q20_train.py runner --n-runners 4 &
+          sleep 5
+          python q20_train.py algo --model qwen4b --ci
+          sleep 5
+
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+          pkill -f q20_train.py && echo "SIGTERM sent to q20_train.py" || echo "No q20_train.py process found"
+          while pgrep -f q20_train.py; do
+            echo "Waiting for q20_train.py to finish..."
+            sleep 5
+          done
+          echo "q20_train.py has finished."
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+          CREWAI_DISABLE_TELEMETRY: true
+          TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
@@ -47,6 +47,7 @@ jobs:
      - 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
@@ -55,10 +56,10 @@ jobs:
        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
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-stable
        if: matrix.setup-script == 'latest'
      - name: Sync dependencies (stable & legacy)
-        run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
+        run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
        if: matrix.setup-script != 'latest'
      - name: Freeze dependencies
        run: |
@@ -81,6 +82,52 @@ jobs:
      - name: Build dashboard
        run: cd dashboard && npm run build

+      - name: Setup Docker environments
+        run: |
+          set -euo pipefail
+
+          cd docker
+
+          # Setup data directories
+          ./setup.sh
+
+          # Start Dockers
+          docker compose -f compose.mongo.yml up -d
+
+          SERVICE_NAME=mongo
+          TIMEOUT=60   # seconds
+          SLEEP=2
+
+          cid="$(docker compose -f compose.mongo.yml ps -q "$SERVICE_NAME")"
+          if [ -z "$cid" ]; then
+            echo "Service $SERVICE_NAME is not running"
+            exit 1
+          fi
+
+          echo "Waiting for $SERVICE_NAME to become healthy..."
+          end=$((SECONDS + TIMEOUT))
+
+          while [ "$SECONDS" -lt "$end" ]; do
+            status="$(docker inspect -f '{{.State.Health.Status}}' "$cid")"
+            echo "Current status: $status"
+
+            if [ "$status" = "healthy" ]; then
+              echo "$SERVICE_NAME is healthy ✅"
+              exit 0
+            elif [ "$status" = "unhealthy" ]; then
+              echo "$SERVICE_NAME is unhealthy ❌"
+              docker logs "$cid" || true
+              exit 1
+            fi
+
+            sleep "$SLEEP"
+          done
+
+          echo "Timed out waiting for $SERVICE_NAME to become healthy after ${TIMEOUT}s"
+          docker logs "$cid" || true
+          exit 1
+        shell: bash
+
      - name: Launch LiteLLM Proxy
        run: |
          ./scripts/litellm_run.sh
@@ -95,6 +142,8 @@ jobs:
          PYTEST_ADDOPTS: "--color=yes"
          OPENAI_BASE_URL: http://localhost:12306/
          OPENAI_API_KEY: dummy
+          AGL_TEST_MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
+

  minimal-examples:
    if: >
@@ -160,6 +209,7 @@ jobs:
          source .venv/bin/activate
          cd examples/minimal
          python write_traces.py otel
+          sleep 5

      - name: Write Traces via AgentOps Tracer
        env:
@@ -170,6 +220,7 @@ jobs:
          source .venv/bin/activate
          cd examples/minimal
          python write_traces.py agentops
+          sleep 5

      - name: Write Traces via Otel Tracer with Client
        run: |
@@ -37,6 +37,7 @@ jobs:
          uv sync --frozen \
            --extra apo \
            --extra verl \
+            --extra mongo \
            --group dev \
            --group torch-cpu \
            --group torch-stable \
@@ -166,7 +167,7 @@ jobs:

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

@@ -213,3 +213,6 @@ agentlightning/dashboard/**/*.css
 agentlightning/dashboard/**/*.js
 agentlightning/dashboard/**/*.html
 agentlightning/dashboard/**/*.svg
+
+# Docker data
+docker/data/
@@ -99,7 +99,7 @@ If you find Agent Lightning useful in your research or projects, please cite our

 ## ⚡ Contributing

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

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

@@ -18,6 +18,7 @@ logger = logging.getLogger(__name__)

 def main(argv: Iterable[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run a LightningStore server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the server to")
    parser.add_argument("--port", type=int, default=4747, help="Port to run the server on")
    parser.add_argument(
        "--cors-origin",
@@ -31,17 +32,60 @@ def main(argv: Iterable[str] | None = None) -> int:
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Configure the logging level for the store.",
    )
+    parser.add_argument(
+        "--prometheus",
+        action="store_true",
+        help="Enable Prometheus metrics.",
+    )
+    parser.add_argument(
+        "--n-workers",
+        default=1,
+        type=int,
+        help=(
+            "Number of workers to run in the server. When it's greater than 1, the server will be run using `mp` launch mode. "
+            "Only applicable for zero-copy stores such as MongoDB backend."
+        ),
+    )
+
+    parser.add_argument(
+        "--backend",
+        choices=["memory", "mongo"],
+        default="memory",
+        help="Backend to use for the store.",
+    )
+    parser.add_argument(
+        "--mongo-uri",
+        default="mongodb://localhost:27017/?replicaSet=rs0",
+        help="MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.",
+    )
+
    args = parser.parse_args(list(argv) if argv is not None else None)

    setup_logging(args.log_level)

-    store = InMemoryLightningStore()
+    if args.backend == "memory":
+        store = InMemoryLightningStore()
+    elif args.backend == "mongo":
+        from agentlightning.store.mongo import MongoLightningStore
+
+        store = MongoLightningStore(client=args.mongo_uri)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    if args.n_workers > 1:
+        logger.info(f"Running the server using `mp` launch mode with {args.n_workers} workers.")
+        launch_mode = "mp"
+    else:
+        logger.info("Running the server using `asyncio` launch mode.")
+        launch_mode = "asyncio"
    server = LightningStoreServer(
        store,
-        host="0.0.0.0",
+        host=args.host,
        port=args.port,
        cors_allow_origins=args.cors_origins,
-        launch_mode="asyncio",
+        launch_mode=launch_mode,
+        prometheus=args.prometheus,
+        n_workers=args.n_workers,
    )
    try:
        asyncio.run(server.run_forever())
@@ -853,7 +853,6 @@ class StreamConversionMiddleware(BaseHTTPMiddleware):
        )  # e.g., "stop", "length", "tool_calls", "content_filter"

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

        # 1) initial chunk with the role
@@ -138,7 +138,7 @@ class LitAgentRunner(Runner[T_task]):
        self._store = store
        self.worker_id = worker_id

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

    def teardown(self, *args: Any, **kwargs: Any) -> None:
        """Teardown the runner and clean up all resources.
@@ -469,7 +469,7 @@ class LitAgentRunner(Runner[T_task]):

            start_time = time.time()
            async with self._tracer.trace_context(
-                name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
+                name=rollout_id, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
            ):
                await self._trigger_hooks(
                    hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
@@ -219,6 +219,8 @@ class LightningStoreServer(LightningStore):
            which runs the server in a separate thread.
        launcher_args: The arguments to use for the server launcher.
            It's not allowed to set `host`, `port`, `launch_mode` together with `launcher_args`.
+        n_workers: The number of workers to run in the server. Only applicable for `mp` launch mode.
+        prometheus: Whether to enable Prometheus metrics.
    """

    def __init__(
@@ -229,6 +231,8 @@ class LightningStoreServer(LightningStore):
        cors_allow_origins: Sequence[str] | str | None = None,
        launch_mode: LaunchMode = "thread",
        launcher_args: PythonServerLauncherArgs | None = None,
+        n_workers: int = 1,
+        prometheus: bool = False,
    ):
        super().__init__()
        self.store = store
@@ -265,6 +269,7 @@ class LightningStoreServer(LightningStore):
            app=self.app,
            args=self.launcher_args,
        )
+        self._prometheus = prometheus

        self._lock: threading.Lock = threading.Lock()
        self._cors_allow_origins = self._normalize_cors_origins(cors_allow_origins)
@@ -309,6 +314,7 @@ class LightningStoreServer(LightningStore):
        return {
            "launcher_args": self.launcher_args,
            "server_launcher": self.server_launcher,
+            "_prometheus": self._prometheus,
            "_owner_pid": self._owner_pid,
        }

@@ -326,6 +332,7 @@ class LightningStoreServer(LightningStore):
        self.store = None
        self.launcher_args = state["launcher_args"]
        self.server_launcher = state["server_launcher"]
+        self._prometheus = state["_prometheus"]
        self._owner_pid = state["_owner_pid"]
        self._cors_allow_origins = state.get("_cors_allow_origins")
        self._client = None
@@ -407,6 +414,11 @@ class LightningStoreServer(LightningStore):
    def _setup_routes(self):
        """Set up FastAPI routes for all store operations."""
        assert self.app is not None
+        api = APIRouter(prefix=API_V1_PREFIX)
+
+        # The outermost-layer of monitoring
+        if self._prometheus:
+            self._setup_prometheus(api=api, app=self.app)

        @self.app.middleware("http")
        async def _app_exception_handler(  # pyright: ignore[reportUnusedFunction]
@@ -461,8 +473,6 @@ class LightningStoreServer(LightningStore):
            )
            return response

-        api = APIRouter(prefix=API_V1_PREFIX)
-
        def _validate_paginated_request(
            request: Union[
                QueryRolloutsRequest,
@@ -720,6 +730,58 @@ class LightningStoreServer(LightningStore):
        # Finally, mount the dashboard assets
        self._setup_dashboard()

+    def _setup_prometheus(self, api: APIRouter, app: FastAPI):
+        """Setup Prometheus metrics endpoints."""
+        try:
+            from prometheus_client import (
+                CONTENT_TYPE_LATEST,
+                Counter,
+                Histogram,
+                generate_latest,
+            )
+        except ImportError:
+            raise ImportError(
+                "Prometheus client is not installed. Please either install it or set prometheus to False."
+            )
+
+        HTTP_REQUESTS = Counter(
+            "http_requests_total",
+            "Total HTTP requests",
+            ["method", "path", "status_code"],
+        )
+
+        # TODO: For multi-process scenarios, should use prometheus_client.multiprocess mode.
+        HTTP_LATENCY = Histogram(
+            "http_request_duration_seconds",
+            "Latency of HTTP requests",
+            ["method", "path"],
+            buckets=[0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10],
+        )
+
+        @app.middleware("http")
+        async def prometheus_http_middleware(  # pyright: ignore[reportUnusedFunction]
+            request: Request, call_next: Callable[[Request], Awaitable[Response]]
+        ) -> Response:
+            start = time.perf_counter()
+            response = await call_next(request)
+            elapsed = time.perf_counter() - start
+
+            path = request.url.path
+            method = request.method
+            status = response.status_code
+
+            HTTP_REQUESTS.labels(method, path, status).inc()
+            HTTP_LATENCY.labels(method, path).observe(elapsed)
+
+            return response
+
+        @api.get("/prometheus")
+        async def prometheus_metrics():  # pyright: ignore[reportUnusedFunction]
+            return Response(
+                content=generate_latest(),
+                media_type=CONTENT_TYPE_LATEST,
+            )
+
    def _setup_otlp(self, api: APIRouter):
        """Setup OTLP endpoints."""

@@ -794,13 +856,35 @@ class LightningStoreServer(LightningStore):
        - In the owner process: delegate to the in-process store.
        - In a different process: delegate to a HTTP client talking to the server.
        """
+        # If the store is zero-copy, we can just call the method directly.
+        if self.store is not None and self.store.capabilities.get("zero_copy", False):
+            return await getattr(self.store, method_name)(*args, **kwargs)
+
        if os.getpid() == self._owner_pid:
            if method_name == "wait_for_rollouts":
                # wait_for_rollouts can block for a long time; avoid holding the lock
                # so other requests can make progress while we wait.
                return await getattr(self.store, method_name)(*args, **kwargs)
-            with self._lock:
+
+            # If it's already thread-safe, we can just call the method directly.
+            # Acquiring the threading lock directly would block the event loop if it's
+            # already held by another thread (for example, the HTTP server thread).
+            # Potential fix here are needed to make it work. For example:
+            # ```
+            # acquired = self._lock.acquire(blocking=False)
+            # if not acquired:
+            #     await asyncio.to_thread(self._lock.acquire)
+            # try:
+            #     return await getattr(self.store, method_name)(*args, **kwargs)
+            # finally:
+            #     self._lock.release()
+            # ```
+            # Or we can just bypass the lock for thread-safe stores.
+            if self.store is not None and self.store.capabilities.get("thread_safe", False):
                return await getattr(self.store, method_name)(*args, **kwargs)
+            else:
+                with self._lock:
+                    return await getattr(self.store, method_name)(*args, **kwargs)
        if self._client is None:
            self._client = LightningStoreClient(self.endpoint)
        return await getattr(self._client, method_name)(*args, **kwargs)
@@ -1605,6 +1689,7 @@ class LightningStoreClient(LightningStore):
            attempt_id=attempt_id,
            sequence_id=sequence_id,
        )
+        print("created span", span)
        await self.add_span(span)
        return span

@@ -3,17 +3,31 @@
 from __future__ import annotations

 from typing import (
+    TYPE_CHECKING,
    Any,
    AsyncContextManager,
+    Awaitable,
+    Callable,
+    Dict,
    Generic,
+    List,
+    Literal,
+    Mapping,
+    MutableMapping,
    Optional,
    Sequence,
+    Tuple,
    Type,
    TypeVar,
+    cast,
 )

+if TYPE_CHECKING:
+    from typing import Self
+
 from agentlightning.types import (
    Attempt,
+    FilterField,
    FilterOptions,
    PaginatedResult,
    ResourcesUpdate,
@@ -36,13 +50,13 @@ class Collection(Generic[T]):
        raise NotImplementedError()

    def __repr__(self) -> str:
-        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self.size()})>"
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}]>"

    def item_type(self) -> Type[T]:
        """Get the type of the items in the collection."""
        raise NotImplementedError()

-    def size(self) -> int:
+    async def size(self) -> int:
        """Get the number of items in the collection."""
        raise NotImplementedError()

@@ -132,7 +146,7 @@ class Queue(Generic[T]):
    """Behaves like a deque. Supporting appending items to the end and popping items from the front."""

    def __repr__(self) -> str:
-        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self.size()})>"
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}]>"

    def item_type(self) -> Type[T]:
        """Get the type of the items in the queue."""
@@ -177,7 +191,7 @@ class Queue(Generic[T]):
        """
        raise NotImplementedError()

-    def size(self) -> int:
+    async def size(self) -> int:
        """Get the number of items in the queue."""
        raise NotImplementedError()

@@ -186,7 +200,7 @@ class KeyValue(Generic[K, V]):
    """Behaves like a dictionary. Supporting addition, updating, and deletion of items."""

    def __repr__(self) -> str:
-        return f"<{self.__class__.__name__} ({self.size()})>"
+        return f"<{self.__class__.__name__}>"

    async def has(self, key: K) -> bool:
        """Check if the given key is in the dictionary."""
@@ -204,7 +218,7 @@ class KeyValue(Generic[K, V]):
        """Pop the value for the given key, or the default value if the key is not found."""
        raise NotImplementedError()

-    def size(self) -> int:
+    async def size(self) -> int:
        """Get the number of items in the dictionary."""
        raise NotImplementedError()

@@ -251,7 +265,7 @@ class LightningCollections:
        """Dictionary (counter) of span sequence IDs."""
        raise NotImplementedError()

-    def atomic(self, *args: Any, **kwargs: Any) -> AsyncContextManager[None]:
+    def atomic(self, *args: Any, **kwargs: Any) -> AsyncContextManager[Self]:
        """Perform a atomic operation on the collections.

        Subclass may use args and kwargs to support multiple levels of atomicity.
@@ -261,3 +275,82 @@ class LightningCollections:
            **kwargs: Keyword arguments to pass to the operation.
        """
        raise NotImplementedError()
+
+    async def execute(self, callback: Callable[[Self], Awaitable[T]]) -> T:
+        """Execute the given callback within an atomic operation."""
+        async with self.atomic() as collections:
+            return await callback(collections)
+
+
+FilterMap = Mapping[str, FilterField]
+
+
+def merge_must_filters(target: MutableMapping[str, FilterField], definition: Any) -> None:
+    """Normalize a `_must` filter group into the provided mapping.
+
+    Mainly for validation purposes.
+    """
+    if definition is None:
+        return
+
+    entries: List[Mapping[str, FilterField]] = []
+    if isinstance(definition, Mapping):
+        entries.append(cast(Mapping[str, FilterField], definition))
+    elif isinstance(definition, Sequence) and not isinstance(definition, (str, bytes)):
+        for entry in definition:  # type: ignore
+            if not isinstance(entry, Mapping):
+                raise TypeError("Each `_must` entry must be a mapping of field names to operators")
+            entries.append(cast(Mapping[str, FilterField], entry))
+    else:
+        raise TypeError("`_must` filters must be provided as a mapping or sequence of mappings")
+
+    for entry in entries:
+        for field_name, ops in entry.items():
+            existing = target.get(field_name, {})
+            merged_ops: Dict[str, Any] = dict(existing)
+            for op_name, expected in ops.items():
+                if op_name in merged_ops:
+                    raise ValueError(f"Duplicate operator '{op_name}' for field '{field_name}' in must filters")
+                merged_ops[op_name] = expected
+            target[field_name] = cast(FilterField, merged_ops)
+
+
+def normalize_filter_options(
+    filter_options: Optional[FilterOptions],
+) -> Tuple[Optional[FilterMap], Optional[FilterMap], Literal["and", "or"]]:
+    """Convert FilterOptions to the internal structure and resolve aggregate logic."""
+    if not filter_options:
+        return None, None, "and"
+
+    aggregate = cast(Literal["and", "or"], filter_options.get("_aggregate", "and"))
+    if aggregate not in ("and", "or"):
+        raise ValueError(f"Unsupported filter aggregate '{aggregate}'")
+
+    # Extract normalized filters and must filters from the filter options.
+    normalized: Dict[str, FilterField] = {}
+    must_filters: Dict[str, FilterField] = {}
+    for field_name, ops in filter_options.items():
+        if field_name == "_aggregate":
+            continue
+        if field_name == "_must":
+            merge_must_filters(must_filters, ops)
+            continue
+        normalized[field_name] = cast(FilterField, dict(ops))  # type: ignore
+
+    return (normalized or None, must_filters or None, aggregate)
+
+
+def resolve_sort_options(sort: Optional[SortOptions]) -> Tuple[Optional[str], Literal["asc", "desc"]]:
+    """Extract sort field/order from the caller-provided SortOptions."""
+    if not sort:
+        return None, "asc"
+
+    sort_name = sort.get("name")
+    if not sort_name:
+        raise ValueError("Sort options must include a 'name' field")
+
+    sort_order = sort.get("order", "asc")
+    if sort_order not in ("asc", "desc"):
+        raise ValueError(f"Unsupported sort order '{sort_order}'")
+
+    return sort_name, sort_order
@@ -9,7 +9,6 @@ from collections import deque
 from contextlib import asynccontextmanager
 from typing import (
    Any,
-    AsyncGenerator,
    Deque,
    Dict,
    Iterable,
@@ -23,7 +22,6 @@ from typing import (
    Type,
    TypeVar,
    Union,
-    cast,
 )

 from agentlightning.types import (
@@ -40,9 +38,12 @@ from agentlightning.types import (

 from .base import (
    Collection,
+    FilterMap,
    KeyValue,
    LightningCollections,
    Queue,
+    normalize_filter_options,
+    resolve_sort_options,
 )

 T = TypeVar("T")  # Recommended to be a BaseModel, not a dict
@@ -58,81 +59,9 @@ ListBasedCollectionItemType = Union[
    Dict[Any, T],  # leaf node dictionary
 ]

-FilterMap = Mapping[str, FilterField]
 MutationMode = Literal["insert", "update", "upsert", "delete"]


-def _merge_must_filters(target: Dict[str, FilterField], definition: Any) -> None:
-    """Normalize a `_must` filter group into the provided mapping.
-
-    Mainly for validation purposes.
-    """
-    if definition is None:
-        return
-
-    entries: List[Mapping[str, FilterField]] = []
-    if isinstance(definition, Mapping):
-        entries.append(cast(Mapping[str, FilterField], definition))
-    elif isinstance(definition, Sequence) and not isinstance(definition, (str, bytes)):
-        for entry in definition:  # type: ignore
-            if not isinstance(entry, Mapping):
-                raise TypeError("Each `_must` entry must be a mapping of field names to operators")
-            entries.append(cast(Mapping[str, FilterField], entry))
-    else:
-        raise TypeError("`_must` filters must be provided as a mapping or sequence of mappings")
-
-    for entry in entries:
-        for field_name, ops in entry.items():
-            existing = target.get(field_name, {})
-            merged_ops: Dict[str, Any] = dict(existing)
-            for op_name, expected in ops.items():
-                if op_name in merged_ops:
-                    raise ValueError(f"Duplicate operator '{op_name}' for field '{field_name}' in must filters")
-                merged_ops[op_name] = expected
-            target[field_name] = cast(FilterField, merged_ops)
-
-
-def _normalize_filter_options(
-    filter_options: Optional[FilterOptions],
-) -> Tuple[Optional[FilterMap], Optional[FilterMap], Literal["and", "or"]]:
-    """Convert FilterOptions to the internal structure and resolve aggregate logic."""
-    if not filter_options:
-        return None, None, "and"
-
-    aggregate = cast(Literal["and", "or"], filter_options.get("_aggregate", "and"))
-    if aggregate not in ("and", "or"):
-        raise ValueError(f"Unsupported filter aggregate '{aggregate}'")
-
-    # Extract normalized filters and must filters from the filter options.
-    normalized: Dict[str, FilterField] = {}
-    must_filters: Dict[str, FilterField] = {}
-    for field_name, ops in filter_options.items():
-        if field_name == "_aggregate":
-            continue
-        if field_name == "_must":
-            _merge_must_filters(must_filters, ops)
-            continue
-        normalized[field_name] = cast(FilterField, dict(ops))  # type: ignore
-
-    return (normalized or None, must_filters or None, aggregate)
-
-
-def _resolve_sort_options(sort: Optional[SortOptions]) -> Tuple[Optional[str], Literal["asc", "desc"]]:
-    """Extract sort field/order from the caller-provided SortOptions."""
-    if not sort:
-        return None, "asc"
-
-    sort_name = sort.get("name")
-    if not sort_name:
-        raise ValueError("Sort options must include a 'name' field")
-
-    sort_order = sort.get("order", "asc")
-    if sort_order not in ("asc", "desc"):
-        raise ValueError(f"Unsupported sort order '{sort_order}'")
-
-    return sort_name, sort_order
-
-
 def _item_matches_filters(
    item: object,
    filters: Optional[FilterMap],
@@ -280,12 +209,12 @@ class ListBasedCollection(Collection[T]):
        """Return the Pydantic model type of items stored in this collection."""
        return self._item_type

-    def size(self) -> int:
+    async def size(self) -> int:
        """Return the number of items stored in the collection."""
        return self._size

    def __repr__(self) -> str:
-        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self.size()})>"
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self._size})>"

    # -------------------------------------------------------------------------
    # Internal helpers
@@ -520,8 +449,8 @@ class ListBasedCollection(Collection[T]):
            limit: Max number of items to return. Use -1 for "no limit".
            offset: Number of items to skip from the start of the *matching* items.
        """
-        filters, must_filters, filter_logic = _normalize_filter_options(filter)
-        sort_by, sort_order = _resolve_sort_options(sort)
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)

        # No sorting: stream through items and apply pagination on the fly.
@@ -574,8 +503,8 @@ class ListBasedCollection(Collection[T]):
        sort: Optional[SortOptions] = None,
    ) -> Optional[T]:
        """Return the first (or best-sorted) item that matches the given filters, or None."""
-        filters, must_filters, filter_logic = _normalize_filter_options(filter)
-        sort_by, sort_order = _resolve_sort_options(sort)
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)

        if not sort_by:
@@ -655,6 +584,9 @@ class DequeBasedQueue(Queue[T]):
    def item_type(self) -> Type[T]:
        return self._item_type

+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({len(self._items)})>"
+
    async def has(self, item: T) -> bool:
        if not isinstance(item, self._item_type):
            raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
@@ -686,7 +618,7 @@ class DequeBasedQueue(Queue[T]):
            result.append(item)
        return result

-    def size(self) -> int:
+    async def size(self) -> int:
        return len(self._items)


@@ -708,7 +640,7 @@ class DictBasedKeyValue(KeyValue[K, V]):
    async def pop(self, key: K, default: V | None = None) -> V | None:
        return self._values.pop(key, default)

-    def size(self) -> int:
+    async def size(self) -> int:
        return len(self._values)


@@ -759,9 +691,10 @@ class InMemoryLightningCollections(LightningCollections):
        return self._span_sequence_ids

    @asynccontextmanager
-    async def atomic(self, *args: Any, **kwargs: Any) -> AsyncGenerator[None, None]:
+    async def atomic(self, *args: Any, **kwargs: Any):
+        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside."""
        async with self._lock:
-            yield
+            yield self

    async def evict_spans_for_rollout(self, rollout_id: str) -> None:
        """Evict all spans for a given rollout ID.
@@ -19,6 +19,7 @@ from typing import (
    Optional,
    Set,
    TypeVar,
+    Union,
    cast,
 )

@@ -26,7 +27,7 @@ from pydantic import BaseModel

 from agentlightning.types import AttemptedRollout, PaginatedResult, Rollout, Span

-from .base import LightningStoreCapabilities, is_finished, is_running
+from .base import UNSET, LightningStoreCapabilities, Unset, is_finished, is_running
 from .collection import InMemoryLightningCollections
 from .collection_based import CollectionBasedLightningStore

@@ -68,9 +69,6 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
    In-memory implementation of LightningStore using Python data structures.
    Thread-safe and async-compatible but data is not persistent.

-    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.
@@ -127,6 +125,9 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
        # Running rollouts cache, including preparing and running rollouts
        self._running_rollout_ids: Set[str] = set()

+        # Caches the latest resources ID.
+        self._latest_resources_id: Union[str, None, Unset] = UNSET
+
    @property
    def capabilities(self) -> LightningStoreCapabilities:
        """Return the capabilities of the store."""
@@ -139,8 +140,8 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl

    async def wait_for_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
        """Wait for a specific rollout to complete with a timeout."""
-        async with self.collections.atomic():
-            rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+        async with self.collections.atomic() as collections:
+            rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
            if rollout and is_finished(rollout):
                return rollout

@@ -167,8 +168,8 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl

            # If event was set (not timeout), check if rollout is finished
            if result:
-                async with self.collections.atomic():
-                    rollout = await self.collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+                async with self.collections.atomic() as collections:
+                    rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
                    if rollout and is_finished(rollout):
                        return rollout

@@ -192,14 +193,12 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
        if rollout.rollout_id not in self._start_time_by_rollout:
            self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time

-    async def get_running_rollouts(self) -> List[AttemptedRollout]:
+    async def get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
        """Accelerated version of `get_running_rollouts` for in-memory store. Used for healthcheck."""
-        rollouts = await self.collections.rollouts.query(
-            filter={"rollout_id": {"within": list(self._running_rollout_ids)}}
-        )
+        rollouts = await collections.rollouts.query(filter={"rollout_id": {"within": list(self._running_rollout_ids)}})
        running_rollouts: List[AttemptedRollout] = []
        for rollout in rollouts.items:
-            latest_attempt = await self.collections.attempts.get(
+            latest_attempt = await collections.attempts.get(
                filter={"rollout_id": {"exact": rollout.rollout_id}},
                sort={"name": "sequence_id", "order": "desc"},
            )
@@ -220,15 +219,24 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
            raise RuntimeError(f"Spans for rollout {rollout_id} have been evicted")
        return await super().query_spans(rollout_id, attempt_id, **kwargs)

-    async def _add_span_unlocked(self, span: Span) -> Span:
+    async def _add_span_unlocked(self, collections: InMemoryLightningCollections, span: Span) -> Span:
        """In-memory store needs to maintain the span data in memory, and evict spans when memory is low."""

-        await super()._add_span_unlocked(span)
+        await super()._add_span_unlocked(collections, span)
        self._account_span_size(span)
-        await self._maybe_evict_spans()
+        await self._maybe_evict_spans(collections)

        return span

+    async def _get_latest_resources_id(self, collections: InMemoryLightningCollections) -> Optional[str]:
+        if isinstance(self._latest_resources_id, Unset):
+            latest_resources = await collections.resources.get(sort={"name": "update_time", "order": "desc"})
+            if latest_resources:
+                self._latest_resources_id = latest_resources.resources_id
+            else:
+                self._latest_resources_id = None
+        return self._latest_resources_id
+
    @staticmethod
    def _resolve_memory_threshold(
        value: float | int | None,
@@ -269,7 +277,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
        self._total_span_bytes += size
        return size

-    async def _maybe_evict_spans(self) -> None:
+    async def _maybe_evict_spans(self, collections: InMemoryLightningCollections) -> None:
        if self._total_span_bytes <= self._eviction_threshold_bytes:
            return

@@ -288,11 +296,11 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
            if self._total_span_bytes <= self._safe_threshold_bytes:
                break
            logger.debug(f"Evicting spans for rollout {rollout_id} to free up memory...")
-            await self._evict_spans_for_rollout(rollout_id)
+            await self._evict_spans_for_rollout(collections, rollout_id)
        logger.info(f"Freed up {memory_consumed_before - self._total_span_bytes} bytes of memory")

-    async def _evict_spans_for_rollout(self, rollout_id: str) -> None:
-        await self.collections.evict_spans_for_rollout(rollout_id)
+    async def _evict_spans_for_rollout(self, collections: InMemoryLightningCollections, rollout_id: str) -> None:
+        await collections.evict_spans_for_rollout(rollout_id)
        removed_bytes = self._span_bytes_by_rollout.pop(rollout_id, 0)
        if removed_bytes > 0:
            # There is something removed for real
@@ -0,0 +1,82 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import uuid
+from typing import (
+    Any,
+    Callable,
+    Mapping,
+    TypeVar,
+)
+
+from pymongo import AsyncMongoClient
+
+from .base import LightningStoreCapabilities
+from .collection.mongo import MongoClientPool, MongoLightningCollections
+from .collection_based import CollectionBasedLightningStore
+
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+
+def _generate_partition_id() -> str:
+    return "pt-" + hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+
+
+class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollections]):
+    """
+    MongoDB implementation of LightningStore using MongoDB collections.
+    Data is persistent and can be shared between multiple processes.
+
+    Args:
+        client: The MongoDB client. Could be a string URI or an instance of AsyncMongoClient.
+        database: The MongoDB database. Could be a string name or an instance of AsyncDatabase.
+            You must provide at least one of client or database.
+        partition_id: The partition id. Useful when sharing the database among multiple Agent-lightning trainers.
+    """
+
+    def __init__(
+        self,
+        *,
+        client: AsyncMongoClient[Mapping[str, Any]] | str,
+        database_name: str | None = None,
+        partition_id: str | None = None,
+    ) -> None:
+        self._auto_created_client = False
+        if isinstance(client, str):
+            self._client = AsyncMongoClient[Mapping[str, Any]](client)
+            self._auto_created_client = True
+        else:
+            self._client = client
+        if database_name is None:
+            database_name = "agentlightning"
+            logger.info("No database name provided, using default 'agentlightning'")
+
+        if partition_id is None:
+            partition_id = _generate_partition_id()
+            logger.info("No partition id provided, generated a new one: %s", partition_id)
+
+        self._client_pool = MongoClientPool(self._client)
+
+        super().__init__(collections=MongoLightningCollections(self._client_pool, database_name, partition_id))
+
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        return LightningStoreCapabilities(
+            thread_safe=True,
+            async_safe=True,
+            zero_copy=True,
+            otlp_traces=False,
+        )
+
+    async def close(self) -> None:
+        """Close the store by closing the client pool."""
+        await self._client_pool.close()
+        # If I created the client, I should close it too.
+        if self._auto_created_client:
+            await self._client.close()
@@ -4,6 +4,7 @@ from __future__ import annotations

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

@@ -114,6 +115,14 @@ class AgentOpsTracer(OtelTracer):
        Yields:
            The OpenTelemetry tracer instance to collect spans.
        """
+        if store is not None:
+            warnings.warn(
+                "store is deprecated in favor of init_worker(). It will be removed in the future.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+        else:
+            store = self._store
        with self._trace_context_sync(name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id) as tracer:
            yield tracer

@@ -164,9 +173,10 @@ class AgentOpsTracer(OtelTracer):
        try:
            yield
        except Exception as e:
-            # TODO: I'm not sure whether this will catch errors in user code.
+            # This will catch errors in user code.
            status = StatusCode.ERROR  # type: ignore
            logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}: {e}")
+            raise  # should reraise the error here so that runner can handle it
        finally:
            agentops.end_trace(trace, end_state=status)  # type: ignore

@@ -52,6 +52,18 @@ class Tracer(ParallelWorkerBase):
    ```
    """

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

-        If a store is provided, the spans will be added to the store when tracing.
-
        Args:
            name: The name for the root span of this trace context.
-            store: The store to add the spans to.
+            store: The store to add the spans to. Deprecated in favor of passing store to init_worker().
            rollout_id: The rollout ID to add the spans to.
            attempt_id: The attempt ID to add the spans to.
        """
@@ -82,7 +92,6 @@ class Tracer(ParallelWorkerBase):
        self,
        name: Optional[str] = None,
        *,
-        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
    ) -> ContextManager[Any]:
@@ -141,11 +150,14 @@ class Tracer(ParallelWorkerBase):
        return None

    @contextmanager
-    def lifespan(self):
+    def lifespan(self, store: Optional[LightningStore] = None):
        """A context manager to manage the lifespan of the tracer.

        This can be used to set up and tear down any necessary resources
        for the tracer, useful for debugging purposes.
+
+        Args:
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
        """
        has_init = False
        has_init_worker = False
@@ -153,7 +165,7 @@ class Tracer(ParallelWorkerBase):
            self.init()
            has_init = True

-            self.init_worker(0)
+            self.init_worker(0, store)
            has_init_worker = True

            yield
@@ -19,6 +19,8 @@ from opentelemetry.trace.span import (
    TraceState,
 )

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

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

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

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

    @asynccontextmanager
@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import logging
 import threading
+import warnings
 from contextlib import asynccontextmanager
 from typing import Any, AsyncGenerator, Awaitable, List, Optional

@@ -42,8 +43,8 @@ class OtelTracer(Tracer):
        self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
        self._initialized: bool = False

-    def init_worker(self, worker_id: int):
-        super().init_worker(worker_id)
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None):
+        super().init_worker(worker_id, store)
        self._initialize_tracer_provider(worker_id)

    def _initialize_tracer_provider(self, worker_id: int):
@@ -92,7 +93,18 @@ class OtelTracer(Tracer):
        if not self._lightning_span_processor:
            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")

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

    def get_last_trace(self) -> List[ReadableSpan]:
        """
@@ -0,0 +1,39 @@
+# Use a Python image with uv pre-installed
+FROM ghcr.io/astral-sh/uv:python3.12-bookworm
+
+# Setup a non-root user
+RUN groupadd --system --gid 999 nonroot \
+ && useradd --system --gid 999 --uid 999 --create-home nonroot
+
+# Install the project into `/app`
+WORKDIR /app
+
+# Enable bytecode compilation
+ENV UV_COMPILE_BYTECODE=1
+
+# Copy from the cache instead of linking since it's a mounted volume
+ENV UV_LINK_MODE=copy
+
+# Ensure installed tools can be executed out of the box
+ENV UV_TOOL_BIN_DIR=/usr/local/bin
+
+# Install the project's dependencies using the lockfile and settings
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync --locked --no-install-project --group dev --extra mongo --group core-stable
+
+# Then, add the rest of the project source code and install it
+# Installing separately from its dependencies allows optimal layer caching
+COPY . /app
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv sync --locked --group dev --extra mongo --group core-stable
+
+# Place executables in the environment at the front of the path
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Reset the entrypoint, don't invoke `uv`
+ENTRYPOINT []
+
+# Use the non-root user to run our application
+USER nonroot
@@ -0,0 +1,23 @@
+# Docker-compose file to launch a MongoDB server for development.
+# It's used to test the MongoDB store implementation.
+
+services:
+
+  mongo:
+    image: mongo:latest
+    ulimits:
+      nofile:
+        soft: 65535
+        hard: 65535
+    ports:
+      - "27017:27017"
+    command: ["mongod", "--bind_ip_all", "--replSet", "rs0"]
+    volumes:
+      - ../scripts/mongodb_init_rs_host.js:/docker-entrypoint-initdb.d/init-rs.js:ro
+      - ./data/mongo-host:/data/db
+    healthcheck:
+      test: ["CMD", "mongosh", "--eval", "db.adminCommand('ping')"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+      start_period: 30s
@@ -0,0 +1,33 @@
+services:
+
+  app:
+    extends:
+      file: compose.store.yml
+      service: app
+
+    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend memory
+
+  node-exporter:
+    image: prom/node-exporter:latest
+    # In CI you might not have full /proc, but this is OK for container-level stats
+    pid: "host"
+    network_mode: "service:app" # share network with app for simplicity
+    command:
+      - '--path.rootfs=/host'
+    volumes:
+      - '/:/host:ro,rslave'
+
+  prometheus:
+    image: prom/prometheus:latest
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--storage.tsdb.retention.time=1h'
+    volumes:
+      - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./data/prometheus:/prometheus
+    depends_on:
+      - app
+      - node-exporter
+    ports:
+      - "9090:9090"
@@ -0,0 +1,62 @@
+services:
+
+  mongo:
+    extends:
+      file: compose.mongo.yml
+      service: mongo
+
+    hostname: mongo
+    volumes:
+      - ./data/mongo-container:/data/db
+      - ../scripts/mongodb_init_rs_profiling.js:/docker-entrypoint-initdb.d/init-rs.js:ro
+
+    # This forces "mongo" to resolve to localhost ONLY inside this container.
+    # This allows rs.initiate to succeed using the hostname "mongo".
+    extra_hosts:
+      - "mongo:127.0.0.1"
+
+  app:
+    extends:
+      file: compose.store.yml
+      service: app
+
+    depends_on:
+      - mongo
+
+    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend mongo --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 --n-workers 4
+
+  mongodb-exporter:
+    image: percona/mongodb_exporter:0.47.1
+    command:
+      - '--mongodb.uri=mongodb://mongo:27017/'
+      - '--collect-all'
+      - '--mongodb.collstats-colls=agentlightning.rollouts,agentlightning.attempts,agentlightning.spans,agentlightning.resources,agentlightning.workers,agentlightning.rollout_queue,agentlightning.span_sequence_ids'
+    depends_on:
+      - mongo
+    ports:
+      - "9216:9216"
+
+  node-exporter:
+    image: prom/node-exporter:latest
+    # In CI you might not have full /proc, but this is OK for container-level stats
+    pid: "host"
+    command:
+      - '--path.rootfs=/host'
+    volumes:
+      - '/:/host:ro,rslave'
+
+  prometheus:
+    image: prom/prometheus:latest
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--storage.tsdb.retention.time=1h'
+    volumes:
+      - ./prometheus.mongo-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./data/prometheus:/prometheus
+    depends_on:
+      - app
+      - mongodb-exporter
+      - node-exporter
+    ports:
+      - "9090:9090"
@@ -0,0 +1,26 @@
+services:
+
+  app:
+    build:
+      context: ../
+      dockerfile: docker/Dockerfile.dev
+
+    ports:
+      - "4747:4747"
+
+    command: agl store --host 0.0.0.0 --port 4747
+
+    develop:
+      watch:
+        # Sync the working directory with the `/app` directory in the container
+        - action: sync
+          path: ..
+          target: /app
+          # Exclude the project virtual environment — it could be for a
+          # different platform in the container
+          ignore:
+            - .venv/
+
+        # Rebuild the image if dependencies change by checking uv.lock
+        - action: rebuild
+          path: ../uv.lock
@@ -0,0 +1,17 @@
+global:
+  scrape_interval: 2s
+  evaluation_interval: 2s
+
+scrape_configs:
+  - job_name: app
+    static_configs:
+      - targets: ["app:4747"]
+    metrics_path: /v1/prometheus
+
+  - job_name: mongodb
+    static_configs:
+      - targets: ["mongodb-exporter:9216"]
+
+  - job_name: node
+    static_configs:
+      - targets: ["node-exporter:9100"]
@@ -0,0 +1,13 @@
+global:
+  scrape_interval: 2s
+  evaluation_interval: 2s
+
+scrape_configs:
+  - job_name: app
+    static_configs:
+      - targets: ["app:4747"]
+    metrics_path: /v1/prometheus
+
+  - job_name: node
+    static_configs:
+      - targets: ["node-exporter:9100"]
@@ -0,0 +1,9 @@
+#!/bin/bash
+
+set -euo pipefail
+
+# Create data directories
+mkdir -p data/prometheus data/mongo-container data/mongo-host
+
+# Change permissions
+chmod 777 data/prometheus data/mongo-container data/mongo-host
@@ -1,18 +1,98 @@
 # Contributing Guide

-Agent Lightning thrives on community improvements, whether you are polishing docs, fixing bugs, or building new features. This guide shows the shortest path from cloning the repository to shipping a polished pull request.
+Agent Lightning gets better every time someone files a clear bug, polishes docs, improves tests, or lands a new feature. This guide collects the expectations, checklists, and tips that help you go from “I have an idea” to “my pull request just merged.”

-## Step 1. Prepare Your Environment
+## Before You Start

-### Prerequisites
+Agent-lightning is built by a small Microsoft Research team with limited reviewer hours and GPU budget. For any sizeable change (new algorithm, example, or API surface) please first discuss scope with us in [Discord](https://discord.gg/RYk7CdvDR7). Early alignment keeps your effort from being blocked late in the process.

- **Python** 3.10 or newer (we test on 3.10–3.13).
- **uv** for dependency and virtual environment management. Install it from the [official uv docs](https://docs.astral.sh/uv/getting-started/installation/).
+## Where You Can Help
+
+Pick a lane, or combine several. Just keep the discussion-first principle in mind for anything non-trivial.
+
+### Documentation Improvements
+
+Documentation improvements are the easiest way to get started. You can find more about how to write good documentations and organize documentations in the following sections. Here are some general contribution points we can think of:
+
+- Tighten language, fix typos, clarify confusing sections, or add missing links. Fresh eyes catch docs gaps best.
+- Organize content using the directories listed below so readers can actually find it.
+- Avoid duplicate prose, unrelated “how-to” guides, or translations (we cannot maintain them today).
+
+!!! note "Changes that are usually rejected"
+
+    - Copy/pasting existing docs with shallow edits.
+    - Adding a `how-to` guide that is not tied to a new example.
+    - Adding doc translations to other languages (no capacity to review/maintain yet).
+
+### Bug Fixes
+
+Bug fixes are the fastest way to get familiar with the codebase. To get started, you can:
+
+- Browse the ["good first issue"](https://github.com/microsoft/agent-lightning/labels/good%20first%20issue) and ["bug"](https://github.com/microsoft/agent-lightning/labels/bug) labels; drop a comment before you start so we can mark it as taken.
+- For fresh bugs, open an issue with reproduction steps, logs, and expected behavior before submitting a fix.
+- Keep each pull request focused, ideally avoiding breaking API changes. Larger refactors should be discussed via RFC or maintainer sync.
+
+### New Examples
+
+Examples must be curated so that we can maintain them. We generally merge only those that meet at least one (ideally several) of these criteria:
+
+- Demonstrates an agent framework or workflow that is materially different from what already exists. ([LangChain](https://www.langchain.com/) vs. [LlamaIndex](https://www.llamaindex.ai/) is not different enough; [LangChain](https://www.langchain.com/) vs. [n8n](https://n8n.io/) or [Vercel AI SDK](https://ai-sdk.dev/) is, because they either have different orchestration paradigms or differ in programming languages.)
+- Shows measurable performance gains on a **real-world** problem with a **real-world** dataset, such as tuning a search agent with Google Search API or improving a coding agent’s (e.g., Claude Code) SWE-Bench score.
+- Integrates a new algorithm, training backend, or serving stack (see “New Algorithms” below).
+- Validates scenarios that are rarely tested, such as multi-modality agents or long-lived memory/workflow agents.
+
+Bonus points for examples that:
+
+- Ship CI or self-test coverage so we know they still work as the core evolves.  **Otherwise, we would have to mark the example as unmaintained because we won't be able to test the examples manually before each release.**
+- Include a [`docs/how-to/`]({{ src("docs/how-to/") }}) guide (or a detailed README if no how-to exists) without duplicating content in multiple places.
+- Favor simple, dependency-light code over heavy abstractions.
+
+!!! warning "Please discuss first"
+
+    Examples tend to be the most time-consuming contributions for both you and reviewers. Sync with us on Discord or through an issue before diving into a new one.
+
+### Fresh Implementations of Core Modules
+
+If you are looking to extend [`Runner`][agentlightning.Runner], [`Tracer`][agentlightning.Tracer], [`Adapter`][agentlightning.Adapter], [`LightningStore`][agentlightning.LightningStore], or another core interface, here are the steps:
+
+1. File an issue or proposal first.
+2. Explain which interface you are extending, why existing implementations are insufficient, and how you intend to test compatibility with the rest of the stack (unit tests, documentation updates, example refreshes, etc.).
+3. Any API changes must be reviewed up front. DO NOT begin coding large changes before the discussion lands!
+
+### New Algorithms
+
+If you are integrating a new training/serving backend, check whether it already lives in the [Algorithm Zoo](../algorithm-zoo/index.md) or is covered in the [Examples Catalog](../how-to/examples-catalog.md). We especially welcome:
+
+- Currently unsupported or under-tested algorithms such as Supervised Fine-tuning (SFT), Direct Policy Optimization (DPO), or Monte Carlo Tree Search (MCTS).
+- Tuning [Resource][agentlightning.Resource]s that are not supported yet, such as workflows or memory.
+- Expansions of supported stacks, e.g., adding multi-modality to APO or multi-agent prompt tuning.
+- Reinforcement-learning integrations beyond our current stack of [VERL](https://github.com/volcengine/verl), [vLLM](https://vllm.ai/), [Azure OpenAI](https://azure.microsoft.com/en-us/products/ai-foundry/models/openai), and [Tinker](https://tinker-docs.thinkingmachines.ai/). Contributions using [SGLang](https://github.com/sgl-project/sglang), [TRL](https://github.com/huggingface/trl), [SkyRL](https://github.com/NovaSky-AI/SkyRL), [RLinf](https://github.com/RLinf/RLinf), [litgpt](https://github.com/Lightning-AI/litgpt), or similar are welcome.
+
+Most brand-new algorithms ultimately land as “new examples,” so read that section too. Post an issue or design doc to scope the work, reuse existing utilities, and avoid duplicating efforts. Mature, battle-tested examples graduate into the [Algorithm Zoo](../algorithm-zoo/index.md).
+
+### Ecosystem Projects
+
+Have a project that builds on Agent-lightning but does not belong in the main repo? Fork it or depend on it externally, then let us know. We can showcase notable projects in [Community Projects](../index.md) and the main [README]({{ src("README.md") }}).
+
+### Other Contribution Ideas
+
+- **Tests.** Add or improve cases in [`tests/`]({{ src("tests") }}) (unit, integration, or end-to-end).
+- **Benchmarks.** Expand [`tests/benchmark`]({{ src("tests/benchmark") }}) to stress large-scale training or rollouts.
+- **Issue triage.** Reproduce bugs, confirm whether they reproduce on `main`, or suggest short-term mitigations so maintainers can prioritize.
+
+## Contribution Workflow
+
+The steps below keep changes reviewable and CI-friendly. Follow them in order; rerun the relevant pieces if you revisit a branch later.
+
+### 1. Prepare Your Environment
+
+Minimum tooling:
+
+- **Python** 3.10+ (3.12 recommended).
+- **uv** for dependency and virtual-environment management. Install it using the [official uv docs](https://docs.astral.sh/uv/getting-started/installation/).
 - **Git** configured with your GitHub credentials.

-### Clone the Repository
-
-Fork the repo, then clone your fork and register the upstream remote so you can stay current:
+Clone your fork and point `upstream` at the official repo:

 ```bash
 git clone git@github.com:<your-username>/agent-lightning.git
@@ -20,15 +100,13 @@ cd agent-lightning
 git remote add upstream https://github.com/microsoft/agent-lightning.git
 ```

-### Install Dependencies
-
-Install the standard development toolchain:
+Install the default development stack:

 ```bash
 uv sync --group dev
 ```

-Want GPU extras, example dependencies, or other optional features? Pin everything in one pass:
+Need GPU extras or specific optional dependencies? Lock them in with one command:

 ```bash
 uv sync --frozen \
@@ -41,28 +119,22 @@ uv sync --frozen \
    --no-default-groups
 ```

-After `uv sync`, run commands with `uv run ...` (or `uv run --no-sync` once the environment is locked), or activate the virtual environment in `.venv/`.
+After `uv sync`, run commands via `uv run ...` (add `--no-sync` once the environment is locked) or activate `.venv/`.

---
+### 2. Install and Run Pre-commit

-## Step 2. Install and Run Pre-commit
-
-We enforce formatting and linting with [pre-commit](https://pre-commit.com/). Install the hooks once, then run them before every push:
+Formatting and linting are enforced through [pre-commit](https://pre-commit.com/). Install once, then run before each push:

 ```bash
 uv run pre-commit install
-
-# The following will auto-run if you have set up the pre-commit hooks to run automatically on commit.
 uv run pre-commit run --all-files --show-diff-on-failure --color=always
 ```

-Running them locally saves a CI round-trip and keeps diffs tidy.
+Once installed, the hooks run automatically on every `git commit`. Running the pre-commit hooks locally keeps CI green and diffs manageable.

---
+### 3. Branch From a Fresh `main`

-## Step 3. Branching Workflow
-
-Start from a fresh `main`, then branch for your change:
+Start all work from the latest upstream state:

 ```bash
 git fetch upstream
@@ -70,20 +142,32 @@ git checkout main
 git merge upstream/main
 ```

-Create a topic branch with one of these prefixes:
+Branch naming convention:

- `feature/<short-description>` for new features
- `fix/<short-description>` for bug fixes
- `docs/<short-description>` for documentation-only work
- `chore/<short-description>` for tooling or maintenance
+- `feature/<short-description>` for new features.
+- `fix/<short-description>` for bug fixes.
+- `docs/<short-description>` for documentation-only updates.
+- `chore/<short-description>` for tooling or maintenance.

-Stick to lowercase words separated by hyphens, e.g. `feature/async-runner-hooks`.
+Use lowercase with hyphens, e.g., `feature/async-runner-hooks`.

---
+!!! note "Where should docs or examples live?"

-## Step 4. Test Your Changes
+    Many new contributors get confused about what to put in the `docs/how-to/` directory and what to put in the `examples/` directory (particularly README files). Here is a quick reference you can refer to:

-Most updates should ship with automated checks. Preface commands with `uv run` so they use the project environment.
+    | Location | Description |
+    | --- | --- |
+    | `docs/algorithm-zoo/` | Documentation for **built-in algorithms** shipped with Agent-lightning. |
+    | `docs/how-to/` | Step-by-step **how-to guides**, usually tied to an example in `examples/`. |
+    | `docs/tutorials/` | Conceptual walkthroughs for components or workflows. See [debugging](../tutorials/debug.md) or [parallelization](../tutorials/parallelize.md) for examples. |
+    | `docs/deep-dive/` | Advanced explanations and in-depth concepts. |
+    | `examples/<name>/README.md` | Example-specific README. If any related how-to if that exists, link to it avoid duplicating the same instructions twice; write only brief instructions on how to install and run the example. Otherwise, you can make the README more detailed and self-explanatory. |
+
+    Remember to register new docs in [`mkdocs.yml`]({{ src("mkdocs.yml") }}), add examples to [examples/README]({{ src("examples/README.md") }}), and update the [Examples Catalog](../how-to/examples-catalog.md).
+
+### 4. Test and Validate
+
+Most contributions require automated checks. Prefix commands with `uv run` so they use the project environment.

 **Full test suite**

@@ -97,53 +181,43 @@ uv run pytest -v
 uv run pytest tests/path/to/test_file.py -k test_name
 ```

-**Optional/gated tests**
+**Optional/gated tests:** GPU-specific suites or API-dependent tests run automatically when the required hardware or environment variables (such as `OPENAI_API_KEY`) are present.

-GPU-specific suites or API-dependent tests run automatically when the required hardware or environment variables (such as `OPENAI_API_KEY`) are present.
-
-**Static analysis**
+**Static analysis:**

 ```bash
 uv run pyright
 ```

-Touching code under `examples/`? Each directory includes a README with example-specific smoke tests—run those too.
+If you have touched code under `examples/`, you should run the example-specific smoke tests. Each directory includes a README with example-specific smoke tests—run those too.

---
+!!! note "Build documentation when needed"

-## Step 5. Build Documentation (When Applicable)
+    Keep API references under [docs/reference]({{ src("docs/reference/") }}) up to date. Doc-only changes should still build cleanly:

-Doc changes should build cleanly before you push:
+    ```bash
+    uv run mkdocs serve --strict   # live reload
+    uv run mkdocs build --strict   # CI-equivalent
+    ```

-```bash
-uv run mkdocs serve --strict  # live reload while editing
-uv run mkdocs build --strict  # CI-equivalent validation
-```
+    `--strict` elevates warnings to errors so you catch issues before CI.

-`--strict` matches CI and promotes warnings to errors so you catch them early.
-
---
-
-## Step 6. Final Local Checks
+Before opening a PR, double-check the basics:

+- Run `uv lock` if you changed dependencies.
 - Run `uv run pre-commit run --all-files` (hooks installed via `pre-commit install` run automatically on `git commit`, but rerun them if you amended history).
- Execute the relevant test commands from Step 4.
- Validate any affected examples by following the instructions in `examples/<name>/README`.
+- Execute the relevant commands from the test list above.
+- Validate each affected example via its README instructions.

---
+### 5. Open a Pull Request

-## Step 7. Open a Pull Request
-
-1. Push your branch to your fork:
+1. Push your branch:
   ```bash
   git push origin <branch-name>
   ```
 2. Open a PR against `microsoft/agent-lightning:main`.
-3. Complete the PR template with:
-    - A concise summary of the change.
-    - The tests or commands you ran (copy from Step 4/6).
-    - Linked issues (use `Fixes #123` to auto-close).
-4. Attach screenshots or terminal output when it clarifies behavior.
-5. Address review feedback promptly. Use focused commits, and consider `git commit --fixup` for follow-up adjustments.
+3. Fill out the template with a concise summary, the commands/tests you ran, and linked issues (use `Fixes #123` syntax to auto-close).
+4. Include screenshots or logs if they clarify behavior.
+5. Address review feedback promptly. Follow-up tweaks work best as focused commits; `git commit --fixup` is handy for reviewer-suggested edits.

-Thanks for contributing—every improvement grows the Agent Lightning community!
+Thanks for contributing! every improvement strengthens the Agent Lightning community!
@@ -0,0 +1,81 @@
+# Examples Catalog
+
+!!! tip "Want to Contribute?"
+
+    We welcome contributions to the examples catalog! Please refer to the [Contributing](../community/contributing.md) guide for more details.
+
+<div class="grid cards" markdown>
+
+-   :material-robot:{ .lg .middle } __APO room selector__
+
+    ---
+
+    Prompt-optimize a room-booking agent with the built-in APO algorithm, then contrast it with the write-your-own algorithm and debugging workflows in the tutorials. Pairs well with the [Train the First Agent how-to]({{ src("docs/how-to/train-first-agent.md") }}) and the [Write the First Algorithm guide]({{ src("docs/how-to/write-first-algorithm.md") }}).
+
+    [:octicons-repo-24: Browse source]({{ src("examples/apo") }})
+
+-   :material-cloud-sync:{ .lg .middle } __Azure OpenAI SFT__
+
+    ---
+
+    Run a supervised fine-tuning loop against Azure OpenAI: roll out the capital-lookup agent, turn traces into JSONL, launch fine-tunes, and redeploy the resulting checkpoints through Azure CLI.
+
+    [:octicons-repo-24: Browse source]({{ src("examples/azure") }})
+
+-   :material-calculator:{ .lg .middle } __Calc-X VERL math__
+
+    ---
+
+    VERL-based reinforcement learning setup for a math-reasoning agent that uses AutoGen plus an MCP calculator tool to solve Calc-X problems end to end.
+
+    [:octicons-repo-24: Browse source]({{ src("examples/calc_x") }})
+
+-   :material-view-grid:{ .lg .middle } __Minimal building blocks__
+
+    ---
+
+    Bite-sized scripts that isolate Agent-lightning primitives (e.g., LightningStore usage, LLM proxying, minimal vLLM host) so you can study each part before composing larger workflows.
+
+    [:octicons-repo-24: Browse source]({{ src("examples/minimal") }})
+
+-   :material-book-open-page-variant:{ .lg .middle } __RAG (MuSiQue)__
+
+    ---
+
+    Retrieval-Augmented Generation pipeline that preps a Wikipedia retriever via MCP and trains a MuSiQue QA agent with GRPO. Documented for historical reference (verified on Agent-lightning v0.1.x).
+
+    [:octicons-repo-24: Browse source]({{ src("examples/rag") }})
+
+-   :material-magnify:{ .lg .middle } __Search-R1 RL__
+
+    ---
+
+    Reproduction of the Search-R1 workflow that prepares its own retrieval backend, runs the rollout script, and coordinates GRPO-style training without extra orchestration layers (last validated on v0.1.x).
+
+    [:octicons-repo-24: Browse source]({{ src("examples/search_r1") }})
+
+-   :material-database:{ .lg .middle } __Spider SQL agent__
+
+    ---
+
+    LangGraph-powered text-to-SQL workflow for the Spider benchmark, combining LangChain tooling with Agent-lightning rollouts; follow along with the [how-to for training SQL agents]({{ src("docs/how-to/train-sql-agent.md") }}).
+
+    [:octicons-repo-24: Browse source]({{ src("examples/spider") }})
+
+-   :material-thought-bubble:{ .lg .middle } __Tinker integration__
+
+    ---
+
+    Adapter package ([`agl_tinker`]({{ src("examples/tinker/agl_tinker") }})) with Tinker plus sample CrewAI/OpenAI agents that feed Agent-lightning traces into Tinker’s reinforcement-learning backend for both toy and 20-Questions-style workflows.
+
+    [:octicons-repo-24: Browse source]({{ src("examples/tinker") }})
+
+-   :material-fast-forward:{ .lg .middle } __Unsloth SFT__
+
+    ---
+
+    Supervised fine-tuning loop that ranks math-agent rollouts, fine-tunes with Unsloth’s 4-bit LoRA stack, and mirrors the [Fine-tune with Unsloth recipe]({{ src("docs/how-to/unsloth-sft.md") }}).
+
+    [:octicons-repo-24: Browse source]({{ src("examples/unsloth") }})
+
+</div>
@@ -84,3 +84,24 @@ canvas[data-chart] {
  width: 100%;
  display: block;
 }
+
+/* Grid behavior */
+.md-typeset .grid {
+  grid-template-columns: repeat(auto-fit, minmax(24rem, 1fr));
+}
+
+/* Make cards fill equal height and push footer link to bottom */
+.md-typeset .grid.cards > ul > li {
+  display: flex;
+  flex-direction: column;
+  gap: 0;
+}
+
+.md-typeset .grid.cards > ul > li > hr {
+  margin: 0.5em 0;
+}
+
+.md-typeset .grid.cards > ul > li > :last-child {
+  margin-top: auto; /* pushes the last element (Browse source) to bottom */
+  padding-top: 0.5em;
+}
@@ -4,14 +4,14 @@ This catalog highlights the examples shipped with Agent-lightning.

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

-*NOTE: CI status avoid taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. Each example's own README also displays its `examples-*` workflow status whenever the project is maintained by CI.*
+*NOTE: CI status avoids taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. Each example's own README also displays its `examples-*` workflow status whenever the project is maintained by CI.*
@@ -376,14 +376,22 @@ class AzureOpenAIFinetune(Algorithm):
                training_file=train_file_id,
                model=base_model,
                seed=self.seed,
-                hyperparameters={
-                    "batch_size": self.finetune_batch_size,
-                    "learning_rate_multiplier": self.finetune_learning_rate,
-                    "n_epochs": self.finetune_epochs,
+                method={
+                    "type": "supervised",
+                    "supervised": {
+                        "hyperparameters": {
+                            "batch_size": self.finetune_batch_size,
+                            "learning_rate_multiplier": self.finetune_learning_rate,
+                            "n_epochs": self.finetune_epochs,
+                        }
+                    },
                },
                # TODO: continuously adding suffix will make model names very long after a few iterations
                # investigate if we can just specify the fine-tuned model name directly
                suffix=f"v{next_iteration:02d}",
+                # NOTE: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/fine-tuning
+                # Other options are "GlobalStandard" and "Standard"
+                extra_body={"trainingType": "GlobalStandard"},
            )
            job_id = job.id
            self._log_info("Fine-tuning job %s created for base model %s.", job_id, base_model)
@@ -444,6 +452,13 @@ class AzureOpenAIFinetune(Algorithm):

        return LLM(endpoint=self.azure_openai_endpoint, model=deployment_name, api_key=self.azure_openai_api_key)

+    def cleanup_deployments(self) -> None:
+        """Delete all deployments created by this algorithm instance."""
+        for deployment_name in self._created_deployments:
+            self._log_info("Cleaning up deployment %s.", deployment_name)
+            self._delete_deployment(deployment_name)
+        self._created_deployments = []
+
    def _filter_training_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Select the top-performing examples and strip reward metadata.

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

+import argparse
+
 import pandas as pd
 from aoai_finetune import AzureOpenAIFinetune
 from capital_agent import capital_agent
@@ -10,14 +12,23 @@ from agentlightning import TraceToMessages, Trainer, setup_logging
 console = Console()


+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Train Capital Agent with Azure OpenAI Finetuning")
+    parser.add_argument("--n-iterations", type=int, default=3, help="Number of finetuning iterations")
+    parser.add_argument("--cleanup", action="store_true", help="Cleanup finetuned deployments after training")
+    return parser.parse_args()
+
+
 def main():
    setup_logging()
+    args = parse_args()
    finetune_algo = AzureOpenAIFinetune(
        base_deployment_name="gpt-4.1-mini",
        finetuned_deployment_name="gpt-4.1-mini-ft",
        base_model_name="gpt-4.1-mini-2025-04-14",
        finetune_every_n_rollouts=24,
        data_filter_ratio=0.6,
+        n_iterations=args.n_iterations,
    )

    trainer = Trainer(n_runners=2, algorithm=finetune_algo, adapter=TraceToMessages())
@@ -27,11 +38,15 @@ def main():

    console.print(f"Training on {len(train_dataset)} samples, validating on {len(val_dataset)} samples.")  # type: ignore

-    trainer.fit(
-        capital_agent,
-        train_dataset=train_dataset.to_dict(orient="records"),  # type: ignore
-        val_dataset=val_dataset.to_dict(orient="records"),  # type: ignore
-    )
+    try:
+        trainer.fit(
+            capital_agent,
+            train_dataset=train_dataset.to_dict(orient="records"),  # type: ignore
+            val_dataset=val_dataset.to_dict(orient="records"),  # type: ignore
+        )
+    finally:
+        if args.cleanup:
+            finetune_algo.cleanup_deployments()


 if __name__ == "__main__":
@@ -34,7 +34,7 @@ async def send_traces_via_otel(use_client: bool = False):
        store = LightningStoreClient("http://localhost:45993")
    rollout = await store.start_rollout(input={"origin": "write_traces_example"})

-    with tracer.lifespan():
+    with tracer.lifespan(store):
        # Initialize the capture of one single trace for one single rollout
        async with tracer.trace_context(
            "trace-manual", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
@@ -89,10 +89,10 @@ async def send_traces_via_agentops(use_client: bool = False):

    # Initialize the tracer lifespan
    # One lifespan can contain multiple traces
-    with tracer.lifespan():
+    with tracer.lifespan(store):
        # Initialize the capture of one single trace for one single rollout
        async with tracer.trace_context(
-            "trace-1", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+            "trace-1", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
        ):
            openai_client = AsyncOpenAI()
            response = await openai_client.chat.completions.create(
@@ -102,6 +102,7 @@ async def send_traces_via_agentops(use_client: bool = False):
                    {"role": "user", "content": "Hello, what's your name?"},
                ],
            )
+            console.print(response)
            assert response.choices[0].message.content is not None
            assert "chatgpt" in response.choices[0].message.content.lower()

@@ -27,7 +27,8 @@ from tinker_cookbook.renderers import Message as TinkerMessage
 from tinker_cookbook.renderers import Renderer
 from tinker_cookbook.renderers import ToolCall as TinkerToolCall
 from tinker_cookbook.renderers import get_renderer
-from transformers import AutoTokenizer, PreTrainedTokenizer
+from tinker_cookbook.tokenizer_utils import get_tokenizer
+from transformers import PreTrainedTokenizer

 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store import LightningStore
@@ -186,6 +187,13 @@ class TinkerLLM(CustomLLM):
                if not self._validate_role(role):
                    assert False, "This should never happen"
                content = parsed_response["content"]
+                # NOTE(yuge): I thought about adding this to make it more robust to empty responses,
+                # but later I found it's a configuration error in my renderer. So I think it's better
+                # to just log a warning and go with the default path.
+                # if not content:
+                #     raise ValueError("Parsed content is empty. Original response: " + str(response))
+                if not content:
+                    logger.warning("Parsed content is empty. Original response: " + str(response))
                tool_calls = parsed_response.get("tool_calls", None)
                if tool_calls:
                    tool_calls = [self._parse_tool_call(tool_call) for tool_call in tool_calls]
@@ -291,7 +299,8 @@ def create_llm_proxy(
    """
    service_client = tinker.ServiceClient()
    sampling_client = service_client.create_sampling_client(base_model=model_name)
-    tokenizer = cast(PreTrainedTokenizer, AutoTokenizer.from_pretrained(model_name))  # type: ignore
+
+    tokenizer = get_tokenizer(model_name)
    tinker_llm = TinkerLLM(
        model_name=model_name,
        sampling_client=sampling_client,
@@ -306,5 +315,13 @@ def create_llm_proxy(
        num_retries=2,
        # Must use thread mode here because otherwise the Tinker sampling client will hang.
        launch_mode="thread",
-        callbacks=["opentelemetry"] if add_return_token_ids else None,
+        # If not adding return token ids, we need to add the opentelemetry callback.
+        # Otherwise, we set it to default.
+        callbacks=["opentelemetry"] if not add_return_token_ids else None,
+        # Lengthened timeout
+        litellm_config={
+            "router_settings": {
+                "timeout": 300,
+            }
+        },
    )
@@ -112,7 +112,7 @@ def run_algo():
            group_size=4,
            seed=42,
        ),
-        renderer_name="qwen3",
+        renderer_name="qwen3_instruct",
        model_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
        log_path="logs/hello",
        max_tokens=32,
@@ -151,21 +151,32 @@ def spawn_runners(*, n_runners: int) -> None:
        runner.join()


-def oneclick():
+def oneclick(ci: bool = False):
    """Run integrated training with algorithm and runners in one process.

    This is the simplest way to run the example, as it handles spawning
    the store, algorithm, and runners automatically.
+
+    Args:
+        ci: Whether to run in CI mode. Fast verification.
    """
+    if ci:
+        # Use smaller batch size and group size for faster verification.
+        batch_size = 4
+        group_size = 2
+    else:
+        batch_size = 16
+        group_size = 4
+
    config = Config(
        learning_rate=1e-5,
        dataset_builder=AGLDatasetBuilder(
-            batch_size=16,
-            group_size=4,
+            batch_size=batch_size,
+            group_size=group_size,
            seed=42,
            n_epochs=1,
        ),
-        renderer_name="qwen3",
+        renderer_name="qwen3_instruct",
        model_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
        log_path="logs/hello",
        max_tokens=32,
@@ -182,23 +193,36 @@ def oneclick():
        n_runners=8,
        port=_find_available_port(),
    )
-    trainer.fit(hello, train_dataset=[str(i) for i in range(1000)], val_dataset=[str(i) for i in range(1000, 1024)])
+
+    if ci:
+        # For faster verification, use a smaller dataset.
+        train_dataset = [str(i) for i in range(16)]
+        val_dataset = [str(i) for i in range(100, 108)]
+    else:
+        train_dataset = [str(i) for i in range(1000)]
+        val_dataset = [str(i) for i in range(1000, 1024)]
+    trainer.fit(hello, train_dataset=train_dataset, val_dataset=val_dataset)


 def main():
    """Entry point for the hello example script."""
    parser = argparse.ArgumentParser(description="Train a hello echo agent with Agent-lightning + Tinker.")
    parser.add_argument("mode", type=str, choices=["algo", "runner", "oneclick"])
+    parser.add_argument("--ci", action="store_true", help="Run in CI mode. Fast verification.")

    args = parser.parse_args()

+    if args.ci:
+        if args.mode != "oneclick":
+            raise ValueError("CI mode only supports oneclick mode.")
+
    agl.setup_logging()
    if args.mode == "algo":
        run_algo()
    elif args.mode == "runner":
        spawn_runners(n_runners=8)
    elif args.mode == "oneclick":
-        oneclick()
+        oneclick(ci=args.ci)


 if __name__ == "__main__":
@@ -57,6 +57,7 @@ async def evaluate_q20(
    output_file: str,
    dataset_path: str,
    seed: Optional[int] = 42,
+    ci: bool = False,
 ):
    """Evaluate a model on the 20 Questions game.

@@ -67,6 +68,7 @@ async def evaluate_q20(
        output_file: Where to append JSONL results.
        dataset_path: CSV file containing category and answer columns.
        seed: Optional random seed for shuffling the dataset; ``None`` disables deterministic shuffling.
+        ci: Whether to run in CI mode. Fast verification.
    """

    store = LightningStoreThreaded(InMemoryLightningStore())
@@ -79,7 +81,11 @@ async def evaluate_q20(
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if model_name.startswith("Qwen/"):
-        llm_proxy = create_llm_proxy(model_name, "qwen3", port, store, add_return_token_ids=False)
+        llm_proxy = create_llm_proxy(model_name, "qwen3_instruct", port, store, add_return_token_ids=False)
+    elif model_name.startswith("GPT-OSS"):
+        llm_proxy = create_llm_proxy(model_name, "gpt_oss_no_sysprompt", port, store, add_return_token_ids=False)
+    elif model_name.startswith("meta-llama"):
+        llm_proxy = create_llm_proxy(model_name, "llama3", port, store, add_return_token_ids=False)
    else:
        console.print(f"Assuming {model_name} is an OpenAI model.")
        llm_proxy = LLMProxy(
@@ -101,11 +107,17 @@ async def evaluate_q20(
    current_model_list = llm_proxy.model_list.copy()
    if not any(model["model_name"] == answerer_model_name for model in current_model_list):
        current_model_list.append(
-            {"model_name": answerer_model_name, "litellm_params": {"model": "openai/" + answerer_model_name}}
+            {
+                "model_name": answerer_model_name,
+                "litellm_params": {"model": "openai/" + answerer_model_name, "timeout": 180},
+            }
        )
    if not any(model["model_name"] == search_model_name for model in current_model_list):
        current_model_list.append(
-            {"model_name": search_model_name, "litellm_params": {"model": "openai/" + search_model_name}}
+            {
+                "model_name": search_model_name,
+                "litellm_params": {"model": "openai/" + search_model_name, "timeout": 180},
+            }
        )
    llm_proxy.update_model_list(current_model_list)
    console.print("Model list:", llm_proxy.model_list)
@@ -113,7 +125,7 @@ async def evaluate_q20(
    try:
        await llm_proxy.start()
        player_llm = CrewLLM(
-            model="openai/" + model_name, base_url=f"http://localhost:{port}/v1", api_key="dummy", timeout=60.0
+            model="openai/" + model_name, base_url=f"http://localhost:{port}/v1", api_key="dummy", timeout=180.0
        )
        answer_llm = CrewLLM(
            model="openai/" + answerer_model_name,
@@ -121,6 +133,7 @@ async def evaluate_q20(
            api_key="dummy",
            reasoning_effort="low",
            response_format=AnswererResponse,
+            timeout=180.0,
        )
        search_tool = (
            SearchTool(
@@ -129,16 +142,17 @@ async def evaluate_q20(
                    base_url=f"http://localhost:{port}/v1",
                    api_key="dummy",
                    reasoning_effort="none",
-                    timeout=60.0,
+                    timeout=180.0,
                )
            )
            if search
            else None
        )
+        n_samples = len(df) if not ci else 5
        sampled_df = (
-            df.sample(n=len(df), random_state=seed)  # type: ignore
+            df.sample(n=n_samples, random_state=seed)  # type: ignore
            if seed is not None
-            else df.sample(n=len(df))  # type: ignore
+            else df.sample(n=n_samples)  # type: ignore
        )
        for index, row in sampled_df.iterrows():  # type: ignore
            if search_tool:
@@ -154,6 +168,10 @@ async def evaluate_q20(
                )
                result_json: dict[str, Any] = {"index": index, **flow.state.model_dump()}
            except Exception as e:
+                # If on CI, directly raise the exception
+                if ci:
+                    raise
+
                result_json = {
                    "index": index,
                    "answer": row["answer"],
@@ -163,6 +181,12 @@ async def evaluate_q20(
                }
            with output_path.open("a") as f:
                f.write(json.dumps(result_json) + "\n")
+
+        if ci:
+            df_result = pd.read_json(output_path, lines=True)  # type: ignore
+            print(f"Evaluation results:\n{df_result.to_dict(orient='records')}")  # type: ignore
+            assert len(df_result["correct"].dropna()) == n_samples, f"{n_samples} evaluation results are required in CI mode."  # type: ignore
+            assert df_result["correct"].sum() > 0, "At least one correct evaluation result is required in CI mode."  # type: ignore
    finally:
        await llm_proxy.stop()

@@ -206,6 +230,11 @@ def main(argv: Optional[List[str]] = None) -> None:
        default=42,
        help="Random seed for shuffling the dataset. Use -1 to disable deterministic shuffling.",
    )
+    parser.add_argument(
+        "--ci",
+        action="store_true",
+        help="Run in CI mode (smaller dataset, smaller batch).",
+    )

    args = parser.parse_args(argv)
    asyncio.run(
@@ -216,6 +245,7 @@ def main(argv: Optional[List[str]] = None) -> None:
            output_file=args.output_file,
            dataset_path=args.dataset,
            seed=None if args.seed == -1 else args.seed,
+            ci=args.ci,
        )
    )

@@ -128,13 +128,19 @@ async def q20_agent(task: Q20Task, llm: agl.LLM, rollout: agl.Rollout) -> None:
        # agl.emit_reward(0.0)


-def dry_run():
+def dry_run(model: Literal["qwen4b", "qwen30b"]):
    """Run a quick dry-run test of the 20 Questions training setup.

    Uses in-memory store and processes 4 sample tasks to verify the setup works.
    """
    store = agl.LightningStoreThreaded(agl.InMemoryLightningStore())
-    llm_proxy = create_llm_proxy("Qwen/Qwen3-30B-A3B-Instruct-2507", "qwen3_instruct", store=store)
+    if model == "qwen4b":
+        model_name = "Qwen/Qwen3-4B-Instruct-2507"
+    elif model == "qwen30b":
+        model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+    else:
+        raise ValueError(f"Invalid model: {model}")
+    llm_proxy = create_llm_proxy(model_name, "qwen3_instruct", store=store)
    trainer = agl.Trainer(
        n_runners=2,
        initial_resources={"llm": llm_proxy.as_resource()},
@@ -150,7 +156,7 @@ def dry_run():
        asyncio.run(llm_proxy.stop())


-async def algo(search: bool, model: Literal["qwen4b", "qwen30b"], port: int):
+async def algo(search: bool, model: Literal["qwen4b", "qwen30b"], port: int, ci: bool = False):
    """Run the training algorithm for 20 Questions.

    Args:
@@ -167,10 +173,10 @@ async def algo(search: bool, model: Literal["qwen4b", "qwen30b"], port: int):

    if model == "qwen4b":
        model_name = "Qwen/Qwen3-4B-Instruct-2507"
-        renderer_name = "qwen3"
+        renderer_name = "qwen3_instruct"
    elif model == "qwen30b":
        model_name = "Qwen/Qwen3-30B-A3B-Instruct-2507"
-        renderer_name = "qwen3"
+        renderer_name = "qwen3_instruct"
    else:
        raise ValueError(f"Invalid model: {model}")

@@ -178,16 +184,27 @@ async def algo(search: bool, model: Literal["qwen4b", "qwen30b"], port: int):

    llm_proxy_port = _find_available_port()

+    if ci:
+        train_dataset = cast(agl.Dataset[Q20Task], train_dataset[:2])  # type: ignore
+        test_dataset = cast(agl.Dataset[Q20Task], test_dataset[:2])  # type: ignore
+        group_size = 2
+        batch_size = 2
+        n_epochs = 1
+    else:
+        group_size = 16
+        batch_size = 16
+        n_epochs = 10
+
    config = Config(
        learning_rate=1e-4,
        dataset_builder=AGLDatasetBuilder(
            train_dataset=train_dataset,
            val_dataset=test_dataset,
-            batch_size=16,
+            batch_size=batch_size,
            shuffle=True,
-            group_size=16,
+            group_size=group_size,
            seed=17,
-            n_epochs=10,
+            n_epochs=n_epochs,
        ),
        lora_rank=16,
        renderer_name=renderer_name,
@@ -315,12 +332,12 @@ def runner(port: int = 4747, n_runners: int = 2):
    trainer.fit(q20_agent)


-def _run_dryrun(_args: argparse.Namespace) -> None:
-    dry_run()
+def _run_dryrun(args: argparse.Namespace) -> None:
+    dry_run(model=args.model)


 def _run_algo(args: argparse.Namespace) -> None:
-    asyncio.run(algo(search=args.search, model=args.model, port=args.port))
+    asyncio.run(algo(search=args.search, model=args.model, port=args.port, ci=args.ci))


 def _run_runner(args: argparse.Namespace) -> None:
@@ -337,6 +354,9 @@ def main() -> None:
    subparsers = parser.add_subparsers(dest="command", required=True)

    dryrun_parser = subparsers.add_parser("dryrun", help="Run the in-memory dry run.")
+    dryrun_parser.add_argument(
+        "--model", choices=("qwen4b", "qwen30b"), default="qwen30b", help="Model variant to train."
+    )
    dryrun_parser.set_defaults(func=_run_dryrun)

    algo_parser = subparsers.add_parser("algo", help="Launch the full training algorithm.")
@@ -348,6 +368,7 @@ def main() -> None:
        default="qwen30b",
        help="Model variant to train.",
    )
+    algo_parser.add_argument("--ci", action="store_true", help="Run in CI mode (smaller dataset, smaller batch).")
    algo_parser.set_defaults(func=_run_algo)

    algo_verl_parser = subparsers.add_parser("verl", help="Launch the full training algorithm with VERL.")
@@ -44,19 +44,20 @@ async def test_tracer():
    )
    tinker_llm.rewrite_litellm_custom_providers()

-    store = InMemoryLightningStore()
+    store = LightningStoreThreaded(InMemoryLightningStore())
    rollout = await store.start_rollout("dummy", "train")
    llm_proxy = LLMProxy(
        port=4000,
        store=store,
        model_list=tinker_llm.as_model_list(),
        num_retries=0,
+        launch_mode="thread",
    )

    try:
        tracer = AgentOpsTracer()
        tracer.init()
-        tracer.init_worker(0)
+        tracer.init_worker(worker_id=0, store=store)

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

        async with tracer.trace_context(
-            name="test_llm", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+            name="test_llm", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
        ):
            response = client.chat.completions.create(
                model=model_name,
@@ -96,6 +97,8 @@ async def test_tracer():
        adapter = TracerTraceToTriplet()
        trajectory = reconstruct_transitions(spans, adapter, rollout.rollout_id)
        print(trajectory)
+        assert len(trajectory.transitions) > 0
+        assert len(trajectory.transitions[0].ac.tokens) > 0
    finally:
        console.print("Stopping LLM proxy...")
        await llm_proxy.stop()
@@ -162,4 +165,4 @@ async def test_llm_proxy():


 if __name__ == "__main__":
-    asyncio.run(test_llm_proxy())
+    asyncio.run(test_tracer())
@@ -49,6 +49,10 @@ markdown_extensions:
  - toc:
      permalink: true
  - attr_list
+  - md_in_html
+  - pymdownx.emoji:
+      emoji_index: !!python/name:material.extensions.emoji.twemoji
+      emoji_generator: !!python/name:material.extensions.emoji.to_svg

 plugins:
  - search
@@ -98,6 +102,7 @@ nav:
    - Train the First Agent: how-to/train-first-agent.md
    - Write the First Algorithm: how-to/write-first-algorithm.md
  - How-To Recipes:
+    - Examples Catalog: how-to/examples-catalog.md
    - SFT with Unsloth: how-to/unsloth-sft.md
    - Train SQL Agent with RL: how-to/train-sql-agent.md
  - Learning More:
@@ -38,6 +38,11 @@ verl = [
  "vllm>=0.8.4", # Due to interface change of ExternalZeroMQDistributedExecutor
 ]

+# Store-related dependencies.
+mongo = [
+  "pymongo",
+]
+
 [project.scripts]
 agl = "agentlightning.cli:main"

@@ -60,6 +65,7 @@ dev = [
  "mkdocs-git-authors-plugin",
  "mkdocs-macros-plugin",
  "mkdocs-autorefs",
+  "prometheus-client",
 ]
 experiment = [
  "random-word",
@@ -93,7 +99,10 @@ torch-stable = [
  "torch>=2.8.0",
  "torchvision>=0.23.0",
  "transformers>=4.55.0",
-  "vllm>=0.10.2",
+  # vLLM 0.11.1 requires PyTorch 2.9.0, which is incompatible with flash-attn
+  # https://github.com/Dao-AILab/flash-attention/issues/1967
+  # Similar issues with vLLM 0.11.2
+  "vllm>=0.10.2,!=0.11.1,!=0.11.2",
  # LiteLLM can then be upgraded with new vLLM
  "litellm[proxy]>=1.78",
 ]
@@ -182,7 +191,8 @@ sql = [
  "nltk",
 ]
 crewai = [
-  "crewai[tools]>=1.2.0",
+  # https://github.com/crewAIInc/crewAI/issues/3959
+  "crewai[tools]==1.2.0",
 ]

 # Summarize into large installable groups.
@@ -233,7 +243,7 @@ torch = [
  { index = "pytorch-cu128", group = "torch-cu128" },
  { index = "pytorch-cpu", group = "torch-cpu" },
 ]
-tinker_cookbook = { git = "https://github.com/thinking-machines-lab/tinker-cookbook", rev = "20e26a629797188aa8c6f34474b0d4757b20b90d" }
+tinker_cookbook = { git = "https://github.com/thinking-machines-lab/tinker-cookbook" }

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

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

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

  "pythonVersion": "3.12",
@@ -25,7 +25,12 @@ sudo apt-get install -y \
    tmux \
    vim \
    git-lfs \
-    nodejs
+    nodejs \
+    gnupg2 \
+    apt-transport-https \
+    ca-certificates \
+    gnupg \
+    lsb-release

 git lfs install

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

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

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

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

+# Configurable port (first CLI argument, or default to 12306)
+PORT="${1:-12306}"
+
 # Launch LiteLLM Proxy in background
-echo "Starting LiteLLM Proxy on port 12306..."
-nohup uv run litellm --config scripts/litellm_ci.yaml --port 12306 &
+echo "Starting LiteLLM Proxy on port ${PORT}..."
+nohup uv run litellm --config scripts/litellm_ci.yaml --port "${PORT}" &

 # Wait for the server to be ready
 echo "Waiting for LiteLLM Proxy to start..."
 for i in {1..30}; do
-  if curl -s http://localhost:12306/v1/models > /dev/null; then
+  if curl -s "http://localhost:${PORT}/v1/models" > /dev/null; then
    echo "LiteLLM Proxy is up!"
    break
  fi
-  echo "Waiting... ($i)"
+  echo "Waiting... (${i})"
  # Wait for 2 seconds before checking again
  sleep 2
 done

 # Run sanity check
 echo "Running sanity check..."
-export OPENAI_BASE_URL="http://localhost:12306/"
+export OPENAI_BASE_URL="http://localhost:${PORT}/"
 export OPENAI_API_KEY="dummy"
 uv run scripts/litellm_sanity_check.py

@@ -8,7 +8,7 @@ import openai


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

@@ -0,0 +1,9 @@
+# MongoDB Development Setup
+
+This script is used to setup MongoDB for development.
+
+## Usage
+
+```bash
+docker compose up -d
+```
@@ -0,0 +1,10 @@
+services:
+  mongo:
+    image: mongo:latest
+    container_name: mongo-dev
+    ports:
+      - "27017:27017"
+    command: ["mongod", "--bind_ip_all", "--replSet", "rs0"]
+    volumes:
+      - ./data:/data/db
+      - ./init-rs.js:/docker-entrypoint-initdb.d/init-rs.js:ro
@@ -0,0 +1,6 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+rs.initiate({
+  _id: "rs0",
+  members: [{ _id: 0, host: "localhost:27017" }],
+});
@@ -0,0 +1,9 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+// MongoDB replica set initialization script.
+// Use this if you are accessing MongoDB from the **host**.
+
+rs.initiate({
+  _id: "rs0",
+  members: [{ _id: 0, host: "localhost:27017" }],
+});
@@ -0,0 +1,12 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+// MongoDB replica set initialization script.
+// Use this if you are accessing MongoDB from another **container**.
+// `mongodb_init_rs_host.js` is the counterpart if accessing from the host.
+
+rs.initiate({
+  _id: "rs0",
+  members: [{ _id: 0, host: "mongo:27017" }],
+});
+
+db.setProfilingLevel(2);
@@ -22,3 +22,9 @@
 {"request": {"messages": [{"content": "Return 1.0 if the answer is 8, else 0.0.", "role": "system"}, {"role": "user", "content": "8"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"response": {"title": "Response", "type": "number"}}, "required": ["response"], "title": "OutputType", "type": "object", "additionalProperties": false}}}, "stream": false}, "response": {"id": "chatcmpl-BluaMcYjJNKVhfvlLVLuexmE1srjB", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "created": 1750758630, "choices": [{"index": 0, "message": {"role": "assistant", "content": "{\"response\":1.0}"}, "finish_reason": "stop"}], "system_fingerprint": "fp_178c8d546f", "usage": {"prompt_tokens": 65, "completion_tokens": 8, "total_tokens": 73, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}}}}
 {"request": {"messages": [{"content": "If the question is about math, handoff to MathAgent. Otherwise, handoff to HistoryAgent.", "role": "system"}, {"role": "user", "content": "Who was the first president of the US?"}], "model": "gpt-4.1-mini", "stream": false, "tools": [{"type": "function", "function": {"name": "transfer_to_mathagent", "description": "Handoff to the MathAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}, {"type": "function", "function": {"name": "transfer_to_historyagent", "description": "Handoff to the HistoryAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}]}, "response": {"id": "chatcmpl-BluaNLVzmIYmNkUGGuk7iol8HHBkd", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "created": 1750758631, "choices": [{"index": 0, "message": {"role": "assistant", "tool_calls": [{"id": "call_j9TL7tbHC4v6OpqD66g3k6dL", "type": "function", "function": {"name": "transfer_to_historyagent", "arguments": "{}"}}]}, "finish_reason": "tool_calls"}], "system_fingerprint": "fp_178c8d546f", "usage": {"prompt_tokens": 97, "completion_tokens": 13, "total_tokens": 110, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}}}}
 {"request": {"messages": [{"content": "Answer history questions.", "role": "system"}, {"role": "user", "content": "Who was the first president of the US?"}, {"role": "assistant", "tool_calls": [{"id": "call_j9TL7tbHC4v6OpqD66g3k6dL", "type": "function", "function": {"name": "transfer_to_historyagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_j9TL7tbHC4v6OpqD66g3k6dL", "content": "{\"assistant\": \"HistoryAgent\"}"}], "model": "gpt-4.1-mini", "stream": false}, "response": {"id": "chatcmpl-BluaNHD93ybfyjMqAiF4HqKS93GPf", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "created": 1750758631, "choices": [{"index": 0, "message": {"role": "assistant", "content": "The first president of the United States was George Washington. He served as president from 1789 to 1797."}, "finish_reason": "stop"}], "system_fingerprint": "fp_178c8d546f", "usage": {"prompt_tokens": 53, "completion_tokens": 25, "total_tokens": 78, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}}}}
+{"request": {"messages": [{"content": "If the question is about math, handoff to MathAgent. Otherwise, handoff to HistoryAgent.", "role": "system"}, {"role": "user", "content": "What is 3+5?"}], "model": "gpt-4.1-mini", "tools": [{"type": "function", "function": {"name": "transfer_to_mathagent", "description": "Handoff to the MathAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}, {"type": "function", "function": {"name": "transfer_to_historyagent", "description": "Handoff to the HistoryAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}]}, "response": {"choices": [{"content_filter_results": {}, "finish_reason": "tool_calls", "index": 0, "logprobs": null, "message": {"annotations": [], "content": null, "refusal": null, "role": "assistant", "tool_calls": [{"function": {"arguments": "{}", "name": "transfer_to_mathagent"}, "id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "type": "function"}]}}], "created": 1763565690, "id": "chatcmpl-CdeHqLyOC89BDkmhvbxnhuYa0MV9J", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "jailbreak": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 13, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 95, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 108}}}
+{"request": {"messages": [{"content": "Add two numbers.", "role": "system"}, {"role": "assistant", "content": "For context, here is the conversation so far between the user and the previous agent:\n<CONVERSATION HISTORY>\n1. user: What is 3+5?\n2. function_call: {\"arguments\": \"{}\", \"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"name\": \"transfer_to_mathagent\", \"id\": \"__fake_id__\"}\n3. function_call_output: {\"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"output\": \"{\\\"assistant\\\": \\\"MathAgent\\\"}\"}\n</CONVERSATION HISTORY>"}, {"role": "assistant", "tool_calls": [{"id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "type": "function", "function": {"name": "transfer_to_mathagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "content": "{\"assistant\": \"MathAgent\"}"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"answer": {"title": "Answer", "type": "integer"}}, "required": ["answer"], "title": "MathOutput", "type": "object", "additionalProperties": false}}}, "tools": [{"type": "function", "function": {"name": "add", "description": "", "parameters": {"properties": {"a": {"title": "A", "type": "integer"}, "b": {"title": "B", "type": "integer"}}, "required": ["a", "b"], "title": "add_args", "type": "object", "additionalProperties": false}}}]}, "response": {"choices": [{"content_filter_results": {}, "finish_reason": "tool_calls", "index": 0, "logprobs": null, "message": {"annotations": [], "content": null, "refusal": null, "role": "assistant", "tool_calls": [{"function": {"arguments": "{\"a\":3,\"b\":5}", "name": "add"}, "id": "call_l5lgStmtMUiYuqU0piPW7tY8", "type": "function"}]}}], "created": 1763565694, "id": "chatcmpl-CdeHuHa6HcAsIlClzRimaHMtoq77q", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 18, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 253, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 271}}}
+{"request": {"messages": [{"content": "Add two numbers.", "role": "system"}, {"role": "assistant", "content": "For context, here is the conversation so far between the user and the previous agent:\n<CONVERSATION HISTORY>\n1. user: What is 3+5?\n2. function_call: {\"arguments\": \"{}\", \"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"name\": \"transfer_to_mathagent\", \"id\": \"__fake_id__\"}\n3. function_call_output: {\"call_id\": \"call_C3TwIO2oWDIH7IFjD7MdsOvo\", \"output\": \"{\\\"assistant\\\": \\\"MathAgent\\\"}\"}\n</CONVERSATION HISTORY>"}, {"role": "assistant", "tool_calls": [{"id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "type": "function", "function": {"name": "transfer_to_mathagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_C3TwIO2oWDIH7IFjD7MdsOvo", "content": "{\"assistant\": \"MathAgent\"}"}, {"role": "assistant", "tool_calls": [{"id": "call_l5lgStmtMUiYuqU0piPW7tY8", "type": "function", "function": {"name": "add", "arguments": "{\"a\":3,\"b\":5}"}}]}, {"role": "tool", "tool_call_id": "call_l5lgStmtMUiYuqU0piPW7tY8", "content": "8"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"answer": {"title": "Answer", "type": "integer"}}, "required": ["answer"], "title": "MathOutput", "type": "object", "additionalProperties": false}}}, "tools": [{"type": "function", "function": {"name": "add", "description": "", "parameters": {"properties": {"a": {"title": "A", "type": "integer"}, "b": {"title": "B", "type": "integer"}}, "required": ["a", "b"], "title": "add_args", "type": "object", "additionalProperties": false}}}]}, "response": {"choices": [{"content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "protected_material_code": {"filtered": false, "detected": false}, "protected_material_text": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}, "finish_reason": "stop", "index": 0, "logprobs": null, "message": {"annotations": [], "content": "{\"answer\":8}", "refusal": null, "role": "assistant"}}], "created": 1763565699, "id": "chatcmpl-CdeHz4XyD8FskC2OsHRrawVXbUk6A", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 11, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 278, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 289}}}
+{"request": {"messages": [{"content": "Return 1.0 if the answer is 8, else 0.0.", "role": "system"}, {"role": "user", "content": "8"}], "model": "gpt-4.1-mini", "response_format": {"type": "json_schema", "json_schema": {"name": "final_output", "strict": true, "schema": {"properties": {"response": {"title": "Response", "type": "number"}}, "required": ["response"], "title": "OutputType", "type": "object", "additionalProperties": false}}}}, "response": {"choices": [{"content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "protected_material_code": {"filtered": false, "detected": false}, "protected_material_text": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}, "finish_reason": "stop", "index": 0, "logprobs": null, "message": {"annotations": [], "content": "{\"response\":1.0}", "refusal": null, "role": "assistant"}}], "created": 1763565701, "id": "chatcmpl-CdeI1ifu2oOt06oTodWPpSpNFAqIN", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "jailbreak": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 8, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 65, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 73}}}
+{"request": {"messages": [{"content": "If the question is about math, handoff to MathAgent. Otherwise, handoff to HistoryAgent.", "role": "system"}, {"role": "user", "content": "Who was the first president of the US?"}], "model": "gpt-4.1-mini", "tools": [{"type": "function", "function": {"name": "transfer_to_mathagent", "description": "Handoff to the MathAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}, {"type": "function", "function": {"name": "transfer_to_historyagent", "description": "Handoff to the HistoryAgent agent to handle the request. ", "parameters": {"additionalProperties": false, "type": "object", "properties": {}, "required": []}}}]}, "response": {"choices": [{"content_filter_results": {}, "finish_reason": "tool_calls", "index": 0, "logprobs": null, "message": {"annotations": [], "content": null, "refusal": null, "role": "assistant", "tool_calls": [{"function": {"arguments": "{}", "name": "transfer_to_historyagent"}, "id": "call_tXXmrsNYcFlKrG9hMo2Yk0E3", "type": "function"}]}}], "created": 1763565706, "id": "chatcmpl-CdeI6hxlaAwzbn7Sx4ZTXZs0Hx3E5", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "jailbreak": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 13, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 97, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 110}}}
+{"request": {"messages": [{"content": "Answer history questions.", "role": "system"}, {"role": "assistant", "content": "For context, here is the conversation so far between the user and the previous agent:\n<CONVERSATION HISTORY>\n1. user: Who was the first president of the US?\n2. function_call: {\"arguments\": \"{}\", \"call_id\": \"call_tXXmrsNYcFlKrG9hMo2Yk0E3\", \"name\": \"transfer_to_historyagent\", \"id\": \"__fake_id__\"}\n3. function_call_output: {\"call_id\": \"call_tXXmrsNYcFlKrG9hMo2Yk0E3\", \"output\": \"{\\\"assistant\\\": \\\"HistoryAgent\\\"}\"}\n</CONVERSATION HISTORY>"}, {"role": "assistant", "tool_calls": [{"id": "call_tXXmrsNYcFlKrG9hMo2Yk0E3", "type": "function", "function": {"name": "transfer_to_historyagent", "arguments": "{}"}}]}, {"role": "tool", "tool_call_id": "call_tXXmrsNYcFlKrG9hMo2Yk0E3", "content": "{\"assistant\": \"HistoryAgent\"}"}], "model": "gpt-4.1-mini"}, "response": {"choices": [{"content_filter_results": {"hate": {"filtered": false, "severity": "safe"}, "protected_material_text": {"filtered": false, "detected": false}, "self_harm": {"filtered": false, "severity": "safe"}, "sexual": {"filtered": false, "severity": "safe"}, "violence": {"filtered": false, "severity": "safe"}}, "finish_reason": "stop", "index": 0, "logprobs": null, "message": {"annotations": [], "content": "The first president of the United States was George Washington. He served as president from 1789 to 1797. If you have more questions about U.S. history or any other historical events, feel free to ask!", "refusal": null, "role": "assistant"}}], "created": 1763565707, "id": "chatcmpl-CdeI78NuMQmjmeuwLhx3I0Qqcem99", "model": "gpt-4.1-mini-2025-04-14", "object": "chat.completion", "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}], "system_fingerprint": "fp_3dcd5944f5", "usage": {"completion_tokens": 46, "completion_tokens_details": {"accepted_prediction_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0}, "prompt_tokens": 179, "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, "total_tokens": 225}}}
@@ -0,0 +1,262 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Benchmarking store performance by writing and querying spans from the store."""
+
+import asyncio
+import random
+from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple
+
+import agentlightning as agl
+from agentlightning.emitter.utils import get_tracer
+
+from .utils import flatten_dict, random_dict
+
+
+def generate_attributes() -> Dict[str, Any]:
+    return flatten_dict(
+        random_dict(
+            depth=(1, 3),
+            breadth=(2, 6),
+            key_length=(3, 20),
+            value_length=(5, 300),
+        )
+    )
+
+
+@agl.rollout
+async def agent(task: str, llm: agl.LLM):
+    tracer = get_tracer()
+    rounds = random.randint(1, 10)
+    selected_round = random.randint(0, rounds - 1)
+
+    for i in range(rounds):
+        with tracer.start_as_current_span(f"agent{i}") as span:
+            # Nested Span
+            with tracer.start_as_current_span(f"round{i}_1") as span:
+                await asyncio.sleep(random.uniform(0.0, 1.0))
+                span.set_attributes(generate_attributes())
+                if i == selected_round:
+                    span.set_attribute("task", task)
+
+            # Nested Span
+            with tracer.start_as_current_span(f"round{i}_2") as span:
+                await asyncio.sleep(random.uniform(0.0, 1.0))
+                span.set_attributes(generate_attributes())
+
+        if random.uniform(0, 1) < 0.5:
+            agl.emit_reward(random.uniform(0.0, 1.0))
+
+    # Final Span
+    with tracer.start_as_current_span("final") as span:
+        await asyncio.sleep(random.uniform(0.0, 1.0))
+        span.set_attributes(generate_attributes())
+
+    agl.emit_reward(random.uniform(1.0, 2.0))
+
+
+def check_spans(spans: Sequence[agl.Span], task: str) -> None:
+    """Check if the spans contain the task."""
+    found_task = False
+    last_reward_in_12 = None
+    for span in spans:
+        if span.attributes.get("task") == task:
+            found_task = True
+        if span.name == agl.SpanNames.REWARD.value:
+            if span.attributes.get("reward") is None:
+                raise ValueError("Reward is not set for a reward span")
+            rew = float(span.attributes.get("reward"))  # type: ignore
+            if rew >= 1 and rew <= 2:
+                last_reward_in_12 = True
+            else:
+                last_reward_in_12 = False
+    if not found_task:
+        raise ValueError(f"Task {task} is not found in the spans")
+    if last_reward_in_12 is None:
+        raise ValueError("Last reward is not found")
+    elif not last_reward_in_12:
+        raise ValueError("Last reward is not in the range of 1 to 2")
+
+
+class AlgorithmBatch(agl.Algorithm):
+    def __init__(
+        self,
+        mode: Literal["batch", "batch_partial", "single"],
+        total_tasks: int,
+        batch_size: Optional[int] = None,
+        remaining_tasks: Optional[int] = None,
+        concurrency: Optional[int] = None,
+    ):
+        self.mode = mode
+        self.total_tasks = total_tasks
+        self.batch_size = batch_size
+        self.remaining_tasks = remaining_tasks
+        self.concurrency = concurrency
+
+    async def run(
+        self, train_dataset: Optional[agl.Dataset[Any]] = None, val_dataset: Optional[agl.Dataset[Any]] = None
+    ):
+        if self.mode == "batch":
+            assert self.batch_size is not None
+            await self.algorithm_batch(self.total_tasks, self.batch_size)
+        elif self.mode == "batch_partial":
+            assert self.batch_size is not None
+            assert self.remaining_tasks is not None
+            await self.algorithm_batch_with_completion_threshold(
+                self.total_tasks, self.batch_size, self.remaining_tasks
+            )
+        elif self.mode == "single":
+            assert self.concurrency is not None
+            await self.algorithm_batch_single(self.total_tasks, self.concurrency)
+        else:
+            raise ValueError(f"Invalid mode: {self.mode}")
+
+    async def algorithm_batch(self, total_tasks: int, batch_size: int):
+        """
+        At each time, the algorithm will enqueue a batch of rollouts of size `batch_size`.
+        The algorithm will use wait_for_rollouts to wait for all rollouts to complete.
+        It then checks whether all rollouts are successful and check the spans to ensure the task is found
+        and the last reward is in the range of 1 to 2.
+        After that, the algorithm will enqueue a new batch of new tasks, until the total number of tasks is reached.
+        """
+        store = self.get_store()
+        submitted = 0
+
+        while submitted < total_tasks:
+            batch_count = min(batch_size, total_tasks - submitted)
+            batch_rollouts: List[Tuple[str, str]] = []
+            await store.add_resources(
+                {
+                    "llm": agl.LLM(
+                        endpoint=f"http://localhost:{submitted}/v1",
+                        model=f"test-model-{submitted}",
+                    )
+                }
+            )
+            for _ in range(batch_count):
+                task_name = f"task-{submitted}-generated"
+                rollout = await store.enqueue_rollout(input=task_name, mode="train")
+                batch_rollouts.append((rollout.rollout_id, task_name))
+                submitted += 1
+
+            pending = {rollout_id: task_name for rollout_id, task_name in batch_rollouts}
+            completed_ids: Set[str] = set()
+            while len(completed_ids) < len(batch_rollouts):
+                finished_rollouts = await store.wait_for_rollouts(
+                    rollout_ids=[rollout_id for rollout_id, _ in batch_rollouts],
+                    timeout=0.0,
+                )
+                for rollout in finished_rollouts:
+                    rollout_id = rollout.rollout_id
+                    if rollout_id in completed_ids:
+                        continue
+                    if rollout.status != "succeeded":
+                        raise RuntimeError(f"Rollout {rollout_id} finished with status {rollout.status}")
+                    spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+                    check_spans(spans, pending[rollout_id])
+                    completed_ids.add(rollout_id)
+                await asyncio.sleep(5.0)
+
+    async def algorithm_batch_with_completion_threshold(self, total_tasks: int, batch_size: int, remaining_tasks: int):
+        """Different from `algorithm_batch`, this algorithm will use query_rollouts to get rollouts' status.
+        It will enqueue a new batch of new tasks when the number of running rollouts is less than the remaining tasks threshold.
+        """
+        store = self.get_store()
+        submitted = 0
+        completed = 0
+        active_rollouts: Dict[str, str] = {}
+
+        while completed < total_tasks:
+            if submitted < total_tasks and len(active_rollouts) < remaining_tasks:
+                batch_count = min(batch_size, total_tasks - submitted)
+                await store.add_resources(
+                    {
+                        "llm": agl.LLM(
+                            endpoint=f"http://localhost:{submitted}/v1",
+                            model=f"test-model-{submitted}",
+                        )
+                    }
+                )
+                for _ in range(batch_count):
+                    task_name = f"task-{submitted}"
+                    rollout = await store.enqueue_rollout(input=task_name, mode="train")
+                    active_rollouts[rollout.rollout_id] = task_name
+                    submitted += 1
+                continue
+
+            if not active_rollouts:
+                await asyncio.sleep(0.01)
+                continue
+
+            rollouts = await store.query_rollouts(rollout_id_in=list(active_rollouts.keys()))
+            newly_completed = 0
+            for rollout in rollouts:
+                rollout_id = rollout.rollout_id
+                if rollout_id not in active_rollouts:
+                    continue
+                if rollout.status in ("queuing", "preparing", "running", "requeuing"):
+                    continue
+                if rollout.status != "succeeded":
+                    raise RuntimeError(f"Rollout {rollout_id} finished with status {rollout.status}")
+                spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+                check_spans(spans, active_rollouts.pop(rollout_id))
+                completed += 1
+                newly_completed += 1
+
+            if newly_completed == 0:
+                await asyncio.sleep(5.0)
+
+    async def algorithm_batch_single(self, total_tasks: int, concurrency: int):
+        """Different from `algorithm_batch`, this algorithm will use one async function to enqueue one rollout at a time.
+        The function only cares about the rollout it's currently processing.
+        It waits for the rollouts with `get_rollout_by_id` and check the spans to ensure the rollout is successful.
+        The concurrency is managed via a asyncio semaphore.
+        """
+        store = self.get_store()
+        semaphore = asyncio.Semaphore(concurrency)
+
+        async def handle_single(task_index: int) -> None:
+            task_name = f"task-{task_index}"
+            async with semaphore:
+                await store.add_resources(
+                    {
+                        "llm": agl.LLM(
+                            endpoint=f"http://localhost:{task_index}/v1",
+                            model=f"test-model-{task_index}",
+                        )
+                    }
+                )
+                rollout = await store.enqueue_rollout(input=task_name, mode="train")
+                rollout_id = rollout.rollout_id
+                while True:
+                    current = await store.get_rollout_by_id(rollout_id)
+                    if current is not None and current.status in ("failed", "succeeded", "cancelled"):
+                        if current.status != "succeeded":
+                            raise RuntimeError(f"Rollout {rollout_id} finished with status {current.status}")
+                        break
+                    await asyncio.sleep(5.0)
+                spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+                check_spans(spans, task_name)
+
+        all_tasks = [handle_single(i) for i in range(total_tasks)]
+        await asyncio.gather(*all_tasks)
+
+
+def main() -> None:
+    store = agl.LightningStoreClient("http://localhost:4747")
+    try:
+        trainer = agl.Trainer(
+            store=store,
+            algorithm=AlgorithmBatch(mode="batch", total_tasks=1024, batch_size=128),
+            n_runners=32,
+            strategy={
+                "type": "cs",
+                "managed_store": False,
+            },
+        )
+        trainer.fit(agent)
+    finally:
+        asyncio.run(store.close())
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Generating random test data for benchmarking."""
+
+import random
+import string
+from typing import Any, Callable, Dict, Optional, Tuple, Union, cast
+
+
+def random_string(length: int, *, alphabet: Optional[str] = None) -> str:
+    """
+    Generate a random string of fixed length.
+
+    Args:
+        length: Length of the generated string.
+        alphabet: Optional character set to draw from. If None, uses [A-Za-z0-9].
+    """
+    if length < 0:
+        raise ValueError("String length cannot be negative.")
+
+    alphabet = alphabet or (string.ascii_letters + string.digits)
+    return "".join(random.choices(alphabet, k=length))
+
+
+def _resolve_param(value: Union[int, Tuple[int, int]], name: str) -> int:
+    """
+    Convert parameter into a concrete integer.
+    If value is an int, return it.
+    If value is a tuple, interpret it as (low, high) and sample uniformly.
+    """
+    if isinstance(value, int):
+        if value < 0:
+            raise ValueError(f"{name} cannot be negative.")
+        return value
+
+    if (
+        isinstance(value, tuple)  # type: ignore
+        and len(value) == 2
+        and isinstance(value[0], int)  # type: ignore
+        and isinstance(value[1], int)  # type: ignore
+    ):
+        low, high = value
+        if low < 0 or high < 0:
+            raise ValueError(f"{name} range cannot contain negative values.")
+        if low > high:
+            raise ValueError(f"{name} tuple must be (low, high) with low <= high.")
+        return random.randint(low, high)
+
+    raise TypeError(f"{name} must be an int or a 2-tuple of ints.")
+
+
+def default_value_factory(length: int) -> str:
+    """Default value factory for generating string payloads."""
+    return random_string(length)
+
+
+def random_dict(
+    *,
+    depth: Union[int, Tuple[int, int]],
+    breadth: Union[int, Tuple[int, int]],
+    key_length: Union[int, Tuple[int, int]],
+    value_length: Union[int, Tuple[int, int]],
+    value_factory: Optional[Callable[[int], Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Generate a nested dictionary with configurable depth, breadth, and
+    value sizes. Integer or (low, high) tuples are supported for
+    all structural parameters.
+
+    Args:
+        depth: Number of nested levels or a tuple specifying a range.
+        breadth: Number of keys per level or a tuple range.
+        key_length: Length of each key or a tuple range.
+        value_length: Length of each value or a tuple range.
+        value_factory: Function mapping `value_length` → value.
+
+    Returns:
+        A nested dictionary of arbitrary size.
+    """
+    # Default factory
+    if value_factory is None:
+        value_factory = random_string
+
+    def build(level: int) -> Dict[str, Any]:
+        # For each level, breadth/key/value lengths may vary, so draw fresh each time
+        current_breadth = _resolve_param(breadth, "breadth")
+
+        if current_breadth < 0:
+            raise ValueError("Breadth cannot be negative.")
+
+        target_depth = depth if isinstance(depth, int) else _resolve_param((level, depth[1]), "depth")
+
+        if level == target_depth:
+            # leaf nodes
+            return {
+                random_string(_resolve_param(key_length, "key_length")): value_factory(
+                    _resolve_param(value_length, "value_length")
+                )
+                for _ in range(current_breadth)
+            }
+
+        # nested nodes
+        return {
+            random_string(_resolve_param(key_length, "key_length")): build(level + 1) for _ in range(current_breadth)
+        }
+
+    return build(1)
+
+
+def flatten_dict(d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]:
+    """Flatten a nested dictionary into a single level dictionary. Keys are joined by dots."""
+
+    result: Dict[str, Any] = {}
+    for key, value in d.items():
+        if isinstance(value, dict):
+            result.update(flatten_dict(cast(Dict[str, Any], value), f"{prefix}.{key}" if prefix else key))
+        else:
+            result[f"{prefix}.{key}" if prefix else key] = value
+    return result
+
+
+if __name__ == "__main__":
+    # Example usage
+    import json
+
+    structured_dict = random_dict(
+        depth=(1, 3),
+        breadth=(2, 6),
+        key_length=(3, 20),
+        value_length=(5, 300),
+    )
+
+    print(json.dumps(flatten_dict(structured_dict), indent=2))
@@ -1,41 +1,72 @@
 # Copyright (c) Microsoft. All rights reserved.

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

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

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

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


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


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


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


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

    return span
+
+
+class SampleItem(BaseModel):
+    partition: str
+    index: int
+    name: str
+    status: str
+    tags: List[str] = Field(default_factory=list)
+    score: float | None = None
+    rank: int | None = None
+    updated_time: float | None = None
+    payload: Dict[str, int] = Field(default_factory=dict)
+    metadata: str | None = None
+
+
+class QueueItem(BaseModel):
+    idx: int
+
+
+@pytest_asyncio.fixture
+async def mongo_client():
+    from pymongo import AsyncMongoClient
+
+    client = AsyncMongoClient[Any](mongo_uri, serverSelectionTimeoutMS=5000)
+    try:
+        await client.admin.command("ping")
+    except Exception as exc:  # depends on external service
+        await client.close()
+        raise RuntimeError(f"MongoDB not available: {exc}")
+
+    try:
+        yield client
+    finally:
+        await client.close()
+
+
+@pytest_asyncio.fixture
+async def temporary_mongo_database(mongo_client: AsyncMongoClient[Any]):
+    """Yield a temporary MongoDB database for integration tests."""
+    db_name = f"agentlightning-test-{uuid4().hex}"
+    db = mongo_client[db_name]  # type: ignore
+    try:
+        yield db
+    finally:
+        await mongo_client.drop_database(db_name)
+
+
+### Collection fixtures ###
+
+
+@pytest.fixture()
+def sample_items() -> List[SampleItem]:
+    return [
+        SampleItem(
+            partition="alpha",
+            index=1,
+            name="urgent-phase-one",
+            status="new",
+            tags=["core", "urgent"],
+            score=10.5,
+            rank=3,
+            updated_time=12.0,
+            payload={"priority": 10},
+            metadata="alpha-start",
+        ),
+        SampleItem(
+            partition="alpha",
+            index=2,
+            name="phase-two",
+            status="running",
+            tags=["core"],
+            score=5.0,
+            rank=2,
+            updated_time=None,
+            payload={"priority": 5},
+            metadata=None,
+        ),
+        SampleItem(
+            partition="alpha",
+            index=3,
+            name="delayed-phase",
+            status="blocked",
+            tags=["delayed"],
+            score=None,
+            rank=5,
+            updated_time=15.1,
+            payload={"priority": 8},
+            metadata="delayed-phase",
+        ),
+        SampleItem(
+            partition="beta",
+            index=1,
+            name="beta-critical",
+            status="new",
+            tags=["beta", "urgent"],
+            score=8.0,
+            rank=1,
+            updated_time=7.0,
+            payload={"priority": 7},
+            metadata="beta critical",
+        ),
+        SampleItem(
+            partition="beta",
+            index=2,
+            name="beta optional",
+            status="done",
+            tags=["beta"],
+            score=3.0,
+            rank=None,
+            updated_time=2.0,
+            payload={"priority": 1},
+            metadata="optional path",
+        ),
+        SampleItem(
+            partition="gamma",
+            index=1,
+            name="gamma-phase",
+            status="running",
+            tags=[],
+            score=9.5,
+            rank=4,
+            updated_time=None,
+            payload={"priority": 9},
+            metadata="gamma-phase data",
+        ),
+        SampleItem(
+            partition="gamma",
+            index=2,
+            name="gamma-late",
+            status="done",
+            tags=["late", "core"],
+            score=1.0,
+            rank=6,
+            updated_time=20.0,
+            payload={"priority": 2},
+            metadata="gamma late entry",
+        ),
+        SampleItem(
+            partition="delta",
+            index=1,
+            name="delta misc",
+            status="archived",
+            tags=["misc"],
+            score=4.2,
+            rank=7,
+            updated_time=11.0,
+            payload={"priority": 3},
+            metadata="delta misc block",
+        ),
+    ]
+
+
+### Generic collection fixtures ###
+
+
+@pytest.fixture()
+def sample_collection_memory(sample_items: Sequence[SampleItem]) -> ListBasedCollection[SampleItem]:
+    collection: Collection[SampleItem] = ListBasedCollection(list(sample_items), SampleItem, ("partition", "index"))
+    setattr(collection, "_test_backend", "memory")
+    return collection
+
+
+@pytest_asyncio.fixture
+async def sample_collection_mongo(temporary_mongo_database: AsyncDatabase[Any], sample_items: Sequence[SampleItem]):
+    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        collection = MongoBasedCollection(
+            client_pool,
+            temporary_mongo_database.name,
+            "sample-items",
+            "partition-123",
+            ["partition", "index"],
+            SampleItem,
+        )
+        await collection.insert(sample_items)
+        setattr(collection, "_test_backend", "mongo")
+        yield collection
+
+
+@pytest.fixture(
+    params=[
+        "memory",
+        pytest.param("mongo", marks=pytest.mark.mongo),
+    ]
+)
+def sample_collection(request: pytest.FixtureRequest):
+    backend = request.param
+    return request.getfixturevalue("sample_collection_" + backend)
+
+
+### Generic queue fixtures ###
+
+
+@pytest.fixture
+def deque_queue_memory() -> DequeBasedQueue[QueueItem]:
+    return DequeBasedQueue(QueueItem, [QueueItem(idx=i) for i in range(3)])
+
+
+@pytest_asyncio.fixture
+async def deque_queue_mongo(temporary_mongo_database: AsyncDatabase[Any]):
+    from agentlightning.store.collection.mongo import MongoBasedQueue, MongoClientPool
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        queue = MongoBasedQueue[QueueItem](
+            client_pool,
+            temporary_mongo_database.name,
+            "queue-items",
+            "partition-1",
+            QueueItem,
+        )
+        await queue.enqueue([QueueItem(idx=i) for i in range(3)])
+        yield queue
+
+
+@pytest.fixture(
+    params=[
+        "memory",
+        pytest.param("mongo", marks=pytest.mark.mongo),
+    ]
+)
+def deque_queue(request: pytest.FixtureRequest) -> AsyncGenerator[Queue[QueueItem], None]:
+    backend = request.param
+    return request.getfixturevalue("deque_queue_" + backend)
+
+
+### Generic key-value fixtures ###
+
+
+@pytest.fixture()
+def dict_key_value_data() -> Dict[str, int]:
+    return {"alpha": 1, "beta": 2}
+
+
+@pytest.fixture()
+def dict_key_value_memory(dict_key_value_data: Dict[str, int]) -> DictBasedKeyValue[str, int]:
+    return DictBasedKeyValue(dict_key_value_data)
+
+
+@pytest_asyncio.fixture
+async def dict_key_value_mongo(temporary_mongo_database: AsyncDatabase[Any], dict_key_value_data: Dict[str, int]):
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            "key-value-items",
+            "partition-1",
+            str,
+            int,
+        )
+        for key, value in dict_key_value_data.items():
+            await key_value.set(key, value)
+        yield key_value
+
+
+@pytest.fixture(
+    params=[
+        "memory",
+        pytest.param("mongo", marks=pytest.mark.mongo),
+    ]
+)
+def dict_key_value(request: pytest.FixtureRequest) -> AsyncGenerator[KeyValue[str, int], None]:
+    backend = request.param
+    return request.getfixturevalue("dict_key_value_" + backend)
@@ -2,135 +2,29 @@

 from __future__ import annotations

-from typing import Dict, Iterable, List, Literal, Mapping, Sequence, Tuple
+import asyncio
+import time
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Literal, Mapping, Sequence, Tuple, Union
+from uuid import uuid4

 import pytest
-from pydantic import BaseModel, Field
+from pydantic import BaseModel

 import agentlightning.store.collection.memory as memory_module
 from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, ListBasedCollection
+from agentlightning.store.collection.base import Collection
 from agentlightning.store.collection.memory import _item_matches_filters  # pyright: ignore[reportPrivateUsage]
+from agentlightning.types import Rollout
+from tests.store.conftest import QueueItem, SampleItem

-
-class SampleItem(BaseModel):
-    partition: str
-    index: int
-    name: str
-    status: str
-    tags: List[str] = Field(default_factory=list)
-    score: float | None = None
-    rank: int | None = None
-    updated_time: float | None = None
-    payload: Dict[str, int] = Field(default_factory=dict)
-    metadata: str | None = None
+if TYPE_CHECKING:
+    from pymongo.asynchronous.database import AsyncDatabase


 def _build_collection(items: Iterable[SampleItem] = ()) -> ListBasedCollection[SampleItem]:
    return ListBasedCollection(list(items), SampleItem, ("partition", "index"))


-@pytest.fixture()
-def sample_items() -> List[SampleItem]:
-    return [
-        SampleItem(
-            partition="alpha",
-            index=1,
-            name="urgent-phase-one",
-            status="new",
-            tags=["core", "urgent"],
-            score=10.5,
-            rank=3,
-            updated_time=12.0,
-            payload={"priority": 10},
-            metadata="alpha-start",
-        ),
-        SampleItem(
-            partition="alpha",
-            index=2,
-            name="phase-two",
-            status="running",
-            tags=["core"],
-            score=5.0,
-            rank=2,
-            updated_time=None,
-            payload={"priority": 5},
-            metadata=None,
-        ),
-        SampleItem(
-            partition="alpha",
-            index=3,
-            name="delayed-phase",
-            status="blocked",
-            tags=["delayed"],
-            score=None,
-            rank=5,
-            updated_time=15.1,
-            payload={"priority": 8},
-            metadata="delayed-phase",
-        ),
-        SampleItem(
-            partition="beta",
-            index=1,
-            name="beta-critical",
-            status="new",
-            tags=["beta", "urgent"],
-            score=8.0,
-            rank=1,
-            updated_time=7.0,
-            payload={"priority": 7},
-            metadata="beta critical",
-        ),
-        SampleItem(
-            partition="beta",
-            index=2,
-            name="beta optional",
-            status="done",
-            tags=["beta"],
-            score=3.0,
-            rank=None,
-            updated_time=2.0,
-            payload={"priority": 1},
-            metadata="optional path",
-        ),
-        SampleItem(
-            partition="gamma",
-            index=1,
-            name="gamma-phase",
-            status="running",
-            tags=[],
-            score=9.5,
-            rank=4,
-            updated_time=None,
-            payload={"priority": 9},
-            metadata="gamma-phase data",
-        ),
-        SampleItem(
-            partition="gamma",
-            index=2,
-            name="gamma-late",
-            status="done",
-            tags=["late", "core"],
-            score=1.0,
-            rank=6,
-            updated_time=20.0,
-            payload={"priority": 2},
-            metadata="gamma late entry",
-        ),
-        SampleItem(
-            partition="delta",
-            index=1,
-            name="delta misc",
-            status="archived",
-            tags=["misc"],
-            score=4.2,
-            rank=7,
-            updated_time=11.0,
-            payload={"priority": 3},
-            metadata="delta misc block",
-        ),
-    ]
-
-
 BASE_KEY_ORDER: List[Tuple[str, int]] = [
    ("alpha", 1),
    ("alpha", 2),
@@ -143,11 +37,6 @@ BASE_KEY_ORDER: List[Tuple[str, int]] = [
 ]


-@pytest.fixture()
-def sample_collection(sample_items: Sequence[SampleItem]) -> ListBasedCollection[SampleItem]:
-    return _build_collection(sample_items)
-
-
 def _key_pairs(items: Sequence[SampleItem]) -> List[Tuple[str, int]]:
    return [(item.partition, item.index) for item in items]

@@ -161,43 +50,50 @@ def test_list_collection_requires_primary_keys(sample_items: Sequence[SampleItem
        ListBasedCollection(list(sample_items), SampleItem, ())


-def test_list_collection_primary_keys(sample_collection: ListBasedCollection[SampleItem]) -> None:
+@pytest.mark.asyncio()
+async def test_list_collection_primary_keys(sample_collection: Collection[SampleItem]) -> None:
    assert tuple(sample_collection.primary_keys()) == ("partition", "index")


-def test_list_collection_item_type(sample_collection: ListBasedCollection[SampleItem]) -> None:
+@pytest.mark.asyncio()
+async def test_list_collection_item_type(sample_collection: Collection[SampleItem]) -> None:
    assert sample_collection.item_type() is SampleItem


-def test_list_collection_initial_size(
-    sample_collection: ListBasedCollection[SampleItem], sample_items: Sequence[SampleItem]
+@pytest.mark.asyncio()
+async def test_list_collection_initial_size(
+    sample_collection: Collection[SampleItem], sample_items: Sequence[SampleItem]
 ) -> None:
-    assert sample_collection.size() == len(sample_items)
-
-
-def test_list_collection_repr_contains_model_info(sample_collection: ListBasedCollection[SampleItem]) -> None:
-    result = repr(sample_collection)
-    assert "ListBasedCollection" in result and "SampleItem" in result and str(sample_collection.size()) in result
+    assert (await sample_collection.size()) == len(sample_items)


@pytest.mark.asyncio()
-async def test_list_collection_insert_adds_item(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_repr_contains_model_info(sample_collection: Collection[SampleItem]) -> None:
+    result = repr(sample_collection)
+    assert sample_collection.__class__.__name__ in result
+    assert "SampleItem" in result
+    if isinstance(sample_collection, ListBasedCollection):
+        assert str(await sample_collection.size()) in result
+
+
+@pytest.mark.asyncio()
+async def test_list_collection_insert_adds_item(sample_collection: Collection[SampleItem]) -> None:
    new_item = SampleItem(partition="omega", index=1, name="omega", status="new")
    await sample_collection.insert([new_item])
-    assert sample_collection.size() == 9
+    assert (await sample_collection.size()) == 9
    result = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 1}})
    assert result == new_item


@pytest.mark.asyncio()
-async def test_list_collection_insert_duplicate_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_insert_duplicate_raises(sample_collection: Collection[SampleItem]) -> None:
    duplicate = SampleItem(partition="alpha", index=1, name="dup", status="new")
    with pytest.raises(ValueError):
        await sample_collection.insert([duplicate])


@pytest.mark.asyncio()
-async def test_list_collection_insert_wrong_type(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_insert_wrong_type(sample_collection: Collection[SampleItem]) -> None:
    class Another(BaseModel):
        partition: str
        index: int
@@ -208,7 +104,7 @@ async def test_list_collection_insert_wrong_type(sample_collection: ListBasedCol


@pytest.mark.asyncio()
-async def test_list_collection_update_existing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_update_existing(sample_collection: Collection[SampleItem]) -> None:
    updated = SampleItem(partition="alpha", index=1, name="updated", status="new")
    await sample_collection.update([updated])
    result = await sample_collection.get({"partition": {"exact": "alpha"}, "index": {"exact": 1}})
@@ -216,74 +112,74 @@ async def test_list_collection_update_existing(sample_collection: ListBasedColle


@pytest.mark.asyncio()
-async def test_list_collection_update_missing_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_update_missing_raises(sample_collection: Collection[SampleItem]) -> None:
    missing = SampleItem(partition="omega", index=99, name="missing", status="lost")
    with pytest.raises(ValueError):
        await sample_collection.update([missing])


@pytest.mark.asyncio()
-async def test_list_collection_delete_existing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_delete_existing(sample_collection: Collection[SampleItem]) -> None:
    target = SampleItem(partition="alpha", index=1, name="ignored", status="new")
    await sample_collection.delete([target])
-    assert sample_collection.size() == 7
+    assert (await sample_collection.size()) == 7
    result = await sample_collection.get({"partition": {"exact": "alpha"}, "index": {"exact": 1}})
    assert result is None


@pytest.mark.asyncio()
-async def test_list_collection_delete_missing_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_delete_missing_raises(sample_collection: Collection[SampleItem]) -> None:
    missing = SampleItem(partition="omega", index=3, name="x", status="y")
    with pytest.raises(ValueError):
        await sample_collection.delete([missing])


@pytest.mark.asyncio()
-async def test_list_collection_upsert_inserts_when_missing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_upsert_inserts_when_missing(sample_collection: Collection[SampleItem]) -> None:
    created = SampleItem(partition="omega", index=4, name="new", status="queued")
    await sample_collection.upsert([created])
-    assert sample_collection.size() == 9
+    assert (await sample_collection.size()) == 9
    fetched = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 4}})
    assert fetched == created


@pytest.mark.asyncio()
-async def test_list_collection_upsert_updates_when_existing(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_upsert_updates_when_existing(sample_collection: Collection[SampleItem]) -> None:
    replacement = SampleItem(partition="beta", index=2, name="replacement", status="done")
    await sample_collection.upsert([replacement])
-    assert sample_collection.size() == 8
+    assert (await sample_collection.size()) == 8
    fetched = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
    assert fetched == replacement


@pytest.mark.asyncio()
-async def test_list_collection_delete_multiple_items(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_delete_multiple_items(sample_collection: Collection[SampleItem]) -> None:
    await sample_collection.delete(
        [
            SampleItem(partition="alpha", index=1, name="", status=""),
            SampleItem(partition="beta", index=1, name="", status=""),
        ]
    )
-    assert sample_collection.size() == 6
+    assert (await sample_collection.size()) == 6


@pytest.mark.asyncio()
 async def test_list_collection_insert_accepts_tuple_sequence(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
 ) -> None:
    extra = (
        SampleItem(partition="tuple", index=1, name="a", status="pending"),
        SampleItem(partition="tuple", index=2, name="b", status="pending"),
    )
    await sample_collection.insert(extra)
-    assert sample_collection.size() == 10
+    assert (await sample_collection.size()) == 10
    fetched = await sample_collection.query(filter={"partition": {"exact": "tuple"}})
    assert _sorted_pairs(fetched.items) == [("tuple", 1), ("tuple", 2)]


@pytest.mark.asyncio()
 async def test_list_collection_query_without_filters_returns_all(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
 ) -> None:
    result = await sample_collection.query()
    assert result.total == 8
@@ -331,10 +227,17 @@ async def test_list_collection_query_without_filters_returns_all(
    ],
 )
 async def test_list_collection_query_filters(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
    filters: Dict[str, Dict[str, object]],
    expected: Sequence[Tuple[str, int]],
+    request: pytest.FixtureRequest,
 ) -> None:
+    # Mongo implementation raises ValueError for non-iterable values in within filter
+    if request.node.callspec.id == "mongo-within-non-iterable":  # type: ignore
+        with pytest.raises(ValueError):
+            await sample_collection.query(filter=filters)  # type: ignore[arg-type]
+        return
+
    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == sorted(expected)
    assert result.total == len(expected)
@@ -367,7 +270,7 @@ async def test_list_collection_query_filters(
    ],
 )
 async def test_list_collection_filter_logic(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
    filters: Dict[str, Dict[str, object]],
    filter_logic: Literal["and", "or"],
    expected: Sequence[Tuple[str, int]],
@@ -380,7 +283,7 @@ async def test_list_collection_filter_logic(

@pytest.mark.asyncio()
 async def test_list_collection_must_filters_respected_with_or(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
 ) -> None:
    filters = {
        "_aggregate": "or",
@@ -394,7 +297,7 @@ async def test_list_collection_must_filters_respected_with_or(

@pytest.mark.asyncio()
 async def test_list_collection_must_filters_accept_sequence(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
 ) -> None:
    filters = {
        "_aggregate": "or",
@@ -411,9 +314,11 @@ async def test_list_collection_must_filters_accept_sequence(

@pytest.mark.asyncio()
 async def test_list_collection_must_filters_limit_tree_scan_even_with_or(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    if not isinstance(sample_collection, ListBasedCollection):
+        pytest.skip("This test is only valid for pure-memory collections")
    seen: List[Tuple[str, int]] = []
    original = _item_matches_filters

@@ -467,9 +372,11 @@ async def test_list_collection_primary_key_prefix_limits_filter_checks(

@pytest.mark.asyncio()
 async def test_list_collection_full_primary_key_avoids_tree_scan(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    if not isinstance(sample_collection, ListBasedCollection):
+        pytest.skip("This test is only valid for pure-memory collections")
    call_count = 0
    original_iter_items = (  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType,reportUnknownVariableType]
        ListBasedCollection._iter_items  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType]
@@ -504,23 +411,42 @@ async def test_list_collection_full_primary_key_avoids_tree_scan(
        ("rank", "desc", 4, [("delta", 1), ("gamma", 2), ("alpha", 3), ("gamma", 1)]),
        ("score", "asc", 4, [("alpha", 3), ("gamma", 2), ("beta", 2), ("delta", 1)]),
        ("score", "desc", 4, [("alpha", 1), ("gamma", 1), ("beta", 1), ("alpha", 2)]),
-        ("updated_time", "asc", 4, [("beta", 2), ("beta", 1), ("delta", 1), ("alpha", 1)]),
-        ("updated_time", "desc", 4, [("gamma", 1), ("alpha", 2), ("gamma", 2), ("alpha", 3)]),
+        (
+            "updated_time",
+            "asc",
+            4,
+            (
+                [("alpha", 2), ("gamma", 1), ("beta", 2), ("beta", 1)],
+                [("beta", 2), ("beta", 1), ("delta", 1), ("alpha", 1)],
+            ),
+        ),
+        (
+            "updated_time",
+            "desc",
+            4,
+            (
+                [("gamma", 2), ("alpha", 3), ("alpha", 1), ("delta", 1)],
+                [("gamma", 1), ("alpha", 2), ("gamma", 2), ("alpha", 3)],
+            ),
+        ),
    ],
 )
 async def test_list_collection_sorting(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
    sort_by: str,
    sort_order: str,
    limit: int,
-    expected: Sequence[Tuple[str, int]],
+    expected: Union[Sequence[Tuple[str, int]], Tuple[Sequence[Tuple[str, int]], ...]],
 ) -> None:
    result = await sample_collection.query(sort={"name": sort_by, "order": sort_order}, limit=limit)  # type: ignore[arg-type]
-    assert _key_pairs(result.items) == list(expected)
+    if isinstance(expected, tuple):
+        assert any(_key_pairs(result.items) == list(expected) for expected in expected)
+    else:
+        assert _key_pairs(result.items) == list(expected)


@pytest.mark.asyncio()
-async def test_list_collection_sort_by_missing_field_raises(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_sort_by_missing_field_raises(sample_collection: Collection[SampleItem]) -> None:
    with pytest.raises(ValueError):
        await sample_collection.query(sort={"name": "does_not_exist", "order": "asc"})

@@ -538,7 +464,7 @@ async def test_list_collection_sort_by_missing_field_raises(sample_collection: L
    ],
 )
 async def test_list_collection_pagination_without_sort(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
    limit: int,
    offset: int,
    expected: Sequence[Tuple[str, int]],
@@ -550,21 +476,21 @@ async def test_list_collection_pagination_without_sort(


@pytest.mark.asyncio()
-async def test_list_collection_pagination_with_sort(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_pagination_with_sort(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.query(sort={"name": "name", "order": "asc"}, limit=2, offset=3)
    assert _key_pairs(result.items) == [("delta", 1), ("gamma", 2)]
    assert result.total == 8


@pytest.mark.asyncio()
-async def test_list_collection_limit_unbounded_with_sort(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_limit_unbounded_with_sort(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.query(sort={"name": "name", "order": "asc"}, limit=-1, offset=6)
    assert _key_pairs(result.items) == [("alpha", 2), ("alpha", 1)]
    assert result.total == 8


@pytest.mark.asyncio()
-async def test_list_collection_limit_zero_reports_total(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_limit_zero_reports_total(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.query(filter={"status": {"exact": "done"}}, limit=0)
    assert result.items == []
    assert result.total == 2
@@ -572,7 +498,7 @@ async def test_list_collection_limit_zero_reports_total(sample_collection: ListB

@pytest.mark.asyncio()
 async def test_list_collection_offset_beyond_total_returns_empty(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
 ) -> None:
    result = await sample_collection.query(filter={"status": {"exact": "done"}}, offset=10)
    assert result.items == []
@@ -581,7 +507,7 @@ async def test_list_collection_offset_beyond_total_returns_empty(

@pytest.mark.asyncio()
 async def test_list_collection_query_reports_total_with_limit(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
 ) -> None:
    result = await sample_collection.query(filter={"partition": {"exact": "alpha"}}, limit=1)
    assert result.total == 3
@@ -589,28 +515,28 @@ async def test_list_collection_query_reports_total_with_limit(


@pytest.mark.asyncio()
-async def test_list_collection_get_returns_first_match(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_get_returns_first_match(sample_collection: Collection[SampleItem]) -> None:
    item = await sample_collection.get({"status": {"exact": "new"}})
    assert item is not None
-    assert (item.partition, item.index) == ("beta", 1)
+    assert (item.partition, item.index) in [("beta", 1), ("alpha", 1)]


@pytest.mark.asyncio()
-async def test_list_collection_get_returns_none(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_get_returns_none(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.get({"partition": {"exact": "missing"}})
    assert result is None


@pytest.mark.asyncio()
-async def test_list_collection_get_respects_filter_logic(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_get_respects_filter_logic(sample_collection: Collection[SampleItem]) -> None:
    filters = {"status": {"exact": "done"}, "tags": {"contains": "urgent"}, "_aggregate": "or"}
    item = await sample_collection.get(filters)  # type: ignore[arg-type]
    assert item is not None
-    assert (item.partition, item.index) == ("gamma", 2)
+    assert (item.partition, item.index) in [("gamma", 2), ("alpha", 1)]


@pytest.mark.asyncio()
-async def test_list_collection_get_honors_sort_by(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_get_honors_sort_by(sample_collection: Collection[SampleItem]) -> None:
    filters = {"partition": {"exact": "alpha"}}
    item = await sample_collection.get(filters, sort={"name": "rank", "order": "asc"})  # type: ignore[arg-type]
    assert item is not None
@@ -618,7 +544,7 @@ async def test_list_collection_get_honors_sort_by(sample_collection: ListBasedCo


@pytest.mark.asyncio()
-async def test_list_collection_get_honors_sort_order(sample_collection: ListBasedCollection[SampleItem]) -> None:
+async def test_list_collection_get_honors_sort_order(sample_collection: Collection[SampleItem]) -> None:
    filters = {"partition": {"exact": "alpha"}}
    item = await sample_collection.get(filters, sort={"name": "rank", "order": "desc"})  # type: ignore[arg-type]
    assert item is not None
@@ -652,14 +578,14 @@ async def test_list_collection_bulk_delete_and_size() -> None:
    items = [SampleItem(partition="bulk", index=i, name=f"item-{i}", status="bulk") for i in range(40)]
    collection = _build_collection(items)
    await collection.delete(items[:20])
-    assert collection.size() == 20
+    assert (await collection.size()) == 20
    await collection.delete(items[20:])
-    assert collection.size() == 0
+    assert (await collection.size()) == 0


@pytest.mark.asyncio()
 async def test_list_collection_query_rejects_unknown_operator(
-    sample_collection: ListBasedCollection[SampleItem],
+    sample_collection: Collection[SampleItem],
 ) -> None:
    with pytest.raises(ValueError):
        await sample_collection.query(filter={"status": {"invalid": "x"}})  # type: ignore[arg-type]
@@ -673,17 +599,9 @@ async def test_list_collection_query_result_type() -> None:
    assert result.offset == 0


-class QueueItem(BaseModel):
-    idx: int
-
-
-@pytest.fixture()
-def deque_queue() -> DequeBasedQueue[QueueItem]:
-    return DequeBasedQueue(QueueItem, [QueueItem(idx=i) for i in range(3)])
-
-
-def test_deque_queue_initial_size(deque_queue: DequeBasedQueue[QueueItem]) -> None:
-    assert deque_queue.size() == 3
+@pytest.mark.asyncio()
+async def test_deque_queue_initial_size(deque_queue: DequeBasedQueue[QueueItem]) -> None:
+    assert (await deque_queue.size()) == 3


 def test_deque_queue_item_type(deque_queue: DequeBasedQueue[QueueItem]) -> None:
@@ -701,7 +619,7 @@ async def test_deque_queue_enqueue_appends_items(deque_queue: DequeBasedQueue[Qu
    items = [QueueItem(idx=3), QueueItem(idx=4)]
    returned = await deque_queue.enqueue(items)
    assert returned == items
-    assert deque_queue.size() == 5
+    assert (await deque_queue.size()) == 5


@pytest.mark.asyncio()
@@ -718,7 +636,7 @@ async def test_deque_queue_enqueue_rejects_wrong_type(deque_queue: DequeBasedQue
 async def test_deque_queue_dequeue_respects_limit(deque_queue: DequeBasedQueue[QueueItem], limit: int) -> None:
    result = await deque_queue.dequeue(limit)
    assert len(result) == min(limit, 3)
-    assert deque_queue.size() == 3 - min(limit, 3)
+    assert (await deque_queue.size()) == 3 - min(limit, 3)


@pytest.mark.asyncio()
@@ -730,14 +648,14 @@ async def test_deque_queue_dequeue_zero_returns_empty(deque_queue: DequeBasedQue
 async def test_deque_queue_dequeue_more_than_available(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    result = await deque_queue.dequeue(10)
    assert len(result) == 3
-    assert deque_queue.size() == 0
+    assert (await deque_queue.size()) == 0


@pytest.mark.asyncio()
 async def test_deque_queue_peek_preserves_items(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    snapshot = await deque_queue.peek(2)
    assert [item.idx for item in snapshot] == [0, 1]
-    assert deque_queue.size() == 3
+    assert (await deque_queue.size()) == 3


@pytest.mark.asyncio()
@@ -757,25 +675,15 @@ async def test_deque_queue_handles_large_volume() -> None:
    queue = DequeBasedQueue(QueueItem)
    items = [QueueItem(idx=i) for i in range(2000)]
    await queue.enqueue(items)
-    assert queue.size() == 2000
+    assert (await queue.size()) == 2000
    drained = await queue.dequeue(1500)
    assert len(drained) == 1500
-    assert queue.size() == 500
-
-
-@pytest.fixture()
-def dict_key_value_data() -> Dict[str, int]:
-    return {"alpha": 1, "beta": 2}
-
-
-@pytest.fixture()
-def dict_key_value(dict_key_value_data: Dict[str, int]) -> DictBasedKeyValue[str, int]:
-    return DictBasedKeyValue(dict_key_value_data)
+    assert (await queue.size()) == 500


@pytest.mark.asyncio()
 async def test_dict_key_value_initial_state(dict_key_value: DictBasedKeyValue[str, int]) -> None:
-    assert dict_key_value.size() == 2
+    assert await dict_key_value.size() == 2
    assert await dict_key_value.get("alpha") == 1
    assert await dict_key_value.get("missing") is None

@@ -789,20 +697,20 @@ async def test_dict_key_value_has_handles_presence(dict_key_value: DictBasedKeyV
@pytest.mark.asyncio()
 async def test_dict_key_value_set_updates_and_expands(dict_key_value: DictBasedKeyValue[str, int]) -> None:
    await dict_key_value.set("gamma", 3)
-    assert dict_key_value.size() == 3
+    assert await dict_key_value.size() == 3
    await dict_key_value.set("alpha", 99)
    assert await dict_key_value.get("alpha") == 99
-    assert dict_key_value.size() == 3
+    assert await dict_key_value.size() == 3


@pytest.mark.asyncio()
 async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyValue[str, int]) -> None:
    result = await dict_key_value.pop("beta")
    assert result == 2
-    assert dict_key_value.size() == 1
+    assert await dict_key_value.size() == 1
    result = await dict_key_value.pop("missing", 42)
    assert result == 42
-    assert dict_key_value.size() == 1
+    assert await dict_key_value.size() == 1


@pytest.mark.asyncio()
@@ -811,3 +719,136 @@ async def test_dict_key_value_does_not_mutate_input_mapping(dict_key_value_data:
    await key_value.set("gamma", 3)  # type: ignore[arg-type]
    await key_value.pop("alpha")  # type: ignore[arg-type]
    assert dict_key_value_data == {"alpha": 1, "beta": 2}
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import (
+        MongoBasedCollection,
+        MongoBasedKeyValue,
+        MongoBasedQueue,
+        MongoClientPool,
+    )
+
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        collection = MongoBasedCollection[Any](
+            client_pool, temporary_mongo_database.name, "test", "test-123", ["rollout_id"], Rollout
+        )
+        await collection.ensure_collection()
+
+        start_time = time.time()
+        await collection.insert(
+            [Rollout(rollout_id="test-123", input="test-123", start_time=start_time, status="running")]
+        )
+
+        result = await collection.query(filter={"status": {"exact": "running"}})
+        assert result.items == [
+            Rollout(rollout_id="test-123", input="test-123", start_time=start_time, status="running")
+        ]
+
+        rollout_queue = MongoBasedQueue[str](
+            client_pool, temporary_mongo_database.name, "rollout_queue", "partition-1", str
+        )
+        await rollout_queue.ensure_collection()
+
+        await rollout_queue.enqueue(["r1", "r2", "r3"])
+        assert await rollout_queue.size() == 3
+        assert await rollout_queue.peek(2) == ["r1", "r2"]
+        assert await rollout_queue.dequeue(2) == ["r1", "r2"]
+        assert await rollout_queue.size() == 1
+
+        span_kv = MongoBasedKeyValue[str, int](
+            client_pool, temporary_mongo_database.name, "span_sequence_ids", "partition-1", str, int
+        )
+        await span_kv.ensure_collection()
+
+        await span_kv.set("span-123", 1)
+        assert await span_kv.has("span-123")
+        assert await span_kv.get("span-123") == 1
+        assert await span_kv.pop("span-123") == 1
+        assert not await span_kv.has("span-123")
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_ensure_collection_creates_partition_scoped_index(
+    temporary_mongo_database: AsyncDatabase[Any],
+) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
+
+    collection_name = f"ensure-{uuid4().hex}"
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        collection = MongoBasedCollection[Any](
+            client_pool,
+            temporary_mongo_database.name,
+            collection_name,
+            "partition-ensure",
+            ["name", "index"],
+            SampleItem,
+        )
+        await collection.ensure_collection()
+
+        unique_index = None
+        async for index in await temporary_mongo_database[collection_name].list_indexes():  # type: ignore
+            if index["name"] == "uniq_partition_name_index" and index.get("unique"):  # type: ignore
+                unique_index = index  # type: ignore
+                break
+
+        assert unique_index is not None, "expected unique partition/index key"
+        key_pairs = list(unique_index["key"].items())  # type: ignore
+        assert key_pairs == [("partition_id", 1), ("name", 1), ("index", 1)]
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_ensure_collection_survives_concurrent_calls(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
+
+    collection_name = f"ensure-{uuid4().hex}"
+
+    async def ensure_once() -> None:
+        async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+            collection = MongoBasedCollection(
+                client_pool,
+                temporary_mongo_database.name,
+                collection_name,
+                "partition-concurrent",
+                ["index"],
+                SampleItem,
+            )
+            await collection.ensure_collection()
+
+    await asyncio.gather(*(ensure_once() for _ in range(20)))
+
+    names = await temporary_mongo_database.list_collection_names()
+    assert names.count(collection_name) == 1
+
+    unique_indexes = []
+    async for index in await temporary_mongo_database[collection_name].list_indexes():  # type: ignore
+        if index["name"].startswith("uniq_partition_"):  # type: ignore
+            unique_indexes.append(index["name"])  # type: ignore
+    assert unique_indexes == ["uniq_partition_index"]
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_ensure_collection_repeats_without_altering_indexes(
+    temporary_mongo_database: AsyncDatabase[Any],
+) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
+
+    collection_name = f"ensure-{uuid4().hex}"
+    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        collection = MongoBasedCollection(
+            client_pool, temporary_mongo_database.name, collection_name, "partition-repeat", ["index"], SampleItem
+        )
+        await collection.ensure_collection()
+        await collection.ensure_collection()
+
+        unique_indexes = []
+        async for index in await temporary_mongo_database[collection_name].list_indexes():  # type: ignore
+            if index["name"].startswith("uniq_partition_"):  # type: ignore
+                unique_indexes.append((index["name"], list(index["key"].items())))  # type: ignore
+
+        assert unique_indexes == [("uniq_partition_index", [("partition_id", 1), ("index", 1)])]
@@ -882,7 +882,7 @@ async def test_query_resources_returns_history(store_fixture: LightningStore) ->
    )

    history = await store_fixture.query_resources()
-    assert [item.resources_id for item in history] == [first.resources_id, second.resources_id]
+    assert set([item.resources_id for item in history]) == {first.resources_id, second.resources_id}
    assert isinstance(history[0], ResourcesUpdate)
    assert isinstance(history[1], ResourcesUpdate)

@@ -1197,7 +1197,7 @@ async def test_duplicate_span_id_error(
    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)

    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-    assert "Duplicate span added" in caplog.text
+    assert "Duplicated span added" in caplog.text


@pytest.mark.asyncio
@@ -1859,7 +1859,8 @@ async def test_wait_with_timeout_none_polling(store_fixture: LightningStore) ->
    await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded")

    # The wait should complete now
-    completed = await asyncio.wait_for(wait_task, timeout=1.0)
+    timeout = 1.0 if isinstance(store_fixture, InMemoryLightningStore) else 11.0
+    completed = await asyncio.wait_for(wait_task, timeout=timeout)
    assert len(completed) == 1
    assert completed[0].rollout_id == rollout.rollout_id
    assert completed[0].status == "succeeded"
@@ -2013,8 +2014,12 @@ async def test_wait_nonexistent_rollout_with_finite_timeout(store_fixture: Light
    completed = await store_fixture.wait_for_rollouts(rollout_ids=["nonexistent"], timeout=0.1)
    elapsed = time.time() - start

-    # Should timeout quickly (not wait indefinitely)
-    assert elapsed < 0.2
+    if isinstance(store_fixture, InMemoryLightningStore):
+        # Should timeout quickly (not wait indefinitely)
+        assert elapsed < 0.2
+    else:
+        # Should be slower, but not too slow
+        assert elapsed < 2.0
    assert len(completed) == 0


@@ -2080,8 +2085,12 @@ async def test_wait_polling_interval_with_timeout_none(store_fixture: LightningS
    completed = await wait_and_complete()
    elapsed = time.time() - start

-    # Should complete after ~0.5s (when we set the event)
-    assert 0.4 < elapsed < 0.7
+    if isinstance(store_fixture, InMemoryLightningStore):
+        # Should complete after ~0.5s (when we set the event)
+        assert 0.4 < elapsed < 0.7
+    else:
+        # Should be more than 5 seconds
+        assert 5 < elapsed < 15
    assert len(completed) == 1
    assert completed[0].status == "succeeded"

@@ -19,6 +19,7 @@ import pytest
 import pytest_asyncio
 from portpicker import pick_unused_port

+from agentlightning.store import LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.types import (
@@ -69,12 +70,11 @@ async def _run_server_with_cors(cors_origins: List[str] | str | None = None):


@pytest_asyncio.fixture
-async def server_client() -> (
-    AsyncGenerator[Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], None]
-):
-    store = InMemoryLightningStore()
+async def server_client(
+    store_fixture: LightningStore,
+) -> AsyncGenerator[Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], None]:
    port = pick_unused_port()
-    server = LightningStoreServer(store, "127.0.0.1", port)
+    server = LightningStoreServer(store_fixture, "127.0.0.1", port)
    await server.start()
    client = LightningStoreClient(server.endpoint)
    session = aiohttp.ClientSession()
@@ -1,9 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.

 import multiprocessing
+import sys
 from typing import Any, Optional, Union

 import agentops
+import pytest
 from agentops.sdk.core import TraceContext
 from opentelemetry.trace.status import StatusCode

@@ -20,7 +22,8 @@ def _func_without_exception():
    pass


-def test_trace_error_status_from_instance():
+@pytest.mark.parametrize("with_exception", [True, False])
+def test_trace_error_status_from_instance(with_exception: bool):
    """
    Test that AgentOpsTracer correctly sets trace end state based on execution result.

@@ -30,7 +33,7 @@ def test_trace_error_status_from_instance():
    """

    ctx = multiprocessing.get_context("spawn")
-    proc = ctx.Process(target=_test_trace_error_status_from_instance_imp)
+    proc = ctx.Process(target=_test_trace_error_status_from_instance_imp, args=(with_exception,))
    proc.start()
    proc.join(30.0)  # On GPU server, the time is around 10 seconds.

@@ -42,13 +45,19 @@ def test_trace_error_status_from_instance():

        assert False, "Child process hung. Check test output for details."

-    assert proc.exitcode == 0, (
-        f"Child process for test_trace_error_status_from_instance failed with exit code {proc.exitcode}. "
-        "Check child traceback in test output."
-    )
+    if with_exception:
+        assert proc.exitcode != 0, (
+            f"Child process for test_trace_error_status_from_instance with exception exited with exit code {proc.exitcode}. "
+            "Check child traceback in test output."
+        )
+    else:
+        assert proc.exitcode == 0, (
+            f"Child process for test_trace_error_status_from_instance without exception failed with exit code {proc.exitcode}. "
+            "Check child traceback in test output."
+        )


-def _test_trace_error_status_from_instance_imp():
+def _test_trace_error_status_from_instance_imp(with_exception: bool):
    captured_state = {}
    old_end_trace = agentops.end_trace

@@ -65,12 +74,14 @@ def _test_trace_error_status_from_instance_imp():
    tracer.init_worker(0)

    try:
-        tracer.trace_run(_func_with_exception)
-        assert captured_state["state"] == StatusCode.ERROR
-
-        tracer.trace_run(_func_without_exception)
-        assert captured_state["state"] == StatusCode.OK
-
+        if with_exception:
+            tracer.trace_run(_func_with_exception)
+            if captured_state["state"] != StatusCode.ERROR:
+                sys.exit(-1)
+        else:
+            tracer.trace_run(_func_without_exception)
+            if captured_state["state"] != StatusCode.OK:
+                sys.exit(-1)
    finally:
        agentops.end_trace = old_end_trace
        tracer.teardown_worker(0)
@@ -39,6 +39,7 @@ import httpx
 import litellm
 import openai
 import pytest
+import requests
 import uvicorn
 from agents import Agent, AgentHooks, GuardrailFunctionOutput, InputGuardrail, RunConfig, Runner, function_tool
 from agents.mcp import MCPServerStdio
@@ -74,27 +75,21 @@ from agentlightning.types import Span, Triplet
 from ..common.tracer import clear_agentops_init, clear_tracer_provider

 USE_OPENAI = os.environ.get("USE_OPENAI", "false").lower() == "true"
+OPENAI_BASE_URL = "http://127.0.0.1:58000/v1"
+OPENAI_MODEL = "gpt-4.1-mini"
+OPENAI_API_KEY = "token-abc123"
+
+REAL_OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL")
+REAL_OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 if USE_OPENAI:
-    OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", os.environ["OPENAI_API_BASE"])
-    OPENAI_MODEL = "gpt-4.1-mini"
-    OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
-else:
-    OPENAI_BASE_URL = "http://127.0.0.1:58000/v1"
-    OPENAI_MODEL = "gpt-4.1-mini"
-    OPENAI_API_KEY = "token-abc123"
+    assert (
+        REAL_OPENAI_BASE_URL is not None and REAL_OPENAI_API_KEY is not None
+    ), "OPENAI_BASE_URL and OPENAI_API_KEY must be set when USE_OPENAI is true"


 _langchain_callback_handler = None


-class ChatCompletionRequest(BaseModel):
-    model: str
-    messages: List[Dict[str, Any]]
-    stream: bool = False
-    tools: Optional[List[Any]] = None
-    tool_choice: Optional[Any] = None
-
-
 class MockOpenAICompatibleServer:
    """
    A mock server that mimics the OpenAI Chat Completions API for testing purposes.
@@ -113,8 +108,11 @@ class MockOpenAICompatibleServer:
        self.prompt_caches = self._load_prompt_caches()
        self._setup_routes()

+    def _prompt_cache_path(self) -> str:
+        return os.path.join(os.path.dirname(__file__), "../assets/prompt_caches.jsonl")
+
    def _load_prompt_caches(self):
-        cache_path = os.path.join(os.path.dirname(__file__), "../assets/prompt_caches.jsonl")
+        cache_path = self._prompt_cache_path()
        caches = []
        if os.path.exists(cache_path):
            with open(cache_path, "r") as f:
@@ -158,10 +156,23 @@ class MockOpenAICompatibleServer:

    def _setup_routes(self):
        @self.app.post("/v1/chat/completions")
-        def chat_completions(request: ChatCompletionRequest):
+        def chat_completions(request: Dict[str, Any]):
+            if USE_OPENAI:
+                # Call Real OpenAI API to get prompt cache
+                response = requests.post(
+                    REAL_OPENAI_BASE_URL.rstrip("/") + "/chat/completions",
+                    json=request,
+                    headers={"Authorization": f"Bearer {REAL_OPENAI_API_KEY}"},
+                )
+                if response.status_code != 200:
+                    raise ValueError(f"Failed to call OpenAI API: {response.status_code} {response.text}")
+                response_dict = response.json()
+                with open(self._prompt_cache_path(), "a") as f:
+                    f.write(json.dumps({"request": request, "response": response_dict}) + "\n")
+                return response_dict
+
            # Try to find the best match in prompt caches
-            request_dict = request.model_dump()
-            cached_response, score = self._find_best_cache_match(request_dict)
+            cached_response, score = self._find_best_cache_match(request)
            if cached_response and score > 0.8:
                time.sleep(0.1)  # Simulate network delay
                # Return the cached response directly
@@ -781,33 +792,6 @@ def run_with_http_tracer() -> None:
    tracer.teardown()


-def create_prompt_caches() -> None:
-    """Create prompt caches for the agent frameworks.
-    This should only be run once to populate the caches.
-    """
-
-    if USE_OPENAI:
-        tracer = HttpTracer()
-        with tracer._trace_context_sync():
-            run_all()
-
-        with open(os.path.join(os.path.dirname(__file__), "../assets/prompt_caches.jsonl"), "w") as f:
-            for span in tracer._last_records.requests.values():
-                if span.url.startswith(OPENAI_BASE_URL) and span.status_code < 400 and span.response.content:
-                    f.write(
-                        json.dumps(
-                            {
-                                "request": json.loads(span.request.content.decode()),
-                                "response": json.loads(span.response.content.decode()),
-                            }
-                        )
-                        + "\n"
-                    )
-
-    else:
-        run_all()
-
-
@pytest.mark.parametrize("agent_func_name", [f.__name__ for f in iterate_over_agents()], ids=str)
 def test_run_with_agentops_tracer(agent_func_name: str):
    """AgentOps tracer tests are notoriously problematic and does not work well with other tests."""
Author	SHA1	Message	Date
Yuge Zhang	7b6e763079	.	2025-11-24 00:10:28 +08:00
Yuge Zhang	03ed353da9	.	2025-11-24 00:09:05 +08:00
Yuge Zhang	87f459886a	.	2025-11-24 00:06:45 +08:00
Yuge Zhang	a791ef6447	update contributing guide	2025-11-23 23:52:39 +08:00
Yuge Zhang	215cc8fe74	.	2025-11-23 23:28:53 +08:00
Yuge Zhang	bda205128b	.	2025-11-23 11:30:23 +08:00
Yuge Zhang	c71a58fa08	how to contribute	2025-11-23 11:23:13 +08:00
Yuge Zhang	bbc4d35c7f	Add examples catalog	2025-11-23 10:57:51 +08:00
Yuge Zhang	bffc7013f9	Store Benchmark - Part 1 (#328 )	2025-11-22 23:35:29 +08:00
Yuge Zhang	4cf8fb94e7	Github Actions Workflow for Tinker and Azure (#327 )	2025-11-22 01:47:53 +08:00
Yuge Zhang	ab185a5c5a	MongoDB-based Lightning Store (#323 )	2025-11-21 11:49:54 +08:00
Yuge Zhang	d581cbcd63	Upgrade VM image (#325 )	2025-11-20 17:49:16 +08:00
Yuge Zhang	3459caa1de	Fix OpenAI Agents 0.6 compatibility and pin vLLM < 0.11.1 (#322 )	2025-11-20 07:13:15 +08:00
Yuge Zhang	f3fd58e72a	Put store init in the right place of tracer (#321 )	2025-11-19 20:35:27 +08:00
Yuge Zhang	b3cb5e1337	Minor improvements to make RL workflow more robust (#319 )	2025-11-18 15:40:51 +08:00