Fix formatting in store.py for metrics tracking

.github/workflows/badge-examples.yml

@@ -10,6 +10,7 @@ on:
       - Examples - Tinker
       - Examples - Azure
       - Examples - Claude Code
+      - Examples - RAG
     types: [completed]
 
   workflow_dispatch:
@@ -37,5 +38,6 @@ jobs:
               { workflow: 'examples-tinker.yml', label: 'examples-tinker.stable', variants: ['stable'] },
               { workflow: 'examples-azure.yml', label: 'examples-azure.stable', variants: ['stable'] },
               { workflow: 'examples-claude-code.yml', label: 'examples-claude-code.stable', variants: ['stable'] },
+              { workflow: 'examples-rag.yml', label: 'examples-rag.stable', variants: ['stable'] },
             ];
             await badgeAggregation({ github, context, core, dependencies });

.github/workflows/badge-latest.yml

@@ -7,6 +7,8 @@ on:
       - Examples - Spider
       - Examples - APO
       - Examples - Unsloth
+      - Examples - RAG
+      - Examples - Claude Code
       - GPU Test
     types: [completed]
 
@@ -32,6 +34,8 @@ jobs:
               { workflow: 'examples-spider.yml', label: 'spider.latest', variants: ['latest'] },
               { workflow: 'examples-apo.yml', label: 'apo.latest', variants: ['latest'] },
               { workflow: 'examples-unsloth.yml', label: 'unsloth.latest', variants: ['latest'] },
+              { workflow: 'examples-claude-code.yml', label: 'claude-code.latest', variants: ['latest'] },
+              { workflow: 'examples-rag.yml', label: 'rag.latest', variants: ['latest'] },
               { workflow: 'tests-full.yml', label: 'tests-full.latest', variants: ['latest'] },
             ];
             await badgeAggregation({ github, context, core, dependencies });

.github/workflows/badge-rag.yml

@@ -0,0 +1,29 @@
+name: Badge - RAG
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - RAG
+    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-rag.yml', label: 'rag', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/benchmark.yml

@@ -6,9 +6,9 @@ on:
 
 jobs:
   benchmark:
-    name: Benchmark (${{ matrix.backend.id }}, ${{ matrix.scenario.display }})
-    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
-    timeout-minutes: 60
+    name: ${{ matrix.workload.kind }} (${{ matrix.backend.id }}, ${{ matrix.workload.display }})
+    runs-on: ${{ matrix.workload.runner }}
+    timeout-minutes: ${{ matrix.workload.timeout }}
     strategy:
       fail-fast: false
       matrix:
@@ -17,10 +17,15 @@ jobs:
             compose_file: compose.prometheus-memory-store.yml
           - id: mongo
             compose_file: compose.prometheus-mongo-store.yml
-        scenario:
+        workload:
           - id: minimal-production
             display: Minimal production scale
+            kind: scenario
             store_workers: 4
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
             args: >-
               --mode batch
               --total-tasks 4096
@@ -30,7 +35,12 @@ jobs:
               --sleep-seconds 0.5
           - id: medium-production
             display: Medium production scale
+            kind: scenario
             store_workers: 16
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
             args: >-
               --mode batch
               --total-tasks 10000
@@ -40,7 +50,12 @@ jobs:
               --sleep-seconds 0.1
           - id: large-batch
             display: Large batch waves
+            kind: scenario
             store_workers: 32
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
             args: >-
               --mode batch
               --total-tasks 100000
@@ -50,7 +65,12 @@ jobs:
               --sleep-seconds 0.1
           - id: long-queues
             display: Long rollout queues
+            kind: scenario
             store_workers: 32
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
             args: >-
               --mode batch_partial
               --total-tasks 100000
@@ -61,7 +81,12 @@ jobs:
               --sleep-seconds 0.1
           - id: high-concurrency
             display: High-throughput concurrent requests
+            kind: scenario
             store_workers: 32
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
             args: >-
               --mode single
               --total-tasks 100000
@@ -71,7 +96,12 @@ jobs:
               --sleep-seconds 0.1
           - id: heavy-traces
             display: Heavy rollouts with deep traces
+            kind: scenario
             store_workers: 64
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
             args: >-
               --mode batch_partial
               --total-tasks 10000
@@ -80,15 +110,63 @@ jobs:
               --n-runners 512
               --max-rounds 20
               --sleep-seconds 1.0
+
+          - id: micro-worker
+            display: Update worker
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: worker
+          - id: micro-dequeue-empty
+            display: Dequeue empty
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: dequeue-empty
+          - id: micro-rollout
+            display: Rollout + span
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: rollout
+          - id: micro-dequeue-update-attempt
+            display: Dequeue + update attempt
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: dequeue-update-attempt
+          - id: micro-dequeue-only
+            display: Dequeue only
+            kind: micro
+            store_workers: 8
+            runner: ubuntu-latest
+            timeout: 30
+            cli: dequeue-only
+          - id: micro-metrics
+            display: Multi-metric fan-out
+            kind: micro
+            store_workers: 1
+            runner: ubuntu-latest
+            timeout: 15
+            cli: metrics
     env:
       STORE_URL: http://localhost:4747
       STORE_API_URL: http://localhost:4747/v1/agl
       PROM_URL: http://localhost:9090
-      SCENARIO_ID: ${{ matrix.scenario.id }}
+      WORKLOAD_KIND: ${{ matrix.workload.kind }}
+      WORKLOAD_ID: ${{ matrix.workload.id }}
       BACKEND_ID: ${{ matrix.backend.id }}
-      ARTIFACT_DIR: artifacts/${{ matrix.scenario.id }}-${{ matrix.backend.id }}
+      ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.workload.id, matrix.backend.id) }}
       COMPOSE_FILE: ${{ matrix.backend.compose_file }}
-      AGL_STORE_N_WORKERS: ${{ matrix.scenario.store_workers }}
+      AGL_STORE_N_WORKERS: ${{ matrix.workload.store_workers }}
+      ANALYSIS_FILE: ${{ format('analysis-{0}.log', matrix.workload.id) }}
+      SUMMARY_FILE: ${{ format('summary-{0}.log', matrix.workload.id) }}
+      PROM_ARCHIVE_BASENAME: ${{ format('prometheus-{0}-{1}', matrix.workload.id, matrix.backend.id) }}
+      ARTIFACT_NAME: ${{ format('{0}-{1}-{2}', matrix.workload.kind == 'micro' && 'micro-benchmark' || 'benchmark', matrix.workload.id, matrix.backend.id) }}
     steps:
       - uses: actions/checkout@v4
 
@@ -127,164 +205,57 @@ jobs:
             sleep 1
           done
           echo "Store did not become ready in time" >&2
-          docker compose -f "$COMPOSE_FILE" logs app
-          exit 1
-
-      - name: Prepare artifact directory
-        run: mkdir -p "$ARTIFACT_DIR"
-
-      - name: Record benchmark start
-        run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
-
-      - name: Run ${{ matrix.scenario.display }} workload
-        run: |
-          set -euo pipefail
-          uv run --locked --no-sync python -m tests.benchmark.benchmark_store \
-            --store-url "$STORE_URL" \
-            ${{ matrix.scenario.args }}
-
-      - name: Record benchmark end
-        if: ${{ always() }}
-        run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
-
-      - name: Run benchmark analysis
-        if: ${{ always() }}
-        run: |
-          set -euo pipefail
-          mkdir -p "$ARTIFACT_DIR"
-          if [ -z "${BENCHMARK_START:-}" ] || [ -z "${BENCHMARK_END:-}" ]; then
-            echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/analysis.txt"
-            exit 1
-          fi
-          uv run --locked --no-sync python -m tests.benchmark.analysis \
-            --prom-url "$PROM_URL" \
-            --store-url "$STORE_API_URL" \
-            --start "$BENCHMARK_START" \
-            --end "$BENCHMARK_END" \
-            | tee "$ARTIFACT_DIR/analysis.txt"
-
-      - name: Stop ${{ matrix.backend.id }} Prometheus stack
-        if: ${{ always() }}
-        run: |
-          set -euo pipefail
-          cd docker
-          docker compose -f "$COMPOSE_FILE" down -v || true
-
-      - name: Archive Prometheus metrics
-        if: ${{ always() }}
-        run: |
-          set -euo pipefail
-          mkdir -p "$ARTIFACT_DIR"
-          if [ -d docker/data/prometheus ]; then
-            tar -C docker/data -czf "$ARTIFACT_DIR/prometheus-${SCENARIO_ID}-${BACKEND_ID}.tar.gz" prometheus
-          fi
-          if docker compose -f "$COMPOSE_FILE" ps --format '{{.Name}}' >/dev/null 2>&1; then
-            docker compose -f "$COMPOSE_FILE" logs app > "$ARTIFACT_DIR/docker-${SCENARIO_ID}-${BACKEND_ID}.log" || true
-          fi
-
-      - name: Upload benchmark artifacts
-        if: ${{ always() }}
-        uses: actions/upload-artifact@v4
-        with:
-          name: benchmark-${{ matrix.scenario.id }}-${{ matrix.backend.id }}
-          path: ${{ env.ARTIFACT_DIR }}
-          if-no-files-found: error
-
-  micro-benchmark:
-    name: Micro-benchmark (${{ matrix.backend.id }}, ${{ matrix.mode.display }})
-    runs-on: ubuntu-latest
-    timeout-minutes: 30
-    strategy:
-      fail-fast: false
-      matrix:
-        backend:
-          - id: memory
-            compose_file: compose.prometheus-memory-store.yml
-          - id: mongo
-            compose_file: compose.prometheus-mongo-store.yml
-        mode:
-          - id: worker
-            display: Update worker throughput
-            cli: worker
-          - id: dequeue-empty
-            display: Dequeue empty throughput
-            cli: dequeue-empty
-          - id: rollout
-            display: Rollout + span throughput
-            cli: rollout
-    env:
-      STORE_URL: http://localhost:4747
-      STORE_API_URL: http://localhost:4747/v1/agl
-      PROM_URL: http://localhost:9090
-      BACKEND_ID: ${{ matrix.backend.id }}
-      MODE_ID: ${{ matrix.mode.id }}
-      ARTIFACT_DIR: artifacts/micro-${{ matrix.mode.id }}-${{ matrix.backend.id }}
-      COMPOSE_FILE: ${{ matrix.backend.compose_file }}
-      AGL_STORE_N_WORKERS: 8
-    steps:
-      - uses: actions/checkout@v4
-
-      - uses: astral-sh/setup-uv@v7
-        with:
-          enable-cache: true
-          python-version: '3.12'
-
-      - name: Sync dependencies
-        run: uv sync --frozen --extra mongo --group core-stable --group dev
-
-      - name: Reset benchmark data directories
-        run: |
-          set -euo pipefail
-          cd docker
-          rm -rf data
-          bash setup.sh
-
-      - name: Launch ${{ matrix.backend.id }} Prometheus stack
-        run: |
-          set -euo pipefail
-          cd docker
-          docker compose -f "$COMPOSE_FILE" down -v || true
-          docker compose -f "$COMPOSE_FILE" up -d --quiet-pull
-
-      - name: Wait for store readiness
-        run: |
-          set -euo pipefail
-          for attempt in {1..60}; do
-            if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
-              exit 0
-            fi
-            sleep 1
-          done
-          echo "Store did not become ready in time" >&2
+          # show logs for debugging
           cd docker && docker compose -f "$COMPOSE_FILE" logs app
           exit 1
 
       - name: Prepare artifact directory
         run: mkdir -p "$ARTIFACT_DIR"
 
-      - name: Record micro benchmark start
+      - name: Record workload start
         run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
 
-      - name: Run ${{ matrix.mode.display }}
+      - name: (Scenario) Run ${{ matrix.workload.display }} workload
+        if: ${{ matrix.workload.kind == 'scenario' }}
+        run: |
+          set -euo pipefail
+          uv run --locked --no-sync python -m tests.benchmark.benchmark_store \
+            --store-url "$STORE_URL" \
+            ${{ matrix.workload.args }}
+
+      - name: (Micro) Run ${{ matrix.workload.display }}
+        if: ${{ matrix.workload.kind == 'micro' }}
         run: |
           set -euo pipefail
           mkdir -p "$ARTIFACT_DIR"
           uv run --locked --no-sync python -m tests.benchmark.micro_benchmark \
             --store-url "$STORE_URL" \
-            --summary-file "$ARTIFACT_DIR/summary-${MODE_ID}.txt" \
-            "${{ matrix.mode.cli }}" | tee "$ARTIFACT_DIR/micro-${MODE_ID}.txt"
+            --summary-file "$ARTIFACT_DIR/$SUMMARY_FILE" \
+            "${{ matrix.workload.cli }}" | tee "$ARTIFACT_DIR/${{ matrix.workload.id }}.txt"
 
-      - name: Record micro benchmark end
+      - name: Record workload end
         if: ${{ always() }}
         run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
 
-      - name: Run micro benchmark analysis
+      - name: Show micro benchmark summary
+        if: ${{ always() && matrix.workload.kind == 'micro' }}
+        run: |
+          set -euo pipefail
+          summary_file="$ARTIFACT_DIR/$SUMMARY_FILE"
+          if [ -f "$summary_file" ]; then
+            echo "Micro benchmark summary ($WORKLOAD_ID/$BACKEND_ID):"
+            cat "$summary_file"
+          else
+            echo "Summary file not found: $summary_file"
+          fi
+
+      - name: Run workload analysis
         if: ${{ always() }}
         run: |
           set -euo pipefail
           mkdir -p "$ARTIFACT_DIR"
           if [ -z "${BENCHMARK_START:-}" ] || [ -z "${BENCHMARK_END:-}" ]; then
-            echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/analysis-${MODE_ID}.txt"
+            echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/$ANALYSIS_FILE"
             exit 1
           fi
           uv run --locked --no-sync python -m tests.benchmark.analysis \
@@ -292,19 +263,22 @@ jobs:
             --store-url "$STORE_API_URL" \
             --start "$BENCHMARK_START" \
             --end "$BENCHMARK_END" \
-            | tee "$ARTIFACT_DIR/analysis-${MODE_ID}.txt"
+            | tee "$ARTIFACT_DIR/$ANALYSIS_FILE"
 
-      - name: Show micro benchmark summary
+      - name: Collect docker logs
         if: ${{ always() }}
         run: |
           set -euo pipefail
-          summary_file="$ARTIFACT_DIR/summary-${MODE_ID}.txt"
-          if [ -f "$summary_file" ]; then
-            echo "Micro benchmark summary ($MODE_ID/$BACKEND_ID):"
-            cat "$summary_file"
-          else
-            echo "Summary file not found: $summary_file"
+          mkdir -p "$ARTIFACT_DIR"
+          cd docker
+          readarray -t services < <(docker compose -f "$COMPOSE_FILE" config --services)
+          if [ "${#services[@]}" -eq 0 ]; then
+            echo "No services defined in compose file."
+            exit 0
           fi
+          for service in "${services[@]}"; do
+            docker compose -f "$COMPOSE_FILE" logs "$service" > "../$ARTIFACT_DIR/docker-${service}-${WORKLOAD_ID}-${BACKEND_ID}.log" || true
+          done
 
       - name: Stop ${{ matrix.backend.id }} Prometheus stack
         if: ${{ always() }}
@@ -319,16 +293,13 @@ jobs:
           set -euo pipefail
           mkdir -p "$ARTIFACT_DIR"
           if [ -d docker/data/prometheus ]; then
-            tar -C docker/data -czf "$ARTIFACT_DIR/prometheus-micro-${MODE_ID}-${BACKEND_ID}.tar.gz" prometheus
-          fi
-          if docker compose -f "$COMPOSE_FILE" ps --format '{{.Name}}' >/dev/null 2>&1; then
-            docker compose -f "$COMPOSE_FILE" logs app > "$ARTIFACT_DIR/docker-micro-${MODE_ID}-${BACKEND_ID}.log" || true
+            tar -C docker/data -czf "$ARTIFACT_DIR/${PROM_ARCHIVE_BASENAME}.tar.gz" prometheus
           fi
 
-      - name: Upload micro benchmark artifacts
+      - name: Upload workload artifacts
         if: ${{ always() }}
         uses: actions/upload-artifact@v4
         with:
-          name: micro-benchmark-${{ matrix.mode.id }}-${{ matrix.backend.id }}
+          name: ${{ env.ARTIFACT_NAME }}
           path: ${{ env.ARTIFACT_DIR }}
           if-no-files-found: error

.gitignore

@@ -3,6 +3,7 @@ verl_old
 meta-llama/**
 debug/*.png
 requirements-freeze*.txt
+/playground
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

AGENTS.md

@@ -0,0 +1,41 @@
+# Repository Guidelines
+
+## Architecture Overview
+Agent Lightning runs through a continuous loop: runners and tracers emit spans, `LightningStore` (`agentlightning/store/`) keeps them synchronized, and algorithms in `agentlightning/algorithm/` consume those traces to improve behavior.
+
+## Project Structure & Module Organization
+- `agentlightning/`: adapters, execution stack, training loop, tracer, reward logic, and the `agl` CLI.
+- `docs/` & `examples/`: narrative and procedural docs (assets in `docs/assets/`, navigation in `mkdocs.yml`) plus runnable workflows whose READMEs point to their companion how-to guides. `docs/how-to` covers task-focused instructions, while `docs/tutorials` explains concepts and subsystems.
+- `dashboard/`, `scripts/`, `tests/`: UI bundles, release/dataset/CI automation, and mirrored coverage of the runtime tree. Record download steps rather than committing binaries.
+
+## Build, Test, and Development Commands
+- `uv sync --group dev` — provision tooling once per environment.
+- `uv run --no-sync pytest -v` — execute the full suite; add a path or `-k expr` to narrow the run.
+- `uv run --no-sync pyright` — enforce static typing parity with CI.
+- `uv run --no-sync pre-commit run --all-files --show-diff-on-failure` and `uv run --no-sync mkdocs build --strict` — keep formatting tidy and documentation valid.
+Always commit the refreshed `uv.lock` when dependencies shift, and mention optional groups (VERL, APO, GPU) in PR notes.
+
+## Common Issues & Fixes
+- When `uv run` errors with `Permission denied` under `~/.cache`, override both cache locations inline: ``UV_CACHE="$(pwd)/.cache_uv" XDG_CACHE_HOME="$(pwd)/.cache_xdg" uv run --no-sync <command>``.
+
+## Coding Style & Naming Conventions
+- Target `requires-python >= 3.10`, four-space indentation, 120-character lines (though docstrings may run longer), and formatter-owned diffs (Black + isort, `black` profile). Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
+- Maintain exhaustive type hints (pyright enforces them) and prefer shared dataclasses or Pydantic models from `agentlightning.types`.
+- Author Google-style docstrings for new modules or public methods—succinct descriptions, no redundant type info, no redundant `Key features/components` bullet points, and `[][]` syntax for cross-references.
+- Writing logs is encouraged, especially for long functions with multiple steps and try-except blocks that catch all exceptions. Use `logging.getLogger(__name__)` to get loggers. Distinguish between DEBUG, INFO, WARNING, and ERROR logs.
+
+## Testing Guidelines
+- Mirror runtime directories under `tests/` and match filenames for quick traceability.
+- Parametrize pytest cases and apply markers (`openai`, `gpu`, `agentops`, `mongo`, `llmproxy`) so optional suites can be skipped via selectors like `-m "not mongo"` yet still exercised in CI.
+- Lean on fixtures, favor real stores/spans/agents over mocks, and drive coverage across the majority of branches.
+- If an imported module is missing from the environment, check whether `uv sync` has been run with the right groups. Do not make stubs for external dependencies unless necessary.
+
+## Example Contributions
+- Ship each example with a README that includes smoke-test instructions so maintainers can validate quickly. The README must contain an "Included Files" section summarizing every file and its role.
+- Keep runnable example modules self-contained with a module-level docstring describing CLI usage. Document important or educational classes/functions with targeted docstrings and inline comments where clarity matters.
+- Add a CI workflow per example named `examples-<name>.yml` in `.github/workflows/`. Register it in `badge-<name>.yml`, `badge-examples.yml`, and `badge-latest.yml` when applicable so badges stay accurate.
+
+## Commit & Pull Request Guidelines
+- Branch from a fresh `main` using `feature/<slug>`, `fix/<slug>`, `docs/<slug>`, or `chore/<slug>`.
+- Write imperative, scoped commits, reference issues with `Fixes #123`, and rerun pre-commit plus the relevant pytest/doc builds before pushing.
+- Use PR descriptions to summarize intent, list verification commands, call out dependency or docs-navigation updates, and link new docs/examples via `mkdocs.yml` or `examples/README.md`. Include logs for dashboard changes.

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

agentlightning/cli/store.py

@@ -7,11 +7,18 @@ from __future__ import annotations
 import argparse
 import asyncio
 import logging
-from typing import Iterable
+from typing import Iterable, List
 
 from agentlightning import setup_logging
 from agentlightning.store.client_server import LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+    setup_multiprocess_prometheus,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -33,9 +40,10 @@ def main(argv: Iterable[str] | None = None) -> int:
         help="Configure the logging level for the store.",
     )
     parser.add_argument(
-        "--prometheus",
-        action="store_true",
-        help="Enable Prometheus metrics.",
+        "--tracker",
+        nargs="+",
+        choices=["prometheus", "console"],
+        help="Enable metrics tracking. Repeat for multiple trackers.",
     )
     parser.add_argument(
         "--n-workers",
@@ -63,14 +71,36 @@ def main(argv: Iterable[str] | None = None) -> int:
 
     setup_logging(args.log_level)
 
+    trackers: List[MetricsBackend] = []
+    if args.tracker:
+        if "prometheus" in args.tracker:
+            logger.info("Enabling Prometheus metrics tracking.")
+            if args.n_workers > 1:
+                # This has to be done before prometheus_client is imported
+                setup_multiprocess_prometheus()
+                logger.info("Setting up Prometheus multiprocess directory for metrics tracking.")
+            trackers.append(PrometheusMetricsBackend())
+
+        if "console" in args.tracker:
+            logger.info("Enabling console metrics tracking.")
+            trackers.append(ConsoleMetricsBackend())
+
+    if len(trackers) == 0:
+        tracker: MetricsBackend | None = None
+    elif len(trackers) == 1:
+        tracker = trackers[0]
+    else:
+        tracker = MultiMetricsBackend(trackers)
+
     if args.backend == "memory":
         store = InMemoryLightningStore(
-            prometheus=args.prometheus, thread_safe=True
-        )  # Using thread_safe store for server
+            thread_safe=True,  # Using thread_safe store for server
+            tracker=tracker,
+        )
     elif args.backend == "mongo":
         from agentlightning.store.mongo import MongoLightningStore
 
-        store = MongoLightningStore(client=args.mongo_uri, prometheus=args.prometheus)
+        store = MongoLightningStore(client=args.mongo_uri, tracker=tracker)
     else:
         raise ValueError(f"Invalid backend: {args.backend}")
 
@@ -86,7 +116,7 @@ def main(argv: Iterable[str] | None = None) -> int:
         port=args.port,
         cors_allow_origins=args.cors_origins,
         launch_mode=launch_mode,
-        prometheus=args.prometheus,
+        tracker=tracker,
         n_workers=args.n_workers,
     )
     try:

agentlightning/runner/agent.py

@@ -582,6 +582,7 @@ class LitAgentRunner(Runner[T_task]):
                 while not (event is not None and event.is_set()):
                     logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
                     next_rollout = await store.dequeue_rollout(worker_id=self.get_worker_id())
+                    logger.debug(f"{self._log_prefix()} Next rollout retrieved: {next_rollout}")
                     if next_rollout is None:
                         logger.debug(
                             f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds."

agentlightning/store/client_server.py

@@ -59,10 +59,12 @@ from agentlightning.types import (
     Worker,
     WorkerStatus,
 )
+from agentlightning.utils.metrics import MetricsBackend, get_prometheus_registry
 from agentlightning.utils.otlp import handle_otlp_export, spans_from_proto
 from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncher, PythonServerLauncherArgs
 
 from .base import UNSET, LightningStore, LightningStoreCapabilities, LightningStoreStatistics, Unset
+from .collection.base import resolve_error_type
 from .utils import LATENCY_BUCKETS
 
 server_logger = logging.getLogger("agentlightning.store.server")
@@ -238,7 +240,7 @@ class LightningStoreServer(LightningStore):
         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.
+        tracker: The metrics tracker to use for the server.
     """
 
     def __init__(
@@ -250,7 +252,7 @@ class LightningStoreServer(LightningStore):
         launch_mode: LaunchMode = "thread",
         launcher_args: PythonServerLauncherArgs | None = None,
         n_workers: int = 1,
-        prometheus: bool = False,
+        tracker: MetricsBackend | None = None,
     ):
         super().__init__()
         self.store = store
@@ -287,7 +289,7 @@ class LightningStoreServer(LightningStore):
             app=self.app,
             args=self.launcher_args,
         )
-        self._prometheus = prometheus
+        self._tracker = tracker
 
         self._lock: threading.Lock = threading.Lock()
         self._cors_allow_origins = self._normalize_cors_origins(cors_allow_origins)
@@ -332,7 +334,6 @@ class LightningStoreServer(LightningStore):
         return {
             "launcher_args": self.launcher_args,
             "server_launcher": self.server_launcher,
-            "_prometheus": self._prometheus,
             "_owner_pid": self._owner_pid,
         }
 
@@ -350,11 +351,12 @@ class LightningStoreServer(LightningStore):
         self.store = None
         self.launcher_args = state["launcher_args"]
         self.server_launcher = state["server_launcher"]
-        self._prometheus = state["_prometheus"]
+        self._tracker = None
         self._owner_pid = state["_owner_pid"]
         self._cors_allow_origins = state.get("_cors_allow_origins")
         self._client = None
         self._lock = threading.Lock()
+        self._prometheus_registry = None
         # Do NOT reconstruct app, _uvicorn_config, _uvicorn_server
         # to avoid transferring server state to subprocess
 
@@ -435,8 +437,8 @@ class LightningStoreServer(LightningStore):
         api = APIRouter(prefix=API_V1_PREFIX)
 
         # The outermost-layer of monitoring
-        if self._prometheus:
-            self._setup_prometheus(api=api, app=self.app)
+        if self._tracker is not None:
+            self._setup_metrics(api=api, app=self.app)
 
         # TODO: This should only be enabled in development mode.
         @self.app.middleware("http")
@@ -844,41 +846,21 @@ class LightningStoreServer(LightningStore):
         # Finally, mount the dashboard assets
         self._setup_dashboard()
 
-    def _setup_prometheus(self, api: APIRouter, app: FastAPI):
+    def _setup_metrics(self, api: APIRouter, app: FastAPI):
         """Setup Prometheus metrics endpoints."""
-        try:
-            from prometheus_client import make_asgi_app  # type: ignore
-            from prometheus_client import (
-                REGISTRY,
-                CollectorRegistry,
-                Counter,
-                Histogram,
-                multiprocess,
-            )
-        except ImportError:
-            raise ImportError(
-                "Prometheus client is not installed. Please either install it or set prometheus to False."
-            )
+        if self._tracker is None:
+            return
 
-        # Multi-process mode: https://prometheus.github.io/client_python/multiprocess/
-        is_multiprocess = self.launcher_args.launch_mode == "mp" and self.launcher_args.n_workers > 1
-        if is_multiprocess:
-            registry = CollectorRegistry()
-            multiprocess.MultiProcessCollector(registry)
-        else:
-            registry = REGISTRY
-
-        HTTP_REQUESTS = Counter(
-            "http_requests_total",
-            "Total HTTP requests",
-            ["method", "path", "status_code"],
+        self._tracker.register_counter(
+            "agl.http.total",
+            ["path", "method", "status"],
+            group_level=2,
         )
-
-        HTTP_LATENCY = Histogram(
-            "http_request_duration_seconds",
-            "Latency of HTTP requests",
-            ["method", "path", "status_code"],
+        self._tracker.register_histogram(
+            "agl.http.latency",
+            ["path", "method", "status"],
             buckets=LATENCY_BUCKETS,
+            group_level=2,
         )
 
         def get_template_path(path: str) -> str:
@@ -904,9 +886,12 @@ class LightningStoreServer(LightningStore):
             return path
 
         @app.middleware("http")
-        async def prometheus_http_middleware(  # pyright: ignore[reportUnusedFunction]
+        async def tracking_middleware(  # pyright: ignore[reportUnusedFunction]
             request: Request, call_next: Callable[[Request], Awaitable[Response]]
         ) -> Response:
+            if self._tracker is None:
+                return await call_next(request)
+
             start = time.perf_counter()
             status = 520  # Default to 520 if things crash hard
 
@@ -918,9 +903,8 @@ class LightningStoreServer(LightningStore):
                 # Client disconnected (Timeout)
                 status = 499  # Standard Nginx code for "Client Closed Request"
                 raise  # Re-raise to let Uvicorn handle the cleanup
-            except Exception:
-                # TODO: Record the error type
-                status = 500
+            except Exception as exc:
+                status = resolve_error_type(exc)
                 raise
             finally:
                 # This block executes NO MATTER WHAT happens above
@@ -930,13 +914,25 @@ class LightningStoreServer(LightningStore):
                 path = get_template_path(request.url.path)
                 method = request.method
 
-                HTTP_REQUESTS.labels(method, path, status).inc()
-                HTTP_LATENCY.labels(method, path, status).observe(elapsed)
+                await self._tracker.inc_counter(
+                    "agl.http.total",
+                    labels={"method": method, "path": path, "status": str(status)},
+                )
+                await self._tracker.observe_histogram(
+                    "agl.http.latency",
+                    value=elapsed,
+                    labels={"method": method, "path": path, "status": str(status)},
+                )
 
-        metrics_app = make_asgi_app(registry=registry)  # type: ignore
+        if self._tracker.has_prometheus():
+            from prometheus_client import make_asgi_app  # pyright: ignore[reportUnknownVariableType]
 
-        # This App would need to be accessed via /v1/prometheus/ (note the trailing slash)
-        app.mount(api.prefix + "/prometheus", metrics_app)  # pyright: ignore[reportUnknownArgumentType]
+            metrics_app = make_asgi_app(  # pyright: ignore[reportUnknownVariableType]
+                registry=get_prometheus_registry()
+            )
+
+            # This App would need to be accessed via /v1/prometheus/ (note the trailing slash)
+            app.mount(api.prefix + "/prometheus", metrics_app)  # pyright: ignore[reportUnknownArgumentType]
 
     def _setup_otlp(self, api: APIRouter):
         """Setup OTLP endpoints."""

agentlightning/store/collection/base.py

@@ -2,6 +2,10 @@
 
 from __future__ import annotations
 
+import functools
+import inspect
+import time
+from contextlib import asynccontextmanager
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -22,9 +26,13 @@ from typing import (
     cast,
 )
 
+from agentlightning.store.utils import LATENCY_BUCKETS
+from agentlightning.utils.metrics import MetricsBackend
+
 if TYPE_CHECKING:
     from typing import Self
 
+from agentlightning.store.base import LightningStore
 from agentlightning.types import (
     Attempt,
     FilterField,
@@ -40,6 +48,7 @@ from agentlightning.types import (
 T = TypeVar("T")  # Recommended to be a BaseModel
 K = TypeVar("K")
 V = TypeVar("V")
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 
 AtomicMode = Literal["r", "w", "rw"]
 """What is expected within the atomic context. Can be "read", "write", or "read-write"."""
@@ -51,8 +60,135 @@ These labels are used to identify the collections that are affected by the atomi
 """
 
 
-class Collection(Generic[T]):
-    """Behaves like a list of items. Supporting addition, updating, and deletion of items."""
+COLLECTION_TRACKING_STORE_METHODS = frozenset(
+    [name for name in LightningStore.__dict__ if not name.startswith("_")] + ["_healthcheck"]
+)
+
+_UNKNOWN_STORE_METHOD = "unknown"
+
+
+def _nearest_lightning_store_method_from_stack() -> str:
+    """Stack introspection so that we capture the nearest public API method from the
+    call stack whenever metrics are recorded."""
+    frame = inspect.currentframe()
+    try:
+        if frame is None:
+            return _UNKNOWN_STORE_METHOD
+        frame = frame.f_back
+        while frame is not None:
+            self_obj = frame.f_locals.get("self")
+            method_name = frame.f_locals.get("method_name")
+            if method_name in COLLECTION_TRACKING_STORE_METHODS and isinstance(self_obj, LightningStore):
+                return method_name
+            frame = frame.f_back
+        return _UNKNOWN_STORE_METHOD
+    except Exception:
+        return _UNKNOWN_STORE_METHOD
+    finally:
+        del frame
+
+
+def resolve_error_type(exc: BaseException | None) -> str:
+    if exc is None:
+        return "N/A"
+
+    try:
+        from .mongo import resolve_mongo_error_type
+
+        error_type = resolve_mongo_error_type(exc)
+        if error_type is not None:
+            return error_type
+    except ImportError:
+        # If the mongo backend is not available, fall back to using the exception's class name.
+        pass
+
+    return exc.__class__.__name__
+
+
+def tracked(operation: str):
+    """Decorator to track the execution of the decorated method."""
+
+    def decorator(func: T_callable) -> T_callable:
+
+        @functools.wraps(func)
+        async def wrapper(self: TrackedCollection, *args: Any, **kwargs: Any) -> Any:
+            async with self.tracking_context(operation, self.collection_name):
+                return await func(self, *args, **kwargs)
+
+        return cast(T_callable, wrapper)
+
+    return decorator
+
+
+class TrackedCollection:
+    """An object that can be tracked by the metrics backend."""
+
+    def __init__(self, tracker: MetricsBackend | None = None):
+        self._tracker = tracker
+
+    @property
+    def tracker(self) -> MetricsBackend | None:
+        return self._tracker
+
+    @property
+    def collection_name(self) -> str:
+        """The identifier of the collection."""
+        raise NotImplementedError()
+
+    @property
+    def extra_tracking_labels(self) -> Mapping[str, Any]:
+        """Extra labels to add to the tracking context."""
+        return {}
+
+    @asynccontextmanager
+    async def tracking_context(self, operation: str, collection: str):
+        """Context manager to track the execution of the decorated method.
+
+        Args:
+            operation: The operation to track.
+            collection: The collection to track.
+        """
+        if self._tracker is None:
+            # no-op context manager
+            yield
+
+        else:
+            # Enable tracking
+            start_time = time.perf_counter()
+            status: str = "OK"
+            store_method = _nearest_lightning_store_method_from_stack()
+            try:
+                yield
+            except BaseException as exc:
+                status = resolve_error_type(exc)
+                raise
+            finally:
+                elapsed = time.perf_counter() - start_time
+                await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
+                    "agl.collections.total",
+                    labels={
+                        "store_method": store_method,
+                        "operation": operation,
+                        "collection": collection,
+                        "status": status,
+                        **self.extra_tracking_labels,
+                    },
+                )
+                await self._tracker.observe_histogram(  # pyright: ignore[reportPrivateUsage]
+                    "agl.collections.latency",
+                    value=elapsed,
+                    labels={
+                        "store_method": store_method,
+                        "operation": operation,
+                        "collection": collection,
+                        "status": status,
+                        **self.extra_tracking_labels,
+                    },
+                )
+
+
+class Collection(TrackedCollection, Generic[T]):
+    """Standard collection interface. Behaves like a list of items. Supporting addition, updating, and deletion of items."""
 
     def primary_keys(self) -> Sequence[str]:
         """Get the primary keys of the collection."""
@@ -174,7 +310,7 @@ class Collection(Generic[T]):
         raise NotImplementedError()
 
 
-class Queue(Generic[T]):
+class Queue(TrackedCollection, Generic[T]):
     """Behaves like a deque. Supporting appending items to the end and popping items from the front."""
 
     def __repr__(self) -> str:
@@ -228,7 +364,7 @@ class Queue(Generic[T]):
         raise NotImplementedError()
 
 
-class KeyValue(Generic[K, V]):
+class KeyValue(TrackedCollection, Generic[K, V]):
     """Behaves like a dictionary. Supporting addition, updating, and deletion of items."""
 
     def __repr__(self) -> str:
@@ -255,13 +391,35 @@ class KeyValue(Generic[K, V]):
         raise NotImplementedError()
 
 
-class LightningCollections:
+class LightningCollections(TrackedCollection):
     """Collections of rollouts, attempts, spans, resources, and workers.
 
     [LightningStore][agentlightning.LightningStore] implementations can use this as a storage base
     to implement the store API.
     """
 
+    def __init__(self, tracker: MetricsBackend | None = None, extra_labels: Optional[Sequence[str]] = None):
+        super().__init__(tracker=tracker)
+        self.register_collection_metrics(extra_labels)
+
+    def register_collection_metrics(self, extra_labels: Optional[Sequence[str]] = None) -> None:
+        if self._tracker is None:
+            return
+        labels = ["store_method", "operation", "collection", "status"]
+        if extra_labels is not None:
+            labels.extend(extra_labels)
+        self._tracker.register_histogram(
+            "agl.collections.latency",
+            labels,
+            buckets=LATENCY_BUCKETS,
+            group_level=2,
+        )
+        self._tracker.register_counter("agl.collections.total", labels, group_level=2)
+
+    @property
+    def tracker(self) -> MetricsBackend | None:
+        return self._tracker
+
     @property
     def rollouts(self) -> Collection[Rollout]:
         """Collections of rollouts."""

agentlightning/store/collection/memory.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import logging
-import time
+import uuid
 import weakref
 from collections import deque
 from contextlib import AsyncExitStack, asynccontextmanager
@@ -28,7 +28,6 @@ from typing import (
 import aiologic
 from pydantic import BaseModel
 
-from agentlightning.store.utils import LATENCY_BUCKETS
 from agentlightning.types import (
     Attempt,
     FilterField,
@@ -40,6 +39,7 @@ from agentlightning.types import (
     Span,
     Worker,
 )
+from agentlightning.utils.metrics import MetricsBackend
 
 from .base import (
     AtomicMode,
@@ -50,6 +50,7 @@ from .base import (
     Queue,
     normalize_filter_options,
     resolve_sort_options,
+    tracked,
 )
 
 T = TypeVar("T")  # Recommended to be a BaseModel, not a dict
@@ -192,10 +193,19 @@ class ListBasedCollection(Collection[T]):
        if the field is str-like, 0 if the field is int-like, 0.0 if the field is float-like.
     """
 
-    def __init__(self, items: List[T], item_type: Type[T], primary_keys: Sequence[str]):
+    def __init__(
+        self,
+        items: List[T],
+        item_type: Type[T],
+        primary_keys: Sequence[str],
+        id: Optional[str] = None,
+        tracker: Optional[MetricsBackend] = None,
+    ):
+        super().__init__(tracker=tracker)
         if not primary_keys:
             raise ValueError("primary_keys must be non-empty")
 
+        self._id = id if id is not None else str(uuid.uuid4())
         self._items: Dict[Any, Any] = {}
         self._size: int = 0
         if issubclass(item_type, dict):
@@ -207,6 +217,10 @@ class ListBasedCollection(Collection[T]):
         for item in items or []:
             self._mutate_single(item, mode="insert")
 
+    @property
+    def collection_name(self) -> str:
+        return self._id
+
     def primary_keys(self) -> Sequence[str]:
         """Return the primary key field names for this collection."""
         return self._primary_keys
@@ -483,6 +497,7 @@ class ListBasedCollection(Collection[T]):
             # No items exist for this primary-key prefix.
             return ()
 
+    @tracked("query")
     async def query(
         self,
         filter: Optional[FilterOptions] = None,
@@ -546,6 +561,7 @@ class ListBasedCollection(Collection[T]):
             total=total_matched,
         )
 
+    @tracked("get")
     async def get(
         self,
         filter: Optional[FilterOptions] = None,
@@ -582,6 +598,7 @@ class ListBasedCollection(Collection[T]):
 
         return best_item
 
+    @tracked("insert")
     async def insert(self, items: Sequence[T]) -> None:
         """Insert the given items.
 
@@ -603,6 +620,7 @@ class ListBasedCollection(Collection[T]):
         for item in prepared:
             self._mutate_single(item, mode="insert")
 
+    @tracked("update")
     async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Update the given items.
 
@@ -617,6 +635,7 @@ class ListBasedCollection(Collection[T]):
             updated_items.append(updated)
         return updated_items
 
+    @tracked("upsert")
     async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
         """Upsert the given items (insert if missing, otherwise update)."""
         upserted_items: List[T] = []
@@ -627,6 +646,7 @@ class ListBasedCollection(Collection[T]):
             upserted_items.append(upserted)
         return upserted_items
 
+    @tracked("delete")
     async def delete(self, items: Sequence[T]) -> None:
         """Delete the given items.
 
@@ -646,23 +666,37 @@ class DequeBasedQueue(Queue[T]):
     Provides O(1) amortized enqueue (append) and dequeue (popleft).
     """
 
-    def __init__(self, item_type: Type[T], items: Optional[Sequence[T]] = None):
+    def __init__(
+        self,
+        item_type: Type[T],
+        items: Optional[Sequence[T]] = None,
+        id: Optional[str] = None,
+        tracker: Optional[MetricsBackend] = None,
+    ):
+        super().__init__(tracker=tracker)
         self._items: Deque[T] = deque()
         self._item_type: Type[T] = item_type
+        self._id = id if id is not None else str(uuid.uuid4())
         if items:
             self._items.extend(items)
 
     def item_type(self) -> Type[T]:
         return self._item_type
 
+    @property
+    def collection_name(self) -> str:
+        return self._id
+
     def __repr__(self) -> str:
         return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({len(self._items)})>"
 
+    @tracked("has")
     async def has(self, item: T) -> bool:
         if not isinstance(item, self._item_type):
             raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
         return item in self._items
 
+    @tracked("enqueue")
     async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
         for item in items:
             if not isinstance(item, self._item_type):
@@ -670,6 +704,7 @@ class DequeBasedQueue(Queue[T]):
             self._items.append(item)
         return items
 
+    @tracked("dequeue")
     async def dequeue(self, limit: int = 1) -> Sequence[T]:
         if limit <= 0:
             return []
@@ -678,6 +713,7 @@ class DequeBasedQueue(Queue[T]):
             out.append(self._items.popleft())
         return out
 
+    @tracked("peek")
     async def peek(self, limit: int = 1) -> Sequence[T]:
         if limit <= 0:
             return []
@@ -689,6 +725,7 @@ class DequeBasedQueue(Queue[T]):
             result.append(item)
         return result
 
+    @tracked("size")
     async def size(self) -> int:
         return len(self._items)
 
@@ -696,21 +733,34 @@ class DequeBasedQueue(Queue[T]):
 class DictBasedKeyValue(KeyValue[K, V]):
     """KeyValue implementation backed by a plain dictionary."""
 
-    def __init__(self, data: Optional[Mapping[K, V]] = None):
+    def __init__(
+        self, data: Optional[Mapping[K, V]] = None, id: Optional[str] = None, tracker: Optional[MetricsBackend] = None
+    ):
+        super().__init__(tracker=tracker)
         self._values: Dict[K, V] = dict(data) if data else {}
+        self._id = id if id is not None else str(uuid.uuid4())
 
+    @property
+    def collection_name(self) -> str:
+        return self._id
+
+    @tracked("has")
     async def has(self, key: K) -> bool:
         return key in self._values
 
+    @tracked("get")
     async def get(self, key: K, default: V | None = None) -> V | None:
         return self._values.get(key, default)
 
+    @tracked("set")
     async def set(self, key: K, value: V) -> None:
         self._values[key] = value
 
+    @tracked("pop")
     async def pop(self, key: K, default: V | None = None) -> V | None:
         return self._values.pop(key, default)
 
+    @tracked("size")
     async def size(self) -> int:
         return len(self._values)
 
@@ -721,7 +771,8 @@ class InMemoryLightningCollections(LightningCollections):
     Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
     """
 
-    def __init__(self, lock_type: Literal["thread", "asyncio"], prometheus: bool = False):
+    def __init__(self, lock_type: Literal["thread", "asyncio"], tracker: MetricsBackend | None = None):
+        super().__init__(tracker=tracker)
         self._lock = {
             "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
             "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
@@ -731,31 +782,29 @@ class InMemoryLightningCollections(LightningCollections):
             "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
             "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
         }
-        self._rollouts = ListBasedCollection(items=[], item_type=Rollout, primary_keys=["rollout_id"])
-        self._attempts = ListBasedCollection(items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"])
-        self._spans = ListBasedCollection(
-            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"]
+        self._rollouts = ListBasedCollection(
+            items=[], item_type=Rollout, primary_keys=["rollout_id"], id="rollouts", tracker=tracker
         )
-        self._resources = ListBasedCollection(items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"])
-        self._workers = ListBasedCollection(items=[], item_type=Worker, primary_keys=["worker_id"])
-        self._rollout_queue = DequeBasedQueue(items=[], item_type=str)
-        self._span_sequence_ids = DictBasedKeyValue[str, int](data={})  # rollout_id -> sequence_id
+        self._attempts = ListBasedCollection(
+            items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"], id="attempts", tracker=tracker
+        )
+        self._spans = ListBasedCollection(
+            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"], id="spans", tracker=tracker
+        )
+        self._resources = ListBasedCollection(
+            items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"], id="resources", tracker=tracker
+        )
+        self._workers = ListBasedCollection(
+            items=[], item_type=Worker, primary_keys=["worker_id"], id="workers", tracker=tracker
+        )
+        self._rollout_queue = DequeBasedQueue(items=[], item_type=str, id="rollout_queue", tracker=tracker)
+        self._span_sequence_ids = DictBasedKeyValue[str, int](
+            data={}, id="span_sequence_ids", tracker=tracker
+        )  # rollout_id -> sequence_id
 
-        self._prometheus = prometheus
-        if self._prometheus:
-            from prometheus_client import Counter, Histogram
-
-            self._rate_metric = Counter(
-                "memory_collection_lock_rate",
-                "Rate of memory collection locks",
-                ["collection"],
-            )
-            self._latency_metric = Histogram(
-                "memory_collection_lock_latency_seconds",
-                "Latency of memory collection locks",
-                ["collection"],
-                buckets=LATENCY_BUCKETS,
-            )
+    @property
+    def collection_name(self) -> str:
+        return "router"
 
     @property
     def rollouts(self) -> ListBasedCollection[Rollout]:
@@ -807,17 +856,15 @@ class InMemoryLightningCollections(LightningCollections):
         # are trying to acquire the same locks in different orders.
         labels = sorted(labels)
 
-        managers = [(label, self._lock[label]) for label in labels]
-        async with AsyncExitStack() as stack:
-            for label, manager in managers:
-                start_time = time.perf_counter()
-                await stack.enter_async_context(manager)
-                elapsed = time.perf_counter() - start_time
-                if self._prometheus:
-                    self._rate_metric.labels(collection=label).inc()
-                    self._latency_metric.labels(collection=label).observe(elapsed)
-            yield self
+        async with self.tracking_context(operation="atomic", collection=self.collection_name):
+            managers = [(label, self._lock[label]) for label in labels]
+            async with AsyncExitStack() as stack:
+                for label, manager in managers:
+                    async with self.tracking_context(operation="lock", collection=label):
+                        await stack.enter_async_context(manager)
+                yield self
 
+    @tracked("evict_spans_for_rollout")
     async def evict_spans_for_rollout(self, rollout_id: str) -> None:
         """Evict all spans for a given rollout ID.
 

agentlightning/store/collection/mongo.py

@@ -3,9 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-import contextvars
-import functools
-import inspect
 import logging
 import random
 import re
@@ -30,6 +27,8 @@ from typing import (
     cast,
 )
 
+from agentlightning.utils.metrics import MetricsBackend
+
 if TYPE_CHECKING:
     from typing import Self
 
@@ -41,8 +40,6 @@ from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.errors import CollectionInvalid, ConnectionFailure, DuplicateKeyError, OperationFailure, PyMongoError
 from pymongo.read_concern import ReadConcern
 
-from agentlightning.store.base import LightningStore
-from agentlightning.store.utils import LATENCY_BUCKETS
 from agentlightning.types import (
     Attempt,
     FilterOptions,
@@ -62,6 +59,7 @@ from .base import (
     Queue,
     normalize_filter_options,
     resolve_sort_options,
+    tracked,
 )
 
 T_model = TypeVar("T_model", bound=BaseModel)
@@ -77,236 +75,27 @@ V = TypeVar("V")
 
 logger = logging.getLogger(__name__)
 
-_OPERATION_CONTEXT: contextvars.ContextVar["_MongoOperationContext | None"] = contextvars.ContextVar(
-    "_mongo_operation_context", default=None
-)
 
-_LIGHTNING_STORE_PUBLIC_METHODS = frozenset(
-    [name for name, value in LightningStore.__dict__.items() if not name.startswith("_") and callable(value)]
-    + ["_healthcheck"]
-)
-
-_UNKNOWN_STORE_METHOD = "unknown"
-
-
-def _nearest_lightning_store_method_from_stack() -> str:
-    """Stack introspection so that we capture the nearest public API method from the
-    call stack whenever metrics are recorded."""
-    frame = inspect.currentframe()
-    try:
-        if frame is None:
-            return _UNKNOWN_STORE_METHOD
-        frame = frame.f_back
-        while frame is not None:
-            self_obj = frame.f_locals.get("self")
-            method_name = frame.f_locals.get("method_name")
-            if method_name in _LIGHTNING_STORE_PUBLIC_METHODS and isinstance(self_obj, LightningStore):
-                return method_name
-            frame = frame.f_back
-        return _UNKNOWN_STORE_METHOD
-    except Exception:
-        return _UNKNOWN_STORE_METHOD
-    finally:
-        del frame
-
-
-class MongoOperationPrometheusTracker:
-    """A tracker for MongoDB operations metrics.
-
-    All classes should share one single instance of this tracker.
-    """
-
-    def __init__(self, enabled: bool):
-        self._enabled = enabled
-
-        if enabled:
-            from prometheus_client import Counter, Histogram
-
-            base_labels = ["operation", "database", "collection", "store_method"]
-            self._latency_metric = Histogram(
-                "mongo_operation_duration_seconds",
-                "Latency of MongoDB operations",
-                base_labels,
-                buckets=LATENCY_BUCKETS,
-            )
-            self._total_metric = Counter(
-                "mongo_operation_total",
-                "Total MongoDB operations",
-                base_labels + ["status"],
-            )
-            self._error_metric = Counter(
-                "mongo_operation_errors_total",
-                "Total MongoDB operations that failed",
-                base_labels + ["error_type"],
-            )
-            self._num_attempts_metric = Histogram(
-                "mongo_operation_num_attempts",
-                "Number of attempts for MongoDB operations",
-                base_labels,
-                buckets=list(range(10)) + list(range(10, 100, 5)),
-            )
-
-    def track(self, operation: str, database: str, collection: str) -> _MongoOperationContext | _DummyOperationContext:
-        if not self._enabled:
-            return _DummyOperationContext()
-        return _MongoOperationContext(self, operation, database, collection)
-
-    @staticmethod
-    def classify_error(exc: BaseException | None) -> str:
-        if exc is None:
-            return "Other"
-        is_transient = isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError")
-        if isinstance(exc, OperationFailure):
-            if is_transient:
-                return f"OperationFailure-{exc.code}-Transient"
-            else:
-                return f"OperationFailure-{exc.code}"
-        if isinstance(exc, DuplicateKeyError):
-            return "DuplicateKeyError-Transient" if is_transient else "DuplicateKeyError"
-        if isinstance(exc, PyMongoError):
-            if is_transient:
-                return f"{exc.__class__.__name__}-Transient"
-            else:
-                return exc.__class__.__name__
-        if isinstance(exc, ConnectionFailure):
-            return "ConnectionFailure-Transient" if is_transient else "ConnectionFailure"
-        return "Other-Transient" if is_transient else "Other"
-
-    def observe(
-        self,
-        *,
-        operation: str,
-        database: str,
-        collection: str,
-        elapsed: float,
-        status: str,
-        error_type: str | None,
-        num_attempts: int | None = None,
-    ) -> None:
-        if not self._enabled:
-            return
-
-        store_method = _nearest_lightning_store_method_from_stack()
-        self._total_metric.labels(operation, database, collection, store_method, status).inc()
-        self._latency_metric.labels(operation, database, collection, store_method).observe(elapsed)
-        if status == "error" and error_type:
-            self._error_metric.labels(operation, database, collection, store_method, error_type).inc()
-        if num_attempts is not None:
-            self._num_attempts_metric.labels(operation, database, collection, store_method).observe(num_attempts)
-
-
-class _MongoOperationContext:
-    """A context manager for tracking MongoDB operations.
-
-    Used via:
-
-    ```python
-    with self.tracker.track("insert", "database", "collection") as track_context:
-        try:
-            await collection.insert_one({})
-        except Exception as exc:
-            # For errors that can be ignored, report the error to the tracker.
-            track_context.report_error(exc)
-    ```
-    """
-
-    def __init__(self, tracker: MongoOperationPrometheusTracker, operation: str, database: str, collection: str):
-        self._tracker = tracker
-        self._operation = operation
-        self._database = database
-        self._collection = collection
-        self._start: float = 0.0
-        self._active: bool = False
-        self._error_type: str | None = None
-        self._num_attempts: int | None = None
-
-    def __enter__(self) -> "_MongoOperationContext":
-        self._active = True
-        self._start = time.perf_counter()
-        return self
-
-    def __exit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> bool:
-        if not self._active:
-            return False
-
-        elapsed = time.perf_counter() - self._start
-
-        # Try to classify the error
-        if self._error_type is not None:
-            error_type = self._error_type
-        elif exc is not None:
-            error_type = self._tracker.classify_error(exc)
+def resolve_mongo_error_type(exc: BaseException | None) -> str | None:
+    is_transient = isinstance(exc, PyMongoError) and exc.has_error_label("TransientTransactionError")
+    if isinstance(exc, OperationFailure):
+        if is_transient:
+            return f"OperationFailure-{exc.code}-Transient"
         else:
-            error_type = None
-
-        status = "ok" if error_type is None else "error"
-        self._tracker.observe(
-            operation=self._operation,
-            database=self._database,
-            collection=self._collection,
-            elapsed=elapsed,
-            status=status,
-            error_type=error_type,
-            num_attempts=self._num_attempts,
-        )
-        return False
-
-    def report_error(self, exc: BaseException) -> None:
-        """Used to report errors that occurred in the middle of an operation."""
-        self._error_type = self._tracker.classify_error(exc)
-
-    def report_num_attempts(self, num_attempts: int) -> None:
-        """Used to report the number of attempts that occurred in the middle of an operation."""
-        self._num_attempts = num_attempts
-
-
-class _DummyOperationContext:
-    """A dummy context manager that does nothing, but compatible with _MongoOperationContext."""
-
-    def __init__(self):
-        pass
-
-    def __enter__(self) -> "_DummyOperationContext":
-        return self
-
-    def __exit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> bool:
-        return False
-
-    def report_error(self, exc: BaseException) -> None:
-        pass
-
-    def report_num_attempts(self, num_attempts: int) -> None:
-        pass
-
-
-def _mongo_operation(operation: str) -> Callable[[T_callable], T_callable]:
-    def decorator(func: T_callable) -> T_callable:
-        if not asyncio.iscoroutinefunction(func):
-            raise TypeError(f"_mongo_operation decorator requires coroutine functions, got {func.__name__}")
-
-        @functools.wraps(func)
-        async def wrapper(
-            self: (
-                MongoBasedCollection[T_model]
-                | MongoBasedQueue[T_generic]
-                | MongoBasedKeyValue[K, V]
-                | MongoLightningCollections
-            ),
-            *args: Any,
-            **kwargs: Any,
-        ) -> Any:
-            tracker = self._prometheus_tracker  # pyright: ignore[reportPrivateUsage]
-            if not tracker._enabled:  # pyright: ignore[reportPrivateUsage]
-                # Skip the tracking because tracking is not configured
-                return await func(self, *args, **kwargs)
-            with tracker.track(
-                operation, self._database_name, self._collection_name  # pyright: ignore[reportPrivateUsage]
-            ):
-                return await func(self, *args, **kwargs)
-
-        return wrapper  # type: ignore
-
-    return decorator
+            return f"OperationFailure-{exc.code}"
+    if isinstance(exc, DuplicateKeyError):
+        return "DuplicateKeyError-Transient" if is_transient else "DuplicateKeyError"
+    if isinstance(exc, PyMongoError):
+        if is_transient:
+            return f"{exc.__class__.__name__}-Transient"
+        else:
+            return exc.__class__.__name__
+    if isinstance(exc, ConnectionFailure):
+        return "ConnectionFailure-Transient" if is_transient else "ConnectionFailure"
+    if is_transient:
+        return "Other-Transient"
+    else:
+        return None
 
 
 def _field_ops_to_conditions(field: str, ops: Mapping[str, Any]) -> List[Dict[str, Any]]:
@@ -545,6 +334,7 @@ class MongoBasedCollection(Collection[T_model]):
         primary_keys: The primary keys of the collection.
         item_type: The type of the items in the collection.
         extra_indexes: The extra indexes to create on the collection.
+        tracker: The metrics tracker to use.
     """
 
     def __init__(
@@ -556,8 +346,9 @@ class MongoBasedCollection(Collection[T_model]):
         primary_keys: Sequence[str],
         item_type: Type[T_model],
         extra_indexes: Sequence[Sequence[str]] = [],
-        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
+        tracker: MetricsBackend | None = None,
     ):
+        super().__init__(tracker=tracker)
         if isinstance(client_pool, AsyncMongoClient):
             self._client_pool = MongoClientPool(client_pool)
         else:
@@ -568,9 +359,6 @@ class MongoBasedCollection(Collection[T_model]):
         self._collection_created = False
         self._extra_indexes = [list(index) for index in extra_indexes]
         self._session: Optional[AsyncClientSession] = None
-        self._prometheus_tracker = (
-            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
-        )
 
         if not primary_keys:
             raise ValueError("primary_keys must be non-empty")
@@ -580,7 +368,17 @@ class MongoBasedCollection(Collection[T_model]):
             raise ValueError(f"item_type must be a subclass of BaseModel, got {item_type.__name__}")
         self._item_type = item_type
 
-    @_mongo_operation("ensure_collection")
+    @property
+    def collection_name(self) -> str:
+        return self._collection_name
+
+    @property
+    def extra_tracking_labels(self) -> Mapping[str, str]:
+        return {
+            "database": self._database_name,
+        }
+
+    @tracked("ensure_collection")
     async def ensure_collection(self) -> AsyncCollection[Mapping[str, Any]]:
         """Ensure the backing MongoDB collection exists (and optionally its indexes).
 
@@ -606,7 +404,7 @@ class MongoBasedCollection(Collection[T_model]):
             primary_keys=self._primary_keys,
             item_type=self._item_type,
             extra_indexes=self._extra_indexes,
-            prometheus_tracker=self._prometheus_tracker,
+            tracker=self._tracker,
         )
         collection._collection_created = self._collection_created
         collection._session = session
@@ -619,7 +417,7 @@ class MongoBasedCollection(Collection[T_model]):
     def item_type(self) -> Type[T_model]:
         return self._item_type
 
-    @_mongo_operation("size")
+    @tracked("size")
     async def size(self) -> int:
         collection = await self.ensure_collection()
         return await collection.count_documents({"partition_id": self._partition_id}, session=self._session)
@@ -672,7 +470,7 @@ class MongoBasedCollection(Collection[T_model]):
         # Convert Mongo document to Pydantic model
         return self._item_type.model_validate(raw)  # type: ignore[arg-type]
 
-    @_mongo_operation("query")
+    @tracked("query")
     async def query(
         self,
         filter: Optional[FilterOptions] = None,
@@ -718,7 +516,7 @@ class MongoBasedCollection(Collection[T_model]):
 
         return PaginatedResult[T_model](items=items, limit=limit, offset=offset, total=total)
 
-    @_mongo_operation("get")
+    @tracked("get")
     async def get(
         self,
         filter: Optional[FilterOptions] = None,
@@ -747,7 +545,7 @@ class MongoBasedCollection(Collection[T_model]):
 
         return self._model_validate_item(raw)
 
-    @_mongo_operation("insert")
+    @tracked("insert")
     async def insert(self, items: Sequence[T_model]) -> None:
         if not items:
             return
@@ -779,23 +577,20 @@ class MongoBasedCollection(Collection[T_model]):
         else:
             existing_filter = {"partition_id": self._partition_id, "$or": pk_conditions}
 
-        with self._prometheus_tracker.track("insert__find_existing", self._database_name, self._collection_name):
+        async with self.tracking_context("insert.find_existing", self._collection_name):
             existing = await collection.find_one(existing_filter, session=self._session)
             if existing is not None:
                 existing_values = tuple(existing.get(pk) for pk in self._primary_keys)
                 raise ValueError(f"Item with primary key(s) {self._render_pk_values(existing_values)} already exists")
 
-        with self._prometheus_tracker.track(
-            "insert__insert_many", self._database_name, self._collection_name
-        ) as tracker:
-            try:
+        try:
+            async with self.tracking_context("insert.insert_many", self._collection_name):
                 await collection.insert_many(docs, session=self._session)
-            except DuplicateKeyError as exc:
-                # In case the DB enforces uniqueness via index, normalize to ValueError
-                tracker.report_error(exc)
-                raise ValueError("Duplicate key error while inserting items") from exc
+        except DuplicateKeyError as exc:
+            # In case the DB enforces uniqueness via index, normalize to ValueError
+            raise ValueError("Duplicate key error while inserting items") from exc
 
-    @_mongo_operation("update")
+    @tracked("update")
     async def update(self, items: Sequence[T_model], update_fields: Sequence[str] | None = None) -> List[T_model]:
         if not items:
             return []
@@ -813,9 +608,7 @@ class MongoBasedCollection(Collection[T_model]):
 
             # Branch 1: Full Replace
             if update_fields is None:
-                with self._prometheus_tracker.track(
-                    "update__find_one_and_replace", self._database_name, self._collection_name
-                ):
+                async with self.tracking_context("update.find_one_and_replace", self._collection_name):
                     updated_doc = await collection.find_one_and_replace(
                         filter=pk_filter,
                         replacement=doc,
@@ -826,9 +619,7 @@ class MongoBasedCollection(Collection[T_model]):
             # Branch 2: Partial Update
             else:
                 update_doc = {field: doc[field] for field in update_fields if field in doc}
-                with self._prometheus_tracker.track(
-                    "update__find_one_and_update", self._database_name, self._collection_name
-                ):
+                async with self.tracking_context("update.find_one_and_update", self._collection_name):
                     updated_doc = await collection.find_one_and_update(
                         filter=pk_filter,
                         update={"$set": update_doc},
@@ -846,7 +637,7 @@ class MongoBasedCollection(Collection[T_model]):
 
         return updated_items
 
-    @_mongo_operation("upsert")
+    @tracked("upsert")
     async def upsert(self, items: Sequence[T_model], update_fields: Sequence[str] | None = None) -> List[T_model]:
         if not items:
             return []
@@ -878,9 +669,7 @@ class MongoBasedCollection(Collection[T_model]):
             if update_subset:
                 update_spec["$set"] = update_subset
 
-            with self._prometheus_tracker.track(
-                "upsert__find_one_and_update", self._database_name, self._collection_name
-            ):
+            async with self.tracking_context("upsert.find_one_and_update", self._collection_name):
                 result_doc = await collection.find_one_and_update(
                     filter=pk_filter,
                     update=update_spec,
@@ -895,7 +684,7 @@ class MongoBasedCollection(Collection[T_model]):
 
         return upserted_items
 
-    @_mongo_operation("delete")
+    @tracked("delete")
     async def delete(self, items: Sequence[T_model]) -> None:
         if not items:
             return
@@ -904,7 +693,7 @@ class MongoBasedCollection(Collection[T_model]):
         for item in items:
             self._ensure_item_type(item)
             pk_filter = self._pk_filter(item)
-            with self._prometheus_tracker.track("delete__delete_one", self._database_name, self._collection_name):
+            async with self.tracking_context("delete.delete_one", self._collection_name):
                 result = await collection.delete_one(pk_filter, session=self._session)
             if result.deleted_count == 0:
                 raise ValueError(f"Item with primary key(s) {pk_filter} does not exist")
@@ -923,7 +712,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
         collection_name: str,
         partition_id: str,
         item_type: Type[T_generic],
-        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
+        tracker: MetricsBackend | None = None,
     ) -> None:
         """
         Args:
@@ -933,6 +722,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
             partition_id: Partition identifier; allows multiple logical queues in one collection.
             item_type: The Python type of queue items (primitive or BaseModel subclass).
         """
+        super().__init__(tracker=tracker)
         if isinstance(client_pool, AsyncMongoClient):
             self._client_pool = MongoClientPool(client_pool)
         else:
@@ -945,14 +735,21 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
         self._collection_created = False
 
         self._session: Optional[AsyncClientSession] = None
-        self._prometheus_tracker = (
-            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
-        )
 
     def item_type(self) -> Type[T_generic]:
         return self._item_type
 
-    @_mongo_operation("ensure_collection")
+    @property
+    def extra_tracking_labels(self) -> Mapping[str, str]:
+        return {
+            "database": self._database_name,
+        }
+
+    @property
+    def collection_name(self) -> str:
+        return self._collection_name
+
+    @tracked("ensure_collection")
     async def ensure_collection(self) -> AsyncCollection[Mapping[str, Any]]:
         """Ensure the backing collection exists.
 
@@ -972,13 +769,13 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
             collection_name=self._collection_name,
             partition_id=self._partition_id,
             item_type=self._item_type,
-            prometheus_tracker=self._prometheus_tracker,
+            tracker=self._tracker,
         )
         queue._collection_created = self._collection_created
         queue._session = session
         return queue
 
-    @_mongo_operation("has")
+    @tracked("has")
     async def has(self, item: T_generic) -> bool:
         collection = await self.ensure_collection()
         encoded = self._adapter.dump_python(item, mode="python")
@@ -992,7 +789,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
         )
         return doc is not None
 
-    @_mongo_operation("enqueue")
+    @tracked("enqueue")
     async def enqueue(self, items: Sequence[T_generic]) -> Sequence[T_generic]:
         if not items:
             return []
@@ -1011,11 +808,11 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
                 }
             )
 
-        with self._prometheus_tracker.track("enqueue__insert_many", self._database_name, self._collection_name):
+        async with self.tracking_context("enqueue.insert_many", self.collection_name):
             await collection.insert_many(docs, session=self._session)
         return list(items)
 
-    @_mongo_operation("dequeue")
+    @tracked("dequeue")
     async def dequeue(self, limit: int = 1) -> Sequence[T_generic]:
         if limit <= 0:
             return []
@@ -1025,9 +822,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
 
         # Atomic claim loop using find_one_and_update
         for _ in range(limit):
-            with self._prometheus_tracker.track(
-                "dequeue__find_one_and_update", self._database_name, self._collection_name
-            ):
+            async with self.tracking_context("dequeue.find_one_and_update", self.collection_name):
                 doc = await collection.find_one_and_update(
                     {
                         "partition_id": self._partition_id,
@@ -1048,13 +843,13 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
 
         return results
 
-    @_mongo_operation("peek")
+    @tracked("peek")
     async def peek(self, limit: int = 1) -> Sequence[T_generic]:
         if limit <= 0:
             return []
 
         collection = await self.ensure_collection()
-        with self._prometheus_tracker.track("peek__find", self._database_name, self._collection_name):
+        async with self.tracking_context("peek.find", self.collection_name):
             cursor = (
                 collection.find(
                     {
@@ -1074,7 +869,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
 
         return items
 
-    @_mongo_operation("size")
+    @tracked("size")
     async def size(self) -> int:
         collection = await self.ensure_collection()
         return await collection.count_documents(
@@ -1097,7 +892,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         partition_id: str,
         key_type: Type[K],
         value_type: Type[V],
-        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
+        tracker: MetricsBackend | None = None,
     ) -> None:
         """
         Args:
@@ -1107,7 +902,9 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
             partition_id: Partition identifier; allows multiple logical maps in one collection.
             key_type: The Python type of keys (primitive or BaseModel).
             value_type: The Python type of values (primitive or BaseModel).
+            tracker: The metrics tracker to use.
         """
+        super().__init__(tracker=tracker)
         if isinstance(client_pool, AsyncMongoClient):
             self._client_pool = MongoClientPool(client_pool)
         else:
@@ -1122,11 +919,18 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         self._collection_created = False
 
         self._session: Optional[AsyncClientSession] = None
-        self._prometheus_tracker = (
-            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
-        )
 
-    @_mongo_operation("ensure_collection")
+    @property
+    def extra_tracking_labels(self) -> Mapping[str, str]:
+        return {
+            "database": self._database_name,
+        }
+
+    @property
+    def collection_name(self) -> str:
+        return self._collection_name
+
+    @tracked("ensure_collection")
     async def ensure_collection(self, *, create_indexes: bool = True) -> AsyncCollection[Mapping[str, Any]]:
         """Ensure the backing collection exists (and optionally its indexes)."""
         if not self._collection_created:
@@ -1144,14 +948,14 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
             partition_id=self._partition_id,
             key_type=self._key_type,
             value_type=self._value_type,
-            prometheus_tracker=self._prometheus_tracker,
+            tracker=self._tracker,
         )
         key_value._collection_created = self._collection_created
         key_value._session = session
 
         return key_value
 
-    @_mongo_operation("has")
+    @tracked("has")
     async def has(self, key: K) -> bool:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
@@ -1164,7 +968,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         )
         return doc is not None
 
-    @_mongo_operation("get")
+    @tracked("get")
     async def get(self, key: K, default: V | None = None) -> V | None:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
@@ -1181,15 +985,13 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         raw_value = doc["value"]
         return self._value_adapter.validate_python(raw_value)
 
-    @_mongo_operation("set")
+    @tracked("set")
     async def set(self, key: K, value: V) -> None:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
         encoded_value = self._value_adapter.dump_python(value, mode="python")
-        with self._prometheus_tracker.track(
-            "upsert__replace_one", self._database_name, self._collection_name
-        ) as tracer:
-            try:
+        try:
+            async with self.tracking_context("set.replace_one", self.collection_name):
                 await collection.replace_one(
                     {
                         "partition_id": self._partition_id,
@@ -1203,12 +1005,11 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
                     upsert=True,
                     session=self._session,
                 )
-            except DuplicateKeyError as exc:
-                # Very unlikely with replace_one+upsert, but normalize anyway.
-                tracer.report_error(exc)
-                raise ValueError("Duplicate key error while setting key-value item") from exc
+        except DuplicateKeyError as exc:
+            # Very unlikely with replace_one+upsert, but normalize anyway.
+            raise ValueError("Duplicate key error while setting key-value item") from exc
 
-    @_mongo_operation("pop")
+    @tracked("pop")
     async def pop(self, key: K, default: V | None = None) -> V | None:
         collection = await self.ensure_collection()
         encoded_key = self._key_adapter.dump_python(key, mode="python")
@@ -1225,7 +1026,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
         raw_value = doc["value"]
         return self._value_adapter.validate_python(raw_value)
 
-    @_mongo_operation("size")
+    @tracked("size")
     async def size(self) -> int:
         collection = await self.ensure_collection()
         return await collection.count_documents(
@@ -1254,15 +1055,12 @@ class MongoLightningCollections(LightningCollections):
         workers: Optional[MongoBasedCollection[Worker]] = None,
         rollout_queue: Optional[MongoBasedQueue[str]] = None,
         span_sequence_ids: Optional[MongoBasedKeyValue[str, int]] = None,
-        prometheus_tracker: MongoOperationPrometheusTracker | None = None,
+        tracker: MetricsBackend | None = None,
     ):
+        super().__init__(tracker=tracker, extra_labels=["database"])
         self._client_pool = client_pool
         self._database_name = database_name
-        self._collection_name = "collections"  # Special collection name for tracking transactions
         self._partition_id = partition_id
-        self._prometheus_tracker = (
-            prometheus_tracker if prometheus_tracker is not None else MongoOperationPrometheusTracker(enabled=False)
-        )
         self._collection_ensured = False
         self._rollouts = (
             rollouts
@@ -1275,7 +1073,7 @@ class MongoLightningCollections(LightningCollections):
                 ["rollout_id"],
                 Rollout,
                 [["status"]],
-                prometheus_tracker=self._prometheus_tracker,
+                tracker=self._tracker,
             )
         )
         self._attempts = (
@@ -1289,7 +1087,7 @@ class MongoLightningCollections(LightningCollections):
                 ["rollout_id", "attempt_id"],
                 Attempt,
                 [["status"], ["sequence_id"]],
-                prometheus_tracker=self._prometheus_tracker,
+                tracker=self._tracker,
             )
         )
         self._spans = (
@@ -1303,7 +1101,7 @@ class MongoLightningCollections(LightningCollections):
                 ["rollout_id", "attempt_id", "span_id"],
                 Span,
                 [["sequence_id"]],
-                prometheus_tracker=self._prometheus_tracker,
+                tracker=self._tracker,
             )
         )
         self._resources = (
@@ -1317,7 +1115,7 @@ class MongoLightningCollections(LightningCollections):
                 ["resources_id"],
                 ResourcesUpdate,
                 ["update_time"],
-                prometheus_tracker=self._prometheus_tracker,
+                tracker=self._tracker,
             )
         )
         self._workers = (
@@ -1331,7 +1129,7 @@ class MongoLightningCollections(LightningCollections):
                 ["worker_id"],
                 Worker,
                 ["status"],
-                prometheus_tracker=self._prometheus_tracker,
+                tracker=self._tracker,
             )
         )
         self._rollout_queue = (
@@ -1343,7 +1141,7 @@ class MongoLightningCollections(LightningCollections):
                 "rollout_queue",
                 self._partition_id,
                 str,
-                prometheus_tracker=self._prometheus_tracker,
+                tracker=self._tracker,
             )
         )
         self._span_sequence_ids = (
@@ -1356,10 +1154,20 @@ class MongoLightningCollections(LightningCollections):
                 self._partition_id,
                 str,
                 int,
-                prometheus_tracker=self._prometheus_tracker,
+                tracker=self._tracker,
             )
         )
 
+    @property
+    def collection_name(self) -> str:
+        return "router"  # Special collection name for tracking transactions
+
+    @property
+    def extra_tracking_labels(self) -> Mapping[str, str]:
+        return {
+            "database": self._database_name,
+        }
+
     def with_session(self, session: AsyncClientSession) -> Self:
         instance = self.__class__(
             client_pool=self._client_pool,
@@ -1372,7 +1180,7 @@ class MongoLightningCollections(LightningCollections):
             workers=self._workers.with_session(session),
             rollout_queue=self._rollout_queue.with_session(session),
             span_sequence_ids=self._span_sequence_ids.with_session(session),
-            prometheus_tracker=self._prometheus_tracker,
+            tracker=self._tracker,
         )
         instance._collection_ensured = self._collection_ensured
         return instance
@@ -1405,7 +1213,7 @@ class MongoLightningCollections(LightningCollections):
     def span_sequence_ids(self) -> MongoBasedKeyValue[str, int]:
         return self._span_sequence_ids
 
-    @_mongo_operation("ensure_collections")
+    @tracked("ensure_collections")
     async def _ensure_collections(self) -> None:
         """Ensure all collections exist."""
         if self._collection_ensured:
@@ -1426,14 +1234,14 @@ class MongoLightningCollections(LightningCollections):
         """Perform a atomic operation on the collections."""
         if commit:
             raise ValueError("Commit should be used with execute() instead.")
-        with self._prometheus_tracker.track(f"atomic__{mode}", self._database_name, self._collection_name):
+        async with self.tracking_context("atomic", self.collection_name):
             # First step: ensure all collections exist before going into the atomic block
             if not self._collection_ensured:
                 await self._ensure_collections()
             # Execute directly without commit
             yield self
 
-    @_mongo_operation("execute")
+    @tracked("execute")
     async def execute(
         self,
         callback: Callable[[Self], Awaitable[T_generic]],
@@ -1459,31 +1267,24 @@ class MongoLightningCollections(LightningCollections):
 
         async with client.start_session() as session:
             collections = self.with_session(session)
-            with self._prometheus_tracker.track(
-                "execute__transaction", self._database_name, self._collection_name
-            ) as tracker:
-                try:
-                    return await self._with_transaction(
-                        session, collections, callback, read_concern, write_concern, tracker
-                    )
-                except (ConnectionFailure, OperationFailure) as exc:
-                    # Un-retryable errors.
-                    tracker.report_error(exc)
-                    raise RuntimeError("Transaction failed with connection or operation error") from exc
+            try:
+                return await self.with_transaction(session, collections, callback, read_concern, write_concern)
+            except (ConnectionFailure, OperationFailure) as exc:
+                # Un-retryable errors.
+                raise RuntimeError("Transaction failed with connection or operation error") from exc
 
-    async def _with_transaction(
+    @tracked("with_transaction")
+    async def with_transaction(
         self,
         session: AsyncClientSession,
         collections: Self,
         callback: Callable[[Self], Awaitable[T_generic]],
         read_concern: ReadConcern,
         write_concern: Optional[WriteConcern],
-        transaction_tracker: _MongoOperationContext | _DummyOperationContext,
     ) -> T_generic:
         # This will start a transaction, run transaction callback, and commit.
         # It will also transparently retry on some transient errors.
         # Expanded implementation of with_transaction from client_session
-        num_attempts = 0
         read_preference = ReadPreference.PRIMARY
         transaction_retry_time_limit = 120
         start_time = time.monotonic()
@@ -1492,72 +1293,58 @@ class MongoLightningCollections(LightningCollections):
             return time.monotonic() - start_time < transaction_retry_time_limit
 
         async def _jitter_before_retry() -> None:
-            with self._prometheus_tracker.track("execute__jitter", self._database_name, self._collection_name):
+            async with self.tracking_context("execute.jitter", self.collection_name):
                 await asyncio.sleep(random.uniform(0, 0.05))
 
         while True:
             await session.start_transaction(read_concern, write_concern, read_preference)
 
-            with self._prometheus_tracker.track(
-                "execute__callback", self._database_name, self._collection_name
-            ) as callback_tracker:
-                try:
-                    num_attempts += 1
-                    transaction_tracker.report_num_attempts(num_attempts)
-                    # The _session is always the same within one transaction,
-                    # so we can use the same collections object.
+            try:
+                # The _session is always the same within one transaction,
+                # so we can use the same collections object.
+                async with self.tracking_context("execute.callback", self.collection_name):
                     ret = await callback(collections)
-                # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
-                except BaseException as exc:
-                    callback_tracker.report_error(exc)
-                    if session.in_transaction:
-                        await session.abort_transaction()
-                    if (
-                        isinstance(exc, PyMongoError)
-                        and exc.has_error_label("TransientTransactionError")
-                        and _within_time_limit()
-                    ):
-                        # Retry the entire transaction.
-                        await _jitter_before_retry()
-                        continue
-                    raise
+            # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
+            except BaseException as exc:
+                if session.in_transaction:
+                    await session.abort_transaction()
+                if (
+                    isinstance(exc, PyMongoError)
+                    and exc.has_error_label("TransientTransactionError")
+                    and _within_time_limit()
+                ):
+                    # Retry the entire transaction.
+                    await _jitter_before_retry()
+                    continue
+                raise
 
             if not session.in_transaction:
                 # Assume callback intentionally ended the transaction.
                 return ret
 
-            commit_num_attempts = 0
-
             # Tracks the commit operation.
-            with self._prometheus_tracker.track(
-                "execute__commit", self._database_name, self._collection_name
-            ) as commit_tracker:
+            async with self.tracking_context("execute.commit", self.collection_name):
                 # Loop until the commit succeeds or we hit the time limit.
                 while True:
                     # Tracks the commit attempt.
-                    with self._prometheus_tracker.track(
-                        "execute__commit_attempt", self._database_name, self._collection_name
-                    ) as commit_attempt_tracker:
-                        try:
-                            commit_num_attempts += 1
-                            commit_tracker.report_num_attempts(commit_num_attempts)
+                    try:
+                        async with self.tracking_context("execute.commit_once", self.collection_name):
                             await session.commit_transaction()
-                        except PyMongoError as exc:
-                            commit_attempt_tracker.report_error(exc)
-                            if (
-                                exc.has_error_label("UnknownTransactionCommitResult")
-                                and _within_time_limit()
-                                and not (isinstance(exc, OperationFailure) and exc.code == 50)  # max_time_expired_error
-                            ):
-                                # Retry the commit.
-                                await _jitter_before_retry()
-                                continue
+                    except PyMongoError as exc:
+                        if (
+                            exc.has_error_label("UnknownTransactionCommitResult")
+                            and _within_time_limit()
+                            and not (isinstance(exc, OperationFailure) and exc.code == 50)  # max_time_expired_error
+                        ):
+                            # Retry the commit.
+                            await _jitter_before_retry()
+                            continue
 
-                            if exc.has_error_label("TransientTransactionError") and _within_time_limit():
-                                # Retry the entire transaction.
-                                await _jitter_before_retry()
-                                break
-                            raise
+                        if exc.has_error_label("TransientTransactionError") and _within_time_limit():
+                            # Retry the entire transaction.
+                            await _jitter_before_retry()
+                            break
+                        raise
 
-                        # Commit succeeded.
-                        return ret
+                    # Commit succeeded.
+                    return ret

agentlightning/store/collection_based.py

@@ -60,6 +60,7 @@ from agentlightning.types import (
     Worker,
     WorkerStatus,
 )
+from agentlightning.utils.metrics import MetricsBackend
 
 from .base import (
     UNSET,
@@ -71,7 +72,7 @@ from .base import (
     is_queuing,
 )
 from .collection import FilterOptions, LightningCollections
-from .collection.base import AtomicLabels
+from .collection.base import COLLECTION_TRACKING_STORE_METHODS, AtomicLabels
 from .utils import LATENCY_BUCKETS, rollout_status_from_attempt, scan_unhealthy_rollouts
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
@@ -119,34 +120,36 @@ def _with_collections_execute(labels: Sequence[AtomicLabels]):
 def tracked(name: str):
     """Decorator to track the execution of the decorated method with Prometheus."""
 
-    _public_methods = frozenset([name for name in LightningStore.__dict__ if not name.startswith("_")])
-
     def decorator(func: T_callable) -> T_callable:
 
         @functools.wraps(func)
         async def wrapper(self: CollectionBasedLightningStore[T_collections], *args: Any, **kwargs: Any) -> Any:
-            # For backtracking in Mongo-collection methods.
+            # For backtracking in collection methods.
             # Only track the public methods (+healthcheck)
-            if name in _public_methods or name == "_healthcheck":
+            if name in COLLECTION_TRACKING_STORE_METHODS:
                 method_name = name  # pyright: ignore[reportUnusedVariable]
             else:
                 method_name = None  # pyright: ignore[reportUnusedVariable]
 
-            if not self._prometheus:  # pyright: ignore[reportPrivateUsage]
+            if self._tracker is None:  # pyright: ignore[reportPrivateUsage]
                 # Skip the tracking because tracking is not configured
                 return await func(self, *args, **kwargs)
 
             start_time = time.perf_counter()
+            status: str = "OK"
             try:
-                ret = await func(self, *args, **kwargs)
-                self._total_metric.labels(name, "OK").inc()  # pyright: ignore[reportPrivateUsage]
-                return ret
-            except Exception as exc:
-                self._total_metric.labels(name, exc.__class__.__name__).inc()  # pyright: ignore[reportPrivateUsage]
+                return await func(self, *args, **kwargs)
+            except BaseException as exc:
+                status = exc.__class__.__name__
                 raise
             finally:
                 elapsed = time.perf_counter() - start_time
-                self._latency_metric.labels(name).observe(elapsed)  # pyright: ignore[reportPrivateUsage]
+                await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
+                    "agl.store.total", labels={"method": name, "status": status}
+                )
+                await self._tracker.observe_histogram(  # pyright: ignore[reportPrivateUsage]
+                    "agl.store.latency", value=elapsed, labels={"method": name, "status": status}
+                )
 
         return cast(T_callable, wrapper)
 
@@ -215,40 +218,40 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         read_snapshot: Make sure read operations are atomic. If set to true,
             all read operations like `query_rollouts` will have better consistency.
             It may use an isolated snapshot that supports repeatable reads.
-        prometheus: Enable Prometheus tracking.
+        tracker: Enable metrics tracking.
     """
 
-    def __init__(self, collections: T_collections, *, read_snapshot: bool = False, prometheus: bool = False):
+    def __init__(
+        self, collections: T_collections, *, read_snapshot: bool = False, tracker: MetricsBackend | None = None
+    ):
         # rollouts and spans' storage
         self.collections = collections
         self._read_snapshot = read_snapshot
-        self._prometheus = prometheus
+        self._tracker = tracker
         self._launch_time = time.time()
 
-        if prometheus:
-            from prometheus_client import Counter, Histogram
-
-            self._latency_metric = Histogram(
-                "collection_store_latency_seconds",
-                "Latency of CollectionBasedLightningStore methods",
-                ["method"],
+        if self._tracker is not None:
+            self._tracker.register_histogram(
+                "agl.store.latency",
+                ["method", "status"],
                 buckets=LATENCY_BUCKETS,
+                group_level=1,
             )
-            self._total_metric = Counter(
-                "collection_store_total",
-                "Total MongoDB operations",
-                ["method", "error_type"],
+            self._tracker.register_counter(
+                "agl.store.total",
+                ["method", "status"],
+                group_level=1,
             )
-            self._rollout_counter = Counter(
-                "collection_store_rollout_total",
-                "Total rollouts",
+            self._tracker.register_counter(
+                "agl.rollouts.total",
                 ["status", "mode"],
+                group_level=1,
             )
-            self._rollout_duration_metric = Histogram(
-                "collection_store_rollout_duration_seconds",
-                "Duration of rollouts",
+            self._tracker.register_histogram(
+                "agl.rollouts.duration",
                 ["status", "mode"],
                 buckets=LATENCY_BUCKETS,
+                group_level=1,
             )
 
     async def statistics(self) -> LightningStoreStatistics:
@@ -610,6 +613,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             if not dequeued:
                 break
             rollout_id = dequeued[0]
+            logger.debug("Rollout ID %s has been dequeued by Worker ID %s", rollout_id, worker_id)
 
             post_dequeue_result = await self._post_dequeue_rollouts([rollout_id], worker_id)
             if post_dequeue_result:
@@ -617,6 +621,7 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
                 attempted_rollout, _ = post_dequeue_result[0]
                 if worker_id is not None:
                     await self._sync_workers_with_attempts([attempted_rollout.attempt], dequeue=True)
+                logger.debug("Rollout has been prepared for Worker ID %s: %s", worker_id, attempted_rollout)
                 return attempted_rollout
 
             # else continue the loop
@@ -1140,6 +1145,8 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         if successful_spans:
             await self._post_add_spans(successful_spans, rollout_id, attempt_id)
 
+        logger.debug("Added %d spans for rollout %s, attempt %s", len(successful_spans), rollout_id, attempt_id)
+
         return successful_spans
 
     @tracked("_post_add_spans")
@@ -1220,7 +1227,17 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
                 logger.error(f"Error waiting for rollout {rollout_id}: {rollout}")
 
         # Filter out the exceptions
-        return [rollout for rollout in rollouts if isinstance(rollout, Rollout)]
+        ret = [rollout for rollout in rollouts if isinstance(rollout, Rollout)]
+        finished_rollout_ids = set([rollout.rollout_id for rollout in ret])
+        unfinished_rollout_ids = set(rollout_ids) - finished_rollout_ids
+        logger.debug(
+            "Waiting for rollouts. Number of finished rollouts: %d; number of unfinished rollouts: %d",
+            len(finished_rollout_ids),
+            len(unfinished_rollout_ids),
+        )
+        if len(unfinished_rollout_ids) < 30:
+            logger.debug("Unfinished rollouts: %s", unfinished_rollout_ids)
+        return ret
 
     @tracked("wait_for_rollout")
     async def wait_for_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
@@ -1467,10 +1484,14 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         for rollout, updated_fields in rollouts:
             # Sometimes "end_time" is set but it's not really updated.
             if "end_time" in updated_fields and is_finished(rollout):
-                if self._prometheus:
-                    self._rollout_counter.labels(rollout.status, rollout.mode).inc()
-                    self._rollout_duration_metric.labels(rollout.status, rollout.mode).observe(
-                        cast(float, rollout.end_time) - rollout.start_time
+                if self._tracker is not None:
+                    labels = {
+                        "status": rollout.status,
+                        "mode": rollout.mode if rollout.mode is not None else "unknown",
+                    }
+                    await self._tracker.inc_counter("agl.rollouts.total", labels=labels)
+                    await self._tracker.observe_histogram(
+                        "agl.rollouts.duration", value=cast(float, rollout.end_time) - rollout.start_time, labels=labels
                     )
 
         if not skip_enqueue:

agentlightning/store/memory.py

@@ -28,6 +28,7 @@ import aiologic
 from pydantic import BaseModel
 
 from agentlightning.types import AttemptedRollout, NamedResources, PaginatedResult, ResourcesUpdate, Rollout, Span
+from agentlightning.utils.metrics import MetricsBackend
 
 from .base import UNSET, LightningStoreCapabilities, LightningStoreStatistics, Unset, is_finished, is_running
 from .collection import InMemoryLightningCollections
@@ -87,13 +88,11 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
         eviction_memory_threshold: float | int | None = None,
         safe_memory_threshold: float | int | None = None,
         span_size_estimator: Callable[[Span], int] | None = None,
-        prometheus: bool = False,
+        tracker: MetricsBackend | None = None,
     ):
         super().__init__(
-            collections=InMemoryLightningCollections(
-                lock_type="thread" if thread_safe else "asyncio", prometheus=prometheus
-            ),
-            prometheus=prometheus,
+            collections=InMemoryLightningCollections(lock_type="thread" if thread_safe else "asyncio", tracker=tracker),
+            tracker=tracker,
         )
 
         self._thread_safe = thread_safe

agentlightning/store/mongo.py

@@ -22,9 +22,10 @@ from typing import (
 from pymongo import AsyncMongoClient
 
 from agentlightning.types import Attempt, AttemptedRollout, Rollout
+from agentlightning.utils.metrics import MetricsBackend
 
 from .base import LightningStoreCapabilities, is_finished
-from .collection.mongo import MongoClientPool, MongoLightningCollections, MongoOperationPrometheusTracker
+from .collection.mongo import MongoClientPool, MongoLightningCollections
 from .collection_based import CollectionBasedLightningStore, healthcheck_before, tracked
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
@@ -54,9 +55,8 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
         client: AsyncMongoClient[Mapping[str, Any]] | str,
         database_name: str | None = None,
         partition_id: str | None = None,
-        prometheus: bool = False,
+        tracker: MetricsBackend | None = None,
     ) -> None:
-        self._enable_prometheus = prometheus
         self._auto_created_client = False
         if isinstance(client, str):
             self._client = AsyncMongoClient[Mapping[str, Any]](client)
@@ -78,9 +78,9 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
                 self._client_pool,
                 database_name,
                 partition_id,
-                prometheus_tracker=MongoOperationPrometheusTracker(enabled=self._enable_prometheus),
+                tracker=tracker,
             ),
-            prometheus=self._enable_prometheus,
+            tracker=tracker,
         )
 
     @property
@@ -136,6 +136,15 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
             await asyncio.sleep(rest_time)
             current_time = time.time()
 
+        # Logging will help debugging when there are stuck rollouts.
+        logger.debug(
+            "Waiting for rollouts. Number of finished rollouts: %d; number of unfinished rollouts: %d",
+            len(finished_rollouts),
+            len(unfinished_rollout_ids),
+        )
+        if len(unfinished_rollout_ids) < 30:
+            logger.debug("Unfinished rollouts: %s", unfinished_rollout_ids)
+
         # Reorder the rollouts to match the input order
         return [finished_rollouts[rollout_id] for rollout_id in rollout_ids if rollout_id in finished_rollouts]
 

agentlightning/utils/metrics.py

@@ -16,16 +16,18 @@ from __future__ import annotations
 import logging
 import os
 import tempfile
-import threading
 import time
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple
 
+import aiologic
+
 if TYPE_CHECKING:
     from prometheus_client import CollectorRegistry
 
 LabelDict = Dict[str, str]
-LabelKey = Tuple[Tuple[str, str], ...]  # normalized, sorted (key, value) pairs
+# Label metadata
+LabelKey = Tuple[Tuple[str, str], ...]  # normalized (key, value) pairs in registration order
 
 logger = logging.getLogger(__name__)
 
@@ -45,7 +47,7 @@ def _validate_labels(
         expected_names: Expected label names as a tuple.
 
     Returns:
-        A tuple of (key, value) pairs sorted by registered label order.
+        A tuple of (key, value) pairs honoring the registered label order.
 
     Raises:
         ValueError: If label keys do not match expected_names.
@@ -67,11 +69,17 @@ def _normalize_label_names(label_names: Optional[Sequence[str]]) -> Tuple[str, .
         label_names: Iterable of label names or None.
 
     Returns:
-        A tuple of label names sorted alphabetically.
+        A tuple of label names preserving their original order.
     """
     if not label_names:
         return ()
-    return tuple(sorted(label_names))
+    return tuple(label_names)
+
+
+def _normalize_prometheus_metric_name(metric_name: str) -> str:
+    """Normalizes Prometheus metric names by replacing unsupported characters."""
+
+    return metric_name.replace(".", "_")
 
 
 @dataclass(frozen=True)
@@ -80,6 +88,7 @@ class _CounterDef:
 
     name: str
     label_names: Tuple[str, ...]
+    group_level: Optional[int] = None
 
 
 @dataclass(frozen=True)
@@ -89,6 +98,7 @@ class _HistogramDef:
     name: str
     label_names: Tuple[str, ...]
     buckets: Tuple[float, ...]
+    group_level: Optional[int] = None
 
 
 @dataclass
@@ -110,16 +120,25 @@ class _HistogramState:
 class MetricsBackend:
     """Abstract base class for metrics backends."""
 
+    def has_prometheus(self) -> bool:
+        """Check if the backend has prometheus support."""
+        return False
+
     def register_counter(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a counter metric.
 
         Args:
             name: Metric name.
-            label_names: List of label names. Order is not important.
+            label_names: List of label names. Order determines the truncation
+                priority for group-level logging.
+            group_level: Optional per-metric grouping depth for backends that
+                support label grouping (Console). Global backend settings take
+                precedence when provided.
 
         Raises:
             ValueError: If the metric is already registered with a different
@@ -132,14 +151,19 @@ class MetricsBackend:
         name: str,
         label_names: Optional[Sequence[str]] = None,
         buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a histogram metric.
 
         Args:
             name: Metric name.
-            label_names: List of label names. Order is not important.
+            label_names: List of label names. Order determines the truncation
+                priority for group-level logging.
             buckets: Bucket boundaries (exclusive upper bounds). If None, the
                 backend may choose defaults.
+            group_level: Optional per-metric grouping depth for backends that
+                support label grouping (Console). Global backend settings take
+                precedence when provided.
 
         Raises:
             ValueError: If the metric is already registered with a different
@@ -147,7 +171,7 @@ class MetricsBackend:
         """
         raise NotImplementedError()
 
-    def inc_counter(
+    async def inc_counter(
         self,
         name: str,
         amount: float = 1.0,
@@ -166,7 +190,7 @@ class MetricsBackend:
         """
         raise NotImplementedError()
 
-    def observe_histogram(
+    async def observe_histogram(
         self,
         name: str,
         value: float,
@@ -199,22 +223,32 @@ class ConsoleMetricsBackend(MetricsBackend):
 
     Rate is always per second.
 
-    Label grouping: When logging, labels are truncated to the first `group_level` label
-    pairs (according to sorted label key order). For example:
+    Label grouping: When logging, label dictionaries are truncated to the first
+    `group_level` label pairs (following the registered label order) and metrics
+    with identical truncated labels are aggregated together. For example:
 
-        labels = {"method": "GET", "path": "/", "status": "200"}
-        group_level = 2 -> logged labels {"method": "GET", "path": "/"}
+    ```python
+    labels = {"method": "GET", "path": "/", "status": "200"}
+    group_level = 2  # aggregated labels {"method": "GET", "path": "/"}
+    ```
 
-    If `group_level` is None or < 1, all labels are logged.
+    If `group_level` is None or < 1, all label combinations for a metric are
+    merged into a single log entry (equivalent to grouping by zero labels).
+    Individual counters or histograms can set their own `group_level` during
+    registration; those values apply only when the backend-level `group_level`
+    is unset, allowing selective overrides.
 
-    Thread-safety: A single lock protects shared state mutation, pruning, and snapshotting.
-    Percentile computation, formatting, and printing are done after releasing the lock.
+    Thread-safety: Runtime updates and snapshotting use two aiologic locks: one for mutating
+    shared state and another that serializes the global logging decision/snapshot capture so
+    other tasks can continue writing. Metric registration happens during initialization,
+    so it is intentionally left lock-free; this assumption is documented here to avoid
+    blocking writes unnecessarily.
     """
 
     def __init__(
         self,
         window_seconds: Optional[float] = 60.0,
-        log_interval_seconds: float = 5.0,
+        log_interval_seconds: float = 10.0,
         group_level: Optional[int] = None,
     ) -> None:
         """Initializes ConsoleMetricsBackend.
@@ -226,8 +260,9 @@ class ConsoleMetricsBackend(MetricsBackend):
                 When the interval elapses, the next metric event triggers a
                 snapshot and logging of all metrics.
             group_level: Label grouping depth. When logging, only the first
-                `group_level` labels (sorted by key) are included. If None or
-                < 1, all labels are included.
+                `group_level` labels (following registered order) are retained and metric
+                events sharing those labels are aggregated. If None or < 1,
+                all label combinations collapse into a single group per metric.
         """
         self.window_seconds = window_seconds
         self.log_interval_seconds = log_interval_seconds
@@ -243,40 +278,42 @@ class ConsoleMetricsBackend(MetricsBackend):
         # Global last log time (for all metrics)
         self._last_log_time: Optional[float] = None
 
-        self._lock = threading.Lock()
+        self._write_lock = aiologic.Lock()
+        self._snapshot_lock = aiologic.Lock()
 
     def register_counter(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a counter metric.
 
         See base class for argument documentation.
         """
         label_tuple = _normalize_label_names(label_names)
-        with self._lock:
-            existing_counter = self._counters.get(name)
-            existing_hist = self._histograms.get(name)
+        existing_counter = self._counters.get(name)
+        existing_hist = self._histograms.get(name)
 
-            if existing_hist is not None:
-                raise ValueError(f"Metric '{name}' already registered as histogram.")
+        if existing_hist is not None:
+            raise ValueError(f"Metric '{name}' already registered as histogram.")
 
-            if existing_counter is not None:
-                if existing_counter.label_names != label_tuple:
-                    raise ValueError(
-                        f"Counter '{name}' already registered with labels "
-                        f"{existing_counter.label_names}, got {label_tuple}."
-                    )
-                return
+        if existing_counter is not None:
+            if existing_counter.label_names != label_tuple:
+                raise ValueError(
+                    f"Counter '{name}' already registered with labels "
+                    f"{existing_counter.label_names}, got {label_tuple}."
+                )
+            return
 
-            self._counters[name] = _CounterDef(name=name, label_names=label_tuple)
+        self._counters[name] = _CounterDef(name=name, label_names=label_tuple, group_level=group_level)
 
     def register_histogram(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
         buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a histogram metric.
 
@@ -288,29 +325,29 @@ class ConsoleMetricsBackend(MetricsBackend):
         else:
             bucket_tuple = tuple(buckets)
 
-        with self._lock:
-            existing_counter = self._counters.get(name)
-            existing_hist = self._histograms.get(name)
+        existing_counter = self._counters.get(name)
+        existing_hist = self._histograms.get(name)
 
-            if existing_counter is not None:
-                raise ValueError(f"Metric '{name}' already registered as counter.")
+        if existing_counter is not None:
+            raise ValueError(f"Metric '{name}' already registered as counter.")
 
-            if existing_hist is not None:
-                if existing_hist.label_names != label_tuple or existing_hist.buckets != bucket_tuple:
-                    raise ValueError(
-                        f"Histogram '{name}' already registered with "
-                        f"labels={existing_hist.label_names}, "
-                        f"buckets={existing_hist.buckets}."
-                    )
-                return
+        if existing_hist is not None:
+            if existing_hist.label_names != label_tuple or existing_hist.buckets != bucket_tuple:
+                raise ValueError(
+                    f"Histogram '{name}' already registered with "
+                    f"labels={existing_hist.label_names}, "
+                    f"buckets={existing_hist.buckets}."
+                )
+            return
 
-            self._histograms[name] = _HistogramDef(
-                name=name,
-                label_names=label_tuple,
-                buckets=bucket_tuple,
-            )
+        self._histograms[name] = _HistogramDef(
+            name=name,
+            label_names=label_tuple,
+            buckets=bucket_tuple,
+            group_level=group_level,
+        )
 
-    def inc_counter(
+    async def inc_counter(
         self,
         name: str,
         amount: float = 1.0,
@@ -330,7 +367,7 @@ class ConsoleMetricsBackend(MetricsBackend):
         label_key = _validate_labels("counter", name, labels, definition.label_names)
         state_key = (name, label_key)
 
-        with self._lock:
+        async with self._write_lock:
             state = self._counter_state.get(state_key)
             if state is None:
                 state = _CounterState(timestamps=[], amounts=[])
@@ -340,18 +377,19 @@ class ConsoleMetricsBackend(MetricsBackend):
             state.amounts.append(amount)
             self._prune_events(state.timestamps, state.amounts, now)
 
-            should_log = self._should_log_locked(now)
-            if should_log:
-                counter_snaps, hist_snaps = self._snapshot_locked(now)
-                snapshot_time = now
-            else:
-                counter_snaps = hist_snaps = []
-                snapshot_time = now
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]] = []
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]] = []
+        should_log = False
+        snapshot_time = now
 
-        if should_log and (counter_snaps or hist_snaps):
+        async with self._snapshot_lock:
+            should_log = self._should_log_locked(now)
+        if should_log:
+            async with self._write_lock:
+                counter_snaps, hist_snaps = self._snapshot_locked(now)
             self._log_snapshot(counter_snaps, hist_snaps, snapshot_time)
 
-    def observe_histogram(
+    async def observe_histogram(
         self,
         name: str,
         value: float,
@@ -371,7 +409,7 @@ class ConsoleMetricsBackend(MetricsBackend):
         label_key = _validate_labels("histogram", name, labels, definition.label_names)
         state_key = (name, label_key)
 
-        with self._lock:
+        async with self._write_lock:
             state = self._hist_state.get(state_key)
             if state is None:
                 state = _HistogramState(timestamps=[], values=[])
@@ -381,15 +419,17 @@ class ConsoleMetricsBackend(MetricsBackend):
             state.values.append(value)
             self._prune_events(state.timestamps, state.values, now)
 
-            should_log = self._should_log_locked(now)
-            if should_log:
-                counter_snaps, hist_snaps = self._snapshot_locked(now)
-                snapshot_time = now
-            else:
-                counter_snaps = hist_snaps = []
-                snapshot_time = now
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]] = []
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]] = []
+        should_log = False
+        snapshot_time = now
 
-        if should_log and (counter_snaps or hist_snaps):
+        async with self._snapshot_lock:
+            should_log = self._should_log_locked(now)
+
+        if should_log:
+            async with self._write_lock:
+                counter_snaps, hist_snaps = self._snapshot_locked(now)
             self._log_snapshot(counter_snaps, hist_snaps, snapshot_time)
 
     def _prune_events(
@@ -490,21 +530,22 @@ class ConsoleMetricsBackend(MetricsBackend):
 
         return counter_snaps, hist_snaps
 
-    def _truncate_labels_for_logging(self, labels: LabelDict) -> LabelDict:
+    def _truncate_labels_for_logging(self, labels: LabelDict, group_level: Optional[int]) -> LabelDict:
         """Returns a label dict truncated to the configured group depth.
 
         Args:
             labels: Original label dictionary.
+            group_level: Effective grouping depth for this metric.
 
         Returns:
             A new dictionary containing at most `group_level` label pairs,
-            chosen by sorted key order. If group_level is None or < 1, returns
-            a shallow copy of the original labels.
+            chosen by registered label order. If group_level is None or < 1,
+            returns an empty dict so that all label combinations collapse together.
         """
-        if self.group_level is None or self.group_level < 1:
-            return dict(labels)
-        items = sorted(labels.items())
-        return dict(items[: self.group_level])
+        if group_level is None or group_level < 1:
+            return {}
+        items = list(labels.items())
+        return dict(items[:group_level])
 
     def _log(self, message: str) -> None:
         """Logs a message via the module logger."""
@@ -523,21 +564,101 @@ class ConsoleMetricsBackend(MetricsBackend):
             hist_snaps: Histogram snapshot list.
         """
         entries: List[str] = []
-        for name, labels, timestamps, amounts in counter_snaps:
-            truncated_labels = self._truncate_labels_for_logging(labels)
-            line = self._log_counter(name, truncated_labels, timestamps, amounts, snapshot_time)
+        for name, labels, timestamps, amounts in self._group_counter_snapshots(counter_snaps):
+            line = self._log_counter(name, labels, timestamps, amounts, snapshot_time)
             if line:
                 entries.append(line)
 
-        for name, labels, values, buckets in hist_snaps:
-            truncated_labels = self._truncate_labels_for_logging(labels)
-            line = self._log_histogram(name, truncated_labels, values, buckets, snapshot_time)
+        for name, labels, values, buckets in self._group_histogram_snapshots(hist_snaps):
+            line = self._log_histogram(name, labels, values, buckets, snapshot_time)
             if line:
                 entries.append(line)
 
         if entries:
+            entries.sort()
             self._log("  ".join(entries))
 
+    def _effective_group_level(self, metric_name: str, *, is_histogram: bool) -> Optional[int]:
+        """Returns the active group level for a metric, honoring per-metric overrides."""
+        if self.group_level is not None:
+            return self.group_level
+        if is_histogram:
+            definition = self._histograms.get(metric_name)
+        else:
+            definition = self._counters.get(metric_name)
+        if definition is None:
+            return None
+        return definition.group_level
+
+    def _group_counter_snapshots(
+        self,
+        counter_snaps: List[Tuple[str, LabelDict, List[float], List[float]]],
+    ) -> List[Tuple[str, LabelDict, List[float], List[float]]]:
+        grouped: Dict[Tuple[str, Tuple[Tuple[str, str], ...]], Dict[str, Any]] = {}
+        for name, labels, timestamps, amounts in counter_snaps:
+            group_level = self._effective_group_level(name, is_histogram=False)
+            truncated_labels = self._truncate_labels_for_logging(labels, group_level)
+            key = (name, tuple(truncated_labels.items()))
+            entry = grouped.setdefault(
+                key,
+                {"name": name, "labels": truncated_labels, "timestamps": [], "amounts": []},
+            )
+            entry["timestamps"].extend(timestamps)
+            entry["amounts"].extend(amounts)
+
+        grouped_snaps: List[Tuple[str, LabelDict, List[float], List[float]]] = []
+        for entry in grouped.values():
+            timestamps = entry["timestamps"]
+            amounts = entry["amounts"]
+            if not timestamps:
+                continue
+            combined = sorted(zip(timestamps, amounts), key=lambda item: item[0])
+            ordered_timestamps = [ts for ts, _ in combined]
+            ordered_amounts = [amt for _, amt in combined]
+            grouped_snaps.append(
+                (
+                    entry["name"],
+                    entry["labels"],
+                    ordered_timestamps,
+                    ordered_amounts,
+                )
+            )
+
+        return grouped_snaps
+
+    def _group_histogram_snapshots(
+        self,
+        hist_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]],
+    ) -> List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]]:
+        grouped: Dict[Tuple[str, Tuple[Tuple[str, str], ...]], Dict[str, Any]] = {}
+        for name, labels, values, buckets in hist_snaps:
+            group_level = self._effective_group_level(name, is_histogram=True)
+            truncated_labels = self._truncate_labels_for_logging(labels, group_level)
+            key = (name, tuple(truncated_labels.items()))
+            entry = grouped.setdefault(
+                key,
+                {"name": name, "labels": truncated_labels, "values": [], "buckets": buckets},
+            )
+            if entry["buckets"] != buckets:
+                raise ValueError(f"Histogram buckets mismatch for metric '{name}'.")
+            entry["values"].extend(values)
+
+        grouped_snaps: List[Tuple[str, LabelDict, List[float], Tuple[float, ...]]] = []
+        for entry in grouped.values():
+            values = entry["values"]
+            if not values:
+                continue
+            grouped_snaps.append(
+                (
+                    entry["name"],
+                    entry["labels"],
+                    list(values),
+                    entry["buckets"],
+                )
+            )
+
+        return grouped_snaps
+
     def _log_counter(
         self,
         name: str,
@@ -598,8 +719,8 @@ class ConsoleMetricsBackend(MetricsBackend):
 
 def _format_label_string(labels: LabelDict) -> str:
     if not labels:
-        return "{}"
-    ordered = ",".join(f"{key}={value}" for key, value in sorted(labels.items()))
+        return ""
+    ordered = ",".join(f"{key}={value}" for key, value in labels.items())
     return f"{{{ordered}}}"
 
 
@@ -641,46 +762,50 @@ class PrometheusMetricsBackend(MetricsBackend):
         self._histograms: Dict[str, _HistogramDef] = {}
         self._prom_counters: Dict[str, Any] = {}
         self._prom_histograms: Dict[str, Any] = {}
+        self._prom_metric_names: Dict[str, str] = {}
 
-        self._lock = threading.Lock()
+    def has_prometheus(self) -> bool:
+        """Check if the backend has prometheus support."""
+        return True
 
     def register_counter(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a Prometheus counter metric."""
         from prometheus_client import Counter as PromCounter
 
         label_tuple = _normalize_label_names(label_names)
 
-        with self._lock:
-            if name in self._histograms:
-                raise ValueError(f"Metric '{name}' already registered as histogram.")
+        if name in self._histograms:
+            raise ValueError(f"Metric '{name}' already registered as histogram.")
 
-            existing = self._counters.get(name)
-            if existing is not None:
-                if existing.label_names != label_tuple:
-                    raise ValueError(
-                        f"Counter '{name}' already registered with labels "
-                        f"{existing.label_names}, got {label_tuple}."
-                    )
-                return
+        existing = self._counters.get(name)
+        if existing is not None:
+            if existing.label_names != label_tuple:
+                raise ValueError(
+                    f"Counter '{name}' already registered with labels " f"{existing.label_names}, got {label_tuple}."
+                )
+            return
 
-            self._counters[name] = _CounterDef(name=name, label_names=label_tuple)
+        prom_name = self._register_prometheus_metric_name(name)
+        self._counters[name] = _CounterDef(name=name, label_names=label_tuple, group_level=group_level)
 
-            prom_counter = PromCounter(
-                name,
-                f"Counter {name}",
-                labelnames=label_tuple,
-            )
-            self._prom_counters[name] = prom_counter
+        prom_counter = PromCounter(
+            prom_name,
+            f"Counter {name}",
+            labelnames=label_tuple,
+        )
+        self._prom_counters[name] = prom_counter
 
     def register_histogram(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
         buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a Prometheus histogram metric."""
         from prometheus_client import Histogram as PromHistogram
@@ -688,43 +813,44 @@ class PrometheusMetricsBackend(MetricsBackend):
         label_tuple = _normalize_label_names(label_names)
         bucket_tuple = tuple(buckets) if buckets is not None else ()
 
-        with self._lock:
-            if name in self._counters:
-                raise ValueError(f"Metric '{name}' already registered as counter.")
+        if name in self._counters:
+            raise ValueError(f"Metric '{name}' already registered as counter.")
 
-            existing = self._histograms.get(name)
-            if existing is not None:
-                if existing.label_names != label_tuple or existing.buckets != bucket_tuple:
-                    raise ValueError(
-                        f"Histogram '{name}' already registered with "
-                        f"labels={existing.label_names}, "
-                        f"buckets={existing.buckets}."
-                    )
-                return
+        existing = self._histograms.get(name)
+        if existing is not None:
+            if existing.label_names != label_tuple or existing.buckets != bucket_tuple:
+                raise ValueError(
+                    f"Histogram '{name}' already registered with "
+                    f"labels={existing.label_names}, "
+                    f"buckets={existing.buckets}."
+                )
+            return
 
-            self._histograms[name] = _HistogramDef(
-                name=name,
-                label_names=label_tuple,
+        prom_name = self._register_prometheus_metric_name(name)
+        self._histograms[name] = _HistogramDef(
+            name=name,
+            label_names=label_tuple,
+            buckets=bucket_tuple,
+            group_level=group_level,
+        )
+
+        if bucket_tuple:
+            prom_hist = PromHistogram(
+                prom_name,
+                f"Histogram {name}",
+                labelnames=label_tuple,
                 buckets=bucket_tuple,
             )
+        else:
+            prom_hist = PromHistogram(
+                prom_name,
+                f"Histogram {name}",
+                labelnames=label_tuple,
+            )
 
-            if bucket_tuple:
-                prom_hist = PromHistogram(
-                    name,
-                    f"Histogram {name}",
-                    labelnames=label_tuple,
-                    buckets=bucket_tuple,
-                )
-            else:
-                prom_hist = PromHistogram(
-                    name,
-                    f"Histogram {name}",
-                    labelnames=label_tuple,
-                )
+        self._prom_histograms[name] = prom_hist
 
-            self._prom_histograms[name] = prom_hist
-
-    def inc_counter(
+    async def inc_counter(
         self,
         name: str,
         amount: float = 1.0,
@@ -743,7 +869,7 @@ class PrometheusMetricsBackend(MetricsBackend):
         else:
             prom_counter.inc(amount)
 
-    def observe_histogram(
+    async def observe_histogram(
         self,
         name: str,
         value: float,
@@ -762,6 +888,19 @@ class PrometheusMetricsBackend(MetricsBackend):
         else:
             prom_hist.observe(value)
 
+    def _register_prometheus_metric_name(self, name: str) -> str:
+        """Registers the normalized Prometheus metric name and ensures uniqueness."""
+
+        normalized = _normalize_prometheus_metric_name(name)
+        existing = self._prom_metric_names.get(normalized)
+        if existing is not None and existing != name:
+            raise ValueError(
+                f"Prometheus metric name conflict: '{name}' normalizes to '{normalized}', "
+                f"which is already used by '{existing}'. Consider renaming one of the metrics."
+            )
+        self._prom_metric_names.setdefault(normalized, name)
+        return normalized
+
 
 class MultiMetricsBackend(MetricsBackend):
     """Metrics backend that forwards calls to multiple underlying backends."""
@@ -779,20 +918,26 @@ class MultiMetricsBackend(MetricsBackend):
             raise ValueError("MultiMetricsBackend requires at least one backend.")
         self._backends = list(backends)
 
+    def has_prometheus(self) -> bool:
+        """Check if the backend has prometheus support."""
+        return any(backend.has_prometheus() for backend in self._backends)
+
     def register_counter(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a counter metric in all underlying backends."""
         for backend in self._backends:
-            backend.register_counter(name, label_names=label_names)
+            backend.register_counter(name, label_names=label_names, group_level=group_level)
 
     def register_histogram(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
         buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
         """Registers a histogram metric in all underlying backends."""
         for backend in self._backends:
@@ -800,9 +945,10 @@ class MultiMetricsBackend(MetricsBackend):
                 name,
                 label_names=label_names,
                 buckets=buckets,
+                group_level=group_level,
             )
 
-    def inc_counter(
+    async def inc_counter(
         self,
         name: str,
         amount: float = 1.0,
@@ -810,9 +956,9 @@ class MultiMetricsBackend(MetricsBackend):
     ) -> None:
         """Increments a counter metric in all underlying backends."""
         for backend in self._backends:
-            backend.inc_counter(name, amount=amount, labels=labels)
+            await backend.inc_counter(name, amount=amount, labels=labels)
 
-    def observe_histogram(
+    async def observe_histogram(
         self,
         name: str,
         value: float,
@@ -820,9 +966,10 @@ class MultiMetricsBackend(MetricsBackend):
     ) -> None:
         """Records a histogram observation in all underlying backends."""
         for backend in self._backends:
-            backend.observe_histogram(name, value=value, labels=labels)
+            await backend.observe_histogram(name, value=value, labels=labels)
 
 
+# This variable should be carried into forked processes
 _prometheus_multiproc_dir: tempfile.TemporaryDirectory[str] | None = None
 
 
@@ -840,7 +987,7 @@ def setup_multiprocess_prometheus():
         logger.debug("Created PROMETHEUS_MULTIPROC_DIR at %s", _prometheus_multiproc_dir.name)
     else:
         logger.warning(
-            "Found PROMETHEUS_MULTIPROC_DIR was set by user. " "This directory must be wiped between multiple runs."
+            "Found PROMETHEUS_MULTIPROC_DIR was set by user. This directory must be wiped between multiple runs."
         )
 
 
@@ -849,7 +996,7 @@ def get_prometheus_registry() -> CollectorRegistry:
     from prometheus_client import REGISTRY, CollectorRegistry, multiprocess
 
     if os.getenv("PROMETHEUS_MULTIPROC_DIR") is not None:
-        logger.debug("Using multiprocess registry for prometheus metrics")
+        logger.info("Using multiprocess registry for prometheus metrics: %s", os.getenv("PROMETHEUS_MULTIPROC_DIR"))
         registry = CollectorRegistry()
         multiprocess.MultiProcessCollector(registry)
         return registry
@@ -857,17 +1004,19 @@ def get_prometheus_registry() -> CollectorRegistry:
     return REGISTRY
 
 
-def shutdown_metrics():
+def shutdown_metrics(server: Any = None, worker: Any = None, *args: Any, **kwargs: Any) -> None:
     """Shutdown prometheus metrics."""
 
-    from prometheus_client import multiprocess
+    if _prometheus_multiproc_dir is not None:
+        from prometheus_client import multiprocess
 
-    path = _prometheus_multiproc_dir
-    if path is None:
-        return
-    try:
-        pid = os.getpid()
-        multiprocess.mark_process_dead(pid, path.name)  # type: ignore
-        logger.debug("Marked Prometheus metrics for process %d as dead", pid)
-    except Exception as e:
-        logger.error("Error during metrics cleanup: %s", str(e))
+        path = _prometheus_multiproc_dir
+        try:
+            if hasattr(worker, "pid"):
+                pid = worker.pid
+            else:
+                pid = os.getpid()
+            multiprocess.mark_process_dead(pid, path.name)  # type: ignore
+            logger.debug("Marked Prometheus metrics for process %d as dead", pid)
+        except Exception as e:
+            logger.error("Error during metrics cleanup: %s", str(e))

agentlightning/utils/server_launcher.py

@@ -940,9 +940,9 @@ class PythonServerLauncher:
                 ),  # Allow half the timeout for graceful shutdown
             }
             if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
-                from prometheus_client import multiprocess
+                from agentlightning.utils.metrics import shutdown_metrics
 
-                options["child_exit"] = lambda server, worker: multiprocess.mark_process_dead(worker.pid)  # type: ignore
+                options["child_exit"] = shutdown_metrics  # type: ignore
 
             self._gunicorn_app = GunicornApp(self.app, options)
 

dashboard/src/components/TracesTable.story.tsx

@@ -3,6 +3,8 @@
 import { useMemo, useState } from 'react';
 import type { Meta, StoryObj } from '@storybook/react';
 import { IconSearch } from '@tabler/icons-react';
+import { waitFor, within } from '@testing-library/dom';
+import userEvent from '@testing-library/user-event';
 import { Box, Stack, TextInput, Title } from '@mantine/core';
 import type { Span } from '@/types';
 import { compareRecords } from '@/utils/table';
@@ -570,4 +572,19 @@ const sequenceSortTestSpans: Span[] = [
 
 export const SequenceIdSortTest: Story = {
   render: () => <TracesTableStoryWrapper maxWidth={1200} spans={sequenceSortTestSpans} />,
+  play: async ({ canvasElement }) => {
+    const canvas = within(canvasElement);
+    const seqHeader = await canvas.findByRole('button', { name: /seq\./i });
+
+    await userEvent.click(seqHeader);
+
+    await waitFor(() => {
+      const rows = canvas.getAllByRole('row');
+      const firstRow = rows[1];
+      if (!firstRow) {
+        throw new Error('Expected at least one data row after sorting by sequence ID');
+      }
+      within(firstRow).getByText('task_sequence_14');
+    });
+  },
 };

docker/compose.prometheus-memory-store.yml

@@ -4,7 +4,7 @@ services:
       file: compose.store.yml
       service: app
 
-    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend memory
+    command: agl store --host 0.0.0.0 --port 4747 --tracker console prometheus --backend memory
 
   node-exporter:
     image: prom/node-exporter:latest

docker/compose.prometheus-mongo-store.yml

@@ -26,13 +26,10 @@ services:
       - /bin/bash
       - -c
       - |
-          mkdir -p /tmp/prometheus &&
           agl store --host 0.0.0.0 --port 4747 \
-            --prometheus --backend mongo \
+            --tracker console prometheus --backend mongo \
             --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 \
             --n-workers ${AGL_STORE_N_WORKERS:-32}
-    environment:
-      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus
 
   mongodb-exporter:
     image: percona/mongodb_exporter:0.47.1

docs/community/contributing.md

@@ -46,6 +46,7 @@ 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.
+- Ship a README that documents smoke-test instructions and includes an "Included Files" section summarizing every file and its role; keep the runnable module self-contained with a module-level docstring explaining CLI usage, plus targeted docstrings or inline comments for educational functions/classes.
 
 !!! warning "Please discuss first"
 
@@ -126,13 +127,13 @@ After `uv sync`, run commands via `uv run ...` (add `--no-sync` once the environ
 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
-uv run pre-commit run --all-files --show-diff-on-failure --color=always
+uv run --no-sync pre-commit install
+uv run --no-sync pre-commit run --all-files --show-diff-on-failure --color=always
 ```
 
 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`
+### 3. Branch from Fresh `main` and Code
 
 Start all work from the latest upstream state:
 
@@ -165,20 +166,28 @@ Use lowercase with hyphens, e.g., `feature/async-runner-hooks`.
 
     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).
 
+Before you start coding, bring the shared coding conventions with you:
+
+- Target `requires-python >= 3.10`, four-space indentation, ~120-character lines (docstrings may run longer), and formatter-owned diffs (Black + isort with the `black` profile).
+- Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
+- Maintain exhaustive type hints (pyright enforces them), write succinct Google-style docstrings (with `[][]` cross-references).
+- Prefer dataclasses or Pydantic models from `agentlightning.types`.
+- Log via `logging.getLogger(__name__)` with targeted DEBUG/INFO/WARNING/ERROR calls—especially for long multi-step functions or broad `try/except` blocks.
+
 ### 4. Test and Validate
 
-Most contributions require automated checks. Prefix commands with `uv run` so they use the project environment.
+Most contributions require automated checks. Once `uv sync` locks dependencies, prefix commands with `uv run --no-sync ...` so they share the same environment as CI.
 
 **Full test suite**
 
 ```bash
-uv run pytest -v
+uv run --no-sync pytest -v
 ```
 
 **Targeted tests**
 
 ```bash
-uv run pytest tests/path/to/test_file.py -k test_name
+uv run --no-sync pytest tests/path/to/test_file.py -k test_name
 ```
 
 **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.
@@ -186,7 +195,7 @@ uv run pytest tests/path/to/test_file.py -k test_name
 **Static analysis:**
 
 ```bash
-uv run pyright
+uv run --no-sync pyright
 ```
 
 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.
@@ -196,8 +205,8 @@ If you have touched code under `examples/`, you should run the example-specific
     Keep API references under [docs/reference]({{ src("docs/reference/") }}) up to date. Doc-only changes should still build cleanly:
 
     ```bash
-    uv run mkdocs serve --strict   # live reload
-    uv run mkdocs build --strict   # CI-equivalent
+    uv run --no-sync mkdocs serve --strict   # live reload
+    uv run --no-sync mkdocs build --strict   # CI-equivalent
     ```
 
     `--strict` elevates warnings to errors so you catch issues before CI.
@@ -205,7 +214,7 @@ If you have touched code under `examples/`, you should run the example-specific
 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).
+- Run `uv run --no-sync pre-commit run --all-files --show-diff-on-failure` (hooks installed via `pre-commit install` run automatically on `git commit`, but rerun them if you amended history).
 - Execute the relevant commands from the test list above.
 - Validate each affected example via its README instructions.
 

docs/reference/internal.md

@@ -56,6 +56,20 @@
 
 ## Utilities
 
+::: agentlightning.utils.metrics.MetricsBackend
+
+::: agentlightning.utils.metrics.ConsoleMetricsBackend
+
+::: agentlightning.utils.metrics.PrometheusMetricsBackend
+
+::: agentlightning.utils.metrics.MultiMetricsBackend
+
+::: agentlightning.utils.metrics.setup_multiprocess_prometheus
+
+::: agentlightning.utils.metrics.get_prometheus_registry
+
+::: agentlightning.utils.metrics.shutdown_metrics
+
 ::: agentlightning.utils.server_launcher.PythonServerLauncher
 
 ::: agentlightning.utils.server_launcher.PythonServerLauncherArgs

examples/README.md

@@ -9,8 +9,8 @@ This catalog highlights the examples shipped with Agent-lightning.
 | [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) |
 | [claude_code](./claude_code) | Claude Code SWE-bench harness that records Agent-lightning traces across Anthropic, vLLM, and OpenAI-compatible backends. | [![claude_code workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-claude-code.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.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 |
+| [rag](./rag) | Retrieval-Augmented Generation pipeline targeting the MuSiQue dataset with Wikipedia retrieval. | [![rag workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-rag.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml) |
+| [search_r1](./search_r1) | Framework-free Search-R1 reinforcement learning training workflow with a retrieval backend. | **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/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) |

examples/minimal/write_metrics.py

@@ -17,6 +17,7 @@ endpoint binds to `0.0.0.0:9105`.
 from __future__ import annotations
 
 import argparse
+import asyncio
 import random
 import signal
 import sys
@@ -43,7 +44,7 @@ def _register_metrics(backend: MetricsBackend) -> None:
     )
 
 
-def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence[str]) -> None:
+async def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence[str]) -> None:
     statuses = ["200", "404", "500"]
     end_time = time.time() + duration
     random.seed(1337)
@@ -51,9 +52,9 @@ def _emit_metrics(backend: MetricsBackend, duration: float, operations: Sequence
         operation = random.choice(operations)
         status = random.choices(statuses, weights=[0.9, 0.05, 0.05], k=1)[0]
         latency = random.lognormvariate(-4.0, 0.5)
-        backend.inc_counter("minimal_requests_total", labels={"operation": operation, "status": status})
-        backend.observe_histogram("minimal_latency_seconds", value=latency, labels={"operation": operation})
-        time.sleep(0.25)
+        await backend.inc_counter("minimal_requests_total", labels={"operation": operation, "status": status})
+        await backend.observe_histogram("minimal_latency_seconds", value=latency, labels={"operation": operation})
+        await asyncio.sleep(0.25)
 
 
 def main() -> None:
@@ -85,7 +86,7 @@ def main() -> None:
 
     original_handler = signal.signal(signal.SIGINT, _handle_interrupt)
     try:
-        _emit_metrics(backend, duration=args.duration, operations=["search", "summary", "answer"])
+        asyncio.run(_emit_metrics(backend, duration=args.duration, operations=["search", "summary", "answer"]))
     finally:
         signal.signal(signal.SIGINT, original_handler)
 

examples/rag/README.md

@@ -1,5 +1,7 @@
 # RAG Agent Example
 
+[![rag workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml)
+
 This example demonstrates training a Retrieval-Augmented Generation (RAG) agent using Agent-Lightning with retrieval capabilities. The agent answers multi-hop questions from a tiny MuSiQue dataset by retrieving and reasoning over Wikipedia passages.
 
 ## Overview

tests/benchmark/analysis.py

@@ -364,9 +364,9 @@ def gather_store_methods(
     peak_window: str,
     subquery_step: str,
 ) -> Tuple[List[StoreMethodStats], StatsSummary]:
-    mean_expr = f"(sum by (method)(increase(collection_store_total[{window}]))) / {window_seconds}"
+    mean_expr = f"(sum by (method)(increase(agl_store_total[{window}]))) / {window_seconds}"
     ops_mean = vector_to_map(safe_vector(client, mean_expr), ("method",))
-    peak_expr = f"sum by (method)(irate(collection_store_total[{peak_window}]))"
+    peak_expr = f"sum by (method)(irate(agl_store_total[{peak_window}]))"
     ops_max = vector_to_map(
         safe_vector(client, f"max_over_time(({peak_expr})[{window}:{subquery_step}])"),
         ("method",),
@@ -378,21 +378,21 @@ def gather_store_methods(
     p50 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.50, sum by (le, method)(rate(collection_store_latency_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.50, sum by (le, method)(rate(agl_store_latency_bucket[{window}])))",
         ),
         ("method",),
     )
     p95 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.95, sum by (le, method)(rate(collection_store_latency_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.95, sum by (le, method)(rate(agl_store_latency_bucket[{window}])))",
         ),
         ("method",),
     )
     p99 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.99, sum by (le, method)(rate(collection_store_latency_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.99, sum by (le, method)(rate(agl_store_latency_bucket[{window}])))",
         ),
         ("method",),
     )
@@ -416,30 +416,30 @@ def gather_store_methods(
     overall: StatsSummary = {
         "ops_mean": safe_scalar(
             client,
-            f"(sum(increase(collection_store_total[{window}]))) / {window_seconds}",
+            f"(sum(increase(agl_store_total[{window}]))) / {window_seconds}",
         )
         or 0.0,
         "ops_max": safe_scalar(
             client,
-            f"max_over_time(((sum(irate(collection_store_total[{peak_window}]))))[{window}:{subquery_step}])",
+            f"max_over_time(((sum(irate(agl_store_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "ops_min": safe_scalar(
             client,
-            f"min_over_time(((sum(irate(collection_store_total[{peak_window}]))))[{window}:{subquery_step}])",
+            f"min_over_time(((sum(irate(agl_store_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "p50": safe_scalar(
             client,
-            f"histogram_quantile(0.50, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.50, sum by (le)(rate(agl_store_latency_bucket[{window}])))",
         )
         or 0.0,
         "p95": safe_scalar(
             client,
-            f"histogram_quantile(0.95, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.95, sum by (le)(rate(agl_store_latency_bucket[{window}])))",
         )
         or 0.0,
         "p99": safe_scalar(
             client,
-            f"histogram_quantile(0.99, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.99, sum by (le)(rate(agl_store_latency_bucket[{window}])))",
         )
         or 0.0,
     }
@@ -454,39 +454,35 @@ def gather_rollout_outcomes(
     rate_map = vector_to_map(
         safe_vector(
             client,
-            f"(sum by (status)(increase(collection_store_rollout_total[{window}]))) / {window_seconds}",
+            f"(sum by (status)(increase(agl_rollouts_total[{window}]))) / {window_seconds}",
         ),
         ("status",),
     )
     p25_map = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.25, "
-            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.25, " f"sum by (le, status)(increase(agl_rollouts_duration_bucket[{window}])))",
         ),
         ("status",),
     )
     p50_map = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.50, "
-            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.50, " f"sum by (le, status)(increase(agl_rollouts_duration_bucket[{window}])))",
         ),
         ("status",),
     )
     p75_map = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.75, "
-            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.75, " f"sum by (le, status)(increase(agl_rollouts_duration_bucket[{window}])))",
         ),
         ("status",),
     )
     max_map = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(1.00, "
-            f"sum by (le, status)(increase(collection_store_rollout_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(1.00, " f"sum by (le, status)(increase(agl_rollouts_duration_bucket[{window}])))",
         ),
         ("status",),
     )
@@ -527,7 +523,7 @@ class HttpPathStats:
 class HttpPathStatusStats:
     method: str
     path: str
-    status_code: str
+    status: str
     qps_mean: float
     qps_max: Optional[float]
     qps_min: Optional[float]
@@ -543,12 +539,12 @@ def gather_http_paths(
     peak_window: str,
     subquery_step: str,
 ) -> Tuple[List[HttpPathStats], StatsSummary]:
-    mean_expr = f"(sum by (method, path)(increase(http_requests_total[{window}]))) / {window_seconds}"
+    mean_expr = f"(sum by (method, path)(increase(agl_http_total[{window}]))) / {window_seconds}"
     qps_mean = vector_to_map(
         safe_vector(client, mean_expr),
         ("method", "path"),
     )
-    peak_expr = f"sum by (method, path)(irate(http_requests_total[{peak_window}]))"
+    peak_expr = f"sum by (method, path)(irate(agl_http_total[{peak_window}]))"
     qps_max = vector_to_map(
         safe_vector(client, f"max_over_time(({peak_expr})[{window}:{subquery_step}])"),
         ("method", "path"),
@@ -560,21 +556,21 @@ def gather_http_paths(
     p50 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.50, sum by (le, method, path)(increase(http_request_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.50, sum by (le, method, path)(increase(agl_http_latency_bucket[{window}])))",
         ),
         ("method", "path"),
     )
     p95 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.95, sum by (le, method, path)(increase(http_request_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.95, sum by (le, method, path)(increase(agl_http_latency_bucket[{window}])))",
         ),
         ("method", "path"),
     )
     p99 = vector_to_map(
         safe_vector(
             client,
-            f"histogram_quantile(0.99, sum by (le, method, path)(increase(http_request_duration_seconds_bucket[{window}])))",
+            f"histogram_quantile(0.99, sum by (le, method, path)(increase(agl_http_latency_bucket[{window}])))",
         ),
         ("method", "path"),
     )
@@ -616,27 +612,27 @@ def gather_http_paths(
     overall: StatsSummary = {
         "qps_mean": safe_scalar(
             client,
-            f"(sum(increase(http_requests_total[{window}]))) / {window_seconds}",
+            f"(sum(increase(agl_http_total[{window}]))) / {window_seconds}",
         )
         or 0.0,
         "qps_max": safe_scalar(
             client,
-            f"max_over_time(((sum(irate(http_requests_total[{peak_window}]))))[{window}:{subquery_step}])",
+            f"max_over_time(((sum(irate(agl_http_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "qps_min": safe_scalar(
             client,
-            f"min_over_time(((sum(irate(http_requests_total[{peak_window}]))))[{window}:{subquery_step}])",
+            f"min_over_time(((sum(irate(agl_http_total[{peak_window}]))))[{window}:{subquery_step}])",
         ),
         "p50": safe_scalar(
-            client, f"histogram_quantile(0.50, sum by (le)(increase(http_request_duration_seconds_bucket[{window}])))"
+            client, f"histogram_quantile(0.50, sum by (le)(increase(agl_http_latency_bucket[{window}])))"
         )
         or 0.0,
         "p95": safe_scalar(
-            client, f"histogram_quantile(0.95, sum by (le)(increase(http_request_duration_seconds_bucket[{window}])))"
+            client, f"histogram_quantile(0.95, sum by (le)(increase(agl_http_latency_bucket[{window}])))"
         )
         or 0.0,
         "p99": safe_scalar(
-            client, f"histogram_quantile(0.99, sum by (le)(increase(http_request_duration_seconds_bucket[{window}])))"
+            client, f"histogram_quantile(0.99, sum by (le)(increase(agl_http_latency_bucket[{window}])))"
         )
         or 0.0,
     }
@@ -650,7 +646,7 @@ def gather_http_paths_with_status(
     peak_window: str,
     subquery_step: str,
 ) -> List[HttpPathStatusStats]:
-    qps_expr = f"sum by (method, path, status_code)(irate(http_requests_total[{peak_window}]))"
+    qps_expr = f"sum by (method, path, status)(irate(agl_http_total[{peak_window}]))"
 
     def fetch_status_metric(expr: str) -> Dict[Tuple[str, str, str], Optional[float]]:
         samples = safe_vector(client, expr)
@@ -666,7 +662,7 @@ def gather_http_paths_with_status(
                 metric_map = {}
             method_raw = metric_map.get("method")
             path_raw = metric_map.get("path")
-            status_raw = metric_map.get("status_code")
+            status_raw = metric_map.get("status")
             key = (
                 _normalize_label(method_raw),
                 _normalize_label(path_raw),
@@ -682,18 +678,18 @@ def gather_http_paths_with_status(
         return text if text else "-"
 
     qps_mean_norm = fetch_status_metric(
-        f"(sum by (method, path, status_code)(increase(http_requests_total[{window}]))) / {window_seconds}"
+        f"(sum by (method, path, status)(increase(agl_http_total[{window}]))) / {window_seconds}"
     )
     qps_max_norm = fetch_status_metric(f"max_over_time(({qps_expr})[{window}:{subquery_step}])")
     qps_min_norm = fetch_status_metric(f"min_over_time(({qps_expr})[{window}:{subquery_step}])")
     p50_norm = fetch_status_metric(
-        f"histogram_quantile(0.50, sum by (le, method, path, status_code)(increase(http_request_duration_seconds_bucket[{window}])))"
+        f"histogram_quantile(0.50, sum by (le, method, path, status)(increase(agl_http_latency_bucket[{window}])))"
     )
     p95_norm = fetch_status_metric(
-        f"histogram_quantile(0.95, sum by (le, method, path, status_code)(increase(http_request_duration_seconds_bucket[{window}])))"
+        f"histogram_quantile(0.95, sum by (le, method, path, status)(increase(agl_http_latency_bucket[{window}])))"
     )
     p99_norm = fetch_status_metric(
-        f"histogram_quantile(0.99, sum by (le, method, path, status_code)(increase(http_request_duration_seconds_bucket[{window}])))"
+        f"histogram_quantile(0.99, sum by (le, method, path, status)(increase(agl_http_latency_bucket[{window}])))"
     )
 
     stats: List[HttpPathStatusStats] = []
@@ -1013,7 +1009,7 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
             [
                 stat.method,
                 stat.path,
-                stat.status_code,
+                stat.status,
                 fmt_rate(stat.qps_mean),
                 fmt_rate(stat.qps_max),
                 fmt_rate(stat.qps_min),

tests/benchmark/benchmark_store.py

@@ -19,7 +19,7 @@ from .utils import flatten_dict, random_dict
 
 console = Console()
 
-MAX_RUNTIME_SECONDS = 45 * 60
+MAX_RUNTIME_SECONDS = 30 * 60
 
 
 def _abort_due_to_timeout() -> None:
@@ -303,6 +303,7 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
 
 def main(argv: Optional[Sequence[str]] = None) -> None:
     args = parse_args(argv)
+    agl.setup_logging()
     store = agl.LightningStoreClient(args.store_url)
     timeout_guard = _start_timeout_guard(MAX_RUNTIME_SECONDS)
     try:

tests/benchmark/micro_benchmark.py

@@ -10,17 +10,32 @@ import multiprocessing
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional, Sequence
+from typing import List, Optional, Sequence
 
 from rich.console import Console
 
 import agentlightning as agl
-from agentlightning.types.tracer import OtelResource, Span, SpanContext, TraceStatus
+from agentlightning.types import EnqueueRolloutRequest, OtelResource, Span, SpanContext, TraceStatus
+from agentlightning.utils.metrics import ConsoleMetricsBackend, MultiMetricsBackend
 from agentlightning.utils.system_snapshot import system_snapshot
 
+from .utils import flatten_dict, random_dict
+
 console = Console()
 
 
+async def _enqueue_rollouts_for_benchmark(store_url: str, *, total_rollouts: int, task_prefix: str) -> None:
+    """Utility that enqueues a fixed number of rollouts for a benchmark."""
+    store = agl.LightningStoreClient(store_url)
+    console.print(f"Enqueuing {total_rollouts} rollouts for {task_prefix} benchmark")
+    try:
+        await store.enqueue_many_rollouts(
+            [EnqueueRolloutRequest(input={"task": f"{task_prefix}-Task-{i}"}) for i in range(total_rollouts)]
+        )
+    finally:
+        await store.close()
+
+
 def _close_store_client(store: agl.LightningStoreClient) -> None:
     try:
         asyncio.run(store.close())
@@ -28,7 +43,7 @@ def _close_store_client(store: agl.LightningStoreClient) -> None:
         pass
 
 
-def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span:
+def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str, attribute_size: int) -> Span:
     trace_hex = f"{sequence_id:032x}"
     span_hex = f"{sequence_id:016x}"
     return Span(
@@ -40,7 +55,14 @@ def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) ->
         parent_id=None,
         name=name,
         status=TraceStatus(status_code="OK"),
-        attributes={},
+        attributes=flatten_dict(
+            random_dict(
+                depth=1,
+                breadth=attribute_size,
+                key_length=(3, 20),
+                value_length=(5, 300),
+            )
+        ),
         events=[],
         links=[],
         start_time=None,
@@ -77,8 +99,8 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
     parser.add_argument("--summary-file", help="File to append final benchmark summary.")
     parser.add_argument(
         "mode",
-        choices=("worker", "dequeue-empty", "rollout"),
-        help="Mode to exercise different operations.",
+        choices=("worker", "dequeue-empty", "dequeue-only", "rollout", "dequeue-update-attempt", "metrics"),
+        help="Mode to exercise different operations (metrics targets MultiMetricsBackend fan-out).",
     )
     args = parser.parse_args(argv)
     return args
@@ -173,6 +195,7 @@ def _rollout_flow_task(args: tuple[str, int, int]) -> bool:
                 attempt_id,
                 task_id * spans_per_attempt + seq,
                 f"micro-span-{seq}",
+                attribute_size=1,
             )
             await store.add_span(span)
         console.print(f"Updating attempt {attempt_id} for task {task_id} with {spans_per_attempt} spans")
@@ -207,6 +230,194 @@ def simulate_rollout_with_spans(store_url: str, spans_per_attempt: int = 4) -> B
     return BenchmarkSummary(mode="rollout", total_tasks=len(task_ids), successes=successes, duration=duration)
 
 
+def _dequeue_only_task(args: tuple[str, str, str]) -> bool:
+    store_url, worker_id, task_id = args
+    console.print(f"[Dequeue-Only Task {task_id}] Dequeueing rollout for worker {worker_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> bool:
+        attempted = await store.dequeue_rollout()  # no worker_id
+        if attempted is None:
+            console.print(f"[Dequeue-Only Task {task_id}] No rollout available to dequeue")
+            return False
+        return True
+
+    try:
+        return asyncio.run(_async_task())
+    except Exception as e:
+        console.print(f"Error dequeueing only worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def dequeue_rollouts(store_url: str) -> BenchmarkSummary:
+    """Benchmark simple dequeues without any additional mutations."""
+    start_time = time.time()
+    total_workers = 512
+    attempts_per_worker = 16
+    total_rollouts = total_workers * attempts_per_worker
+
+    asyncio.run(_enqueue_rollouts_for_benchmark(store_url, total_rollouts=total_rollouts, task_prefix="DequeueOnly"))
+
+    worker_jobs = [
+        (f"Worker-{worker_idx}-Attempt-{attempt_idx}", f"Task-{attempt_idx * total_workers + worker_idx}")
+        for worker_idx in range(total_workers)
+        for attempt_idx in range(attempts_per_worker)
+    ]
+    with multiprocessing.get_context("fork").Pool(processes=total_workers) as pool:
+        successful_tasks = pool.map(
+            _dequeue_only_task, [(store_url, worker_id, task_id) for worker_id, task_id in worker_jobs]
+        )
+
+    async def _query_remaining_rollouts() -> List[str]:
+        store = agl.LightningStoreClient(store_url)
+        try:
+            remaining_rollouts = await store.query_rollouts(status_in=["queuing"])
+            return [item.rollout_id for item in remaining_rollouts]
+        finally:
+            await store.close()
+
+    end_time = time.time()
+    remaining_rollouts = asyncio.run(_query_remaining_rollouts())
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Remaining rollouts: {remaining_rollouts}")
+    console.print(f"Remaining rollouts count: {len(remaining_rollouts)}")
+    console.print(f"Dequeue-only success rate: {successes / len(worker_jobs):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(mode="dequeue-only", total_tasks=len(worker_jobs), successes=successes, duration=duration)
+
+
+def _dequeue_and_update_attempt_task(args: tuple[str, str, str, int]) -> bool:
+    store_url, worker_id, task_id, spans_per_attempt = args
+    console.print(f"Dequeueing and update attempt with worker {worker_id} for task {task_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> bool:
+        console.print(f"[Task {task_id}] Dequeueing rollout")
+        attempted = await store.dequeue_rollout(worker_id=worker_id)
+        if attempted is None:
+            console.print(f"[Task {task_id}] No rollout available to dequeue")
+            return False
+        console.print(f"[Task {task_id}] Retrieving span sequence IDs")
+        sequence_ids = await store.get_many_span_sequence_ids(
+            [(attempted.rollout_id, attempted.attempt.attempt_id) for _ in range(spans_per_attempt)]
+        )
+        if len(sequence_ids) != spans_per_attempt:
+            console.print(
+                f"[Task {task_id}] Unable to retrieve enough span sequence IDs: "
+                f"expected={spans_per_attempt} got={len(sequence_ids)}"
+            )
+            return False
+        console.print(f"[Task {task_id}] Adding {spans_per_attempt} spans")
+        spans = [
+            _make_span(
+                attempted.rollout_id,
+                attempted.attempt.attempt_id,
+                sequence_id,
+                f"micro-span-{sequence_id}",
+                attribute_size=32,
+            )
+            for sequence_id in sequence_ids
+        ]
+        stored_spans = await store.add_many_spans(spans)
+        if len(stored_spans) != len(spans):
+            console.print(
+                f"[Task {task_id}] Only stored {len(stored_spans)}/{len(spans)} spans for "
+                f"rollout_id={attempted.rollout_id} attempt_id={attempted.attempt.attempt_id}"
+            )
+            return False
+        console.print(
+            f"[Task {task_id}] Updating attempt to succeeded: rollout_id={attempted.rollout_id} "
+            f"attempt_id={attempted.attempt.attempt_id}"
+        )
+        await store.update_attempt(attempted.rollout_id, attempted.attempt.attempt_id, status="succeeded")
+        return True
+
+    try:
+        return asyncio.run(_async_task())
+    except Exception as e:
+        console.print(f"Error dequeueing and updating worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def dequeue_and_update_attempts(store_url: str, spans_per_attempt: int = 4) -> BenchmarkSummary:
+    """Simulate dequeueing rollouts and updating attempts with spans."""
+    start_time = time.time()
+    total_workers = 512
+    attempts_per_worker = 16
+    total_rollouts = total_workers * attempts_per_worker
+
+    asyncio.run(_enqueue_rollouts_for_benchmark(store_url, total_rollouts=total_rollouts, task_prefix="Dequeue"))
+
+    worker_jobs = [
+        (f"Worker-{worker_idx}-Attempt-{attempt_idx}", f"Task-{attempt_idx * total_workers + worker_idx}")
+        for worker_idx in range(total_workers)
+        for attempt_idx in range(attempts_per_worker)
+    ]
+    with multiprocessing.get_context("fork").Pool(processes=total_workers) as pool:
+        successful_tasks = pool.map(
+            _dequeue_and_update_attempt_task,
+            [(store_url, worker_id, task_id, spans_per_attempt) for worker_id, task_id in worker_jobs],
+        )
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Dequeue and update attempt success rate: {successes / len(worker_jobs):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(
+        mode="dequeue-update-attempt", total_tasks=len(worker_jobs), successes=successes, duration=duration
+    )
+
+
+def benchmark_multi_metrics_backend(iterations: int = 10_000_000) -> BenchmarkSummary:
+    """Benchmark MultiMetricsBackend fan-out cost."""
+
+    console.print(f"Benchmarking MultiMetricsBackend for {iterations} iterations (2 metric ops per iteration)")
+
+    agl.setup_logging()
+
+    console_backend = ConsoleMetricsBackend(window_seconds=0.5, log_interval_seconds=0.1, group_level=None)
+    console_backend_secondary = ConsoleMetricsBackend(
+        window_seconds=None, log_interval_seconds=1_000_000.0, group_level=None
+    )
+    backend = MultiMetricsBackend([console_backend, console_backend_secondary])
+
+    backend.register_counter("benchmark.metrics.counter", label_names=["worker"])
+    backend.register_histogram(
+        "benchmark.metrics.latency",
+        label_names=["worker"],
+        buckets=(0.001, 0.005, 0.05, 0.5, 1.0),
+    )
+    labels = {"worker": "benchmark"}
+
+    async def _exercise_metrics() -> None:
+        for i in range(iterations):
+            await backend.inc_counter("benchmark.metrics.counter", labels=labels)
+            await backend.observe_histogram(
+                "benchmark.metrics.latency",
+                value=(i % 100) / 100.0,
+                labels=labels,
+            )
+
+    start_time = time.time()
+    asyncio.run(_exercise_metrics())
+    duration = time.time() - start_time
+    total_ops = iterations * 2
+    throughput = total_ops / duration if duration > 0 else 0.0
+
+    console.print(f"Executed {total_ops} metric updates in {duration:.3f}s ({throughput:.1f} ops/s)")
+    return BenchmarkSummary(mode="metrics", total_tasks=total_ops, successes=total_ops, duration=duration)
+
+
 def record_summary(summary: BenchmarkSummary, summary_file: Optional[str]) -> None:
     message = (
         f"[summary] mode={summary.mode} success_rate={summary.success_rate:.3f} "
@@ -227,11 +438,19 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
         summary = simulate_many_update_workers(args.store_url)
     elif args.mode == "dequeue-empty":
         summary = simulate_dequeue_empty_and_update_workers(args.store_url)
+    elif args.mode == "dequeue-only":
+        summary = dequeue_rollouts(args.store_url)
     elif args.mode == "rollout":
         summary = simulate_rollout_with_spans(args.store_url)
+    elif args.mode == "dequeue-update-attempt":
+        summary = dequeue_and_update_attempts(args.store_url)
+    elif args.mode == "metrics":
+        summary = benchmark_multi_metrics_backend()
     else:
         raise ValueError(f"Invalid mode: {args.mode}")
     record_summary(summary, args.summary_file)
+    if summary.success_rate < 1.0:
+        raise ValueError(f"Benchmark failed with success rate {summary.success_rate:.3f}")
 
 
 if __name__ == "__main__":

tests/common/prometheus_stub.py

@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import types
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+
+class _CounterChild:
+    def __init__(self) -> None:
+        self.value = 0.0
+
+    def inc(self, amount: float = 1.0) -> None:
+        self.value += amount
+
+
+class _HistogramChild:
+    def __init__(self) -> None:
+        self.values: List[float] = []
+
+    def observe(self, value: float) -> None:
+        self.values.append(value)
+
+
+class PrometheusStub(types.ModuleType):
+    """Minimal prometheus_client replacement for unit tests."""
+
+    def __init__(self, real_client: Any) -> None:
+        super().__init__("prometheus_client")
+        self.counter_instances: List[_PromCounter] = []
+        self.histogram_instances: List[_PromHistogram] = []
+        self.real_client = real_client
+
+        class CollectorRegistry:
+            pass
+
+        class _Multiprocess:
+            def __init__(self) -> None:
+                self.registry: Optional[CollectorRegistry] = None
+
+            def MultiProcessCollector(self, registry: CollectorRegistry) -> None:
+                self.registry = registry
+
+        self.CollectorRegistry = CollectorRegistry
+        self.REGISTRY = CollectorRegistry()
+        self.multiprocess = _Multiprocess()
+
+        self.Counter = _PromCounterFactory(self)
+        self.Histogram = _PromHistogramFactory(self)
+        self.make_asgi_app = self._make_asgi_app
+
+    def _make_asgi_app(self, *args: Any, **kwargs: Any) -> Any:
+        return self.real_client.make_asgi_app(*args, **kwargs)
+
+
+class _PromCounterFactory:
+    def __init__(self, owner: PrometheusStub) -> None:
+        self._owner = owner
+
+    def __call__(self, name: str, doc: str, labelnames: Sequence[str]) -> _PromCounter:
+        counter = _PromCounter(name, doc, labelnames)
+        counter._register(self._owner.counter_instances)  # pyright: ignore[reportPrivateUsage]
+        return counter
+
+
+class _PromHistogramFactory:
+    def __init__(self, owner: PrometheusStub) -> None:
+        self._owner = owner
+
+    def __call__(
+        self,
+        name: str,
+        doc: str,
+        labelnames: Sequence[str],
+        buckets: Sequence[float] | None = None,
+    ) -> _PromHistogram:
+        histogram = _PromHistogram(name, doc, labelnames, buckets or ())
+        histogram._register(self._owner.histogram_instances)  # pyright: ignore[reportPrivateUsage]
+        return histogram
+
+
+class _PromCounter:
+    def __init__(self, name: str, doc: str, labelnames: Sequence[str]) -> None:
+        self.name = name
+        self.doc = doc
+        self.labelnames = tuple(labelnames)
+        self.default = _CounterChild()
+        self.children: Dict[Tuple[Tuple[str, str], ...], _CounterChild] = {}
+
+    def _register(self, sink: List["_PromCounter"]) -> None:
+        sink.append(self)
+
+    def labels(self, **kwargs: str) -> _CounterChild:
+        key = tuple(sorted(kwargs.items()))
+        return self.children.setdefault(key, _CounterChild())
+
+    def inc(self, amount: float = 1.0) -> None:
+        self.default.inc(amount)
+
+
+class _PromHistogram:
+    def __init__(self, name: str, doc: str, labelnames: Sequence[str], buckets: Sequence[float]) -> None:
+        self.name = name
+        self.doc = doc
+        self.labelnames = tuple(labelnames)
+        self.buckets = tuple(buckets)
+        self.default = _HistogramChild()
+        self.children: Dict[Tuple[Tuple[str, str], ...], _HistogramChild] = {}
+
+    def _register(self, sink: List["_PromHistogram"]) -> None:
+        sink.append(self)
+
+    def labels(self, **kwargs: str) -> _HistogramChild:
+        key = tuple(sorted(kwargs.items()))
+        return self.children.setdefault(key, _HistogramChild())
+
+    def observe(self, value: float) -> None:
+        self.default.observe(value)
+
+
+def make_prometheus_stub() -> PrometheusStub:
+    """Factory helper for tests."""
+    import prometheus_client
+
+    return PrometheusStub(prometheus_client)

tests/store/test_client_server.py

@@ -28,8 +28,16 @@ from agentlightning.types import (
     Span,
     TraceStatus,
 )
+from agentlightning.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+)
 from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncherArgs
 
+from ..common.prometheus_stub import make_prometheus_stub
+
 
 def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span:
     return Span(
@@ -80,6 +88,66 @@ async def server_client(
         await server.stop()
 
 
+async def _exercise_server_metrics_backend(tracker: MetricsBackend) -> None:
+    port = pick_unused_port()
+    store = InMemoryLightningStore(tracker=tracker)
+    server = LightningStoreServer(store, "127.0.0.1", port, tracker=tracker)
+    await server.start()
+    client = LightningStoreClient(server.endpoint)
+    try:
+        await _run_server_side_operations(server)
+        await _run_client_side_operations(client)
+    finally:
+        await client.close()
+        await server.stop()
+
+
+async def _run_server_side_operations(server: LightningStoreServer) -> None:
+    await server.update_resources("metrics-server", {})
+    await server.get_latest_resources()
+    started = await server.start_rollout(input={"origin": "server"}, config=RolloutConfig(timeout_seconds=1.0))
+    queued = await server.enqueue_rollout(input={"origin": "server-queue"})
+    dequeued = await server.dequeue_rollout(worker_id="metrics-server-worker")
+    assert dequeued is not None
+
+    await server.add_span(_make_span(dequeued.rollout_id, dequeued.attempt.attempt_id, 0, "server-span"))
+    await server.update_attempt(queued.rollout_id, dequeued.attempt.attempt_id, status="running")
+    await server.update_attempt(queued.rollout_id, dequeued.attempt.attempt_id, status="succeeded")
+    await server.update_rollout(queued.rollout_id, status="succeeded")
+    await server.wait_for_rollouts(rollout_ids=[queued.rollout_id], timeout=0.1)
+    assert started is not None
+
+
+async def _run_client_side_operations(client: LightningStoreClient) -> None:
+    await client.update_resources("metrics-client", {})
+    await client.get_latest_resources()
+
+    await client.start_rollout(input={"origin": "client"}, mode="train", config=RolloutConfig(timeout_seconds=2.0))
+    queued = await client.enqueue_rollout(
+        input={"origin": "client-queue"}, config=RolloutConfig(unresponsive_seconds=5.0)
+    )
+    dequeued = await client.dequeue_rollout(worker_id="metrics-client-worker")
+    assert dequeued is not None
+
+    span = _make_span(dequeued.rollout_id, dequeued.attempt.attempt_id, 1, "client-span")
+    await client.add_span(span)
+
+    await client.update_attempt(
+        dequeued.rollout_id,
+        dequeued.attempt.attempt_id,
+        status="running",
+        worker_id="metrics-client-worker",
+    )
+    await client.update_attempt(dequeued.rollout_id, dequeued.attempt.attempt_id, status="succeeded")
+    await client.update_rollout(dequeued.rollout_id, status="succeeded")
+
+    await client.wait_for_rollouts(rollout_ids=[dequeued.rollout_id], timeout=0.1)
+    await client.query_rollouts()
+    await client.query_attempts(dequeued.rollout_id)
+    await client.get_worker_by_id("metrics-client-worker")
+    assert queued.rollout_id == dequeued.rollout_id
+
+
 @pytest.mark.asyncio
 async def test_mp_server_does_not_work_with_inmemory_store() -> None:
     store = InMemoryLightningStore()
@@ -200,6 +268,51 @@ async def test_client_start_attempt_propagates_worker_id(
     assert worker.current_attempt_id == retry.attempt.attempt_id
 
 
+@pytest.mark.asyncio
+async def test_console_metrics_backend_tracks_http_and_store_metrics() -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=3600.0, group_level=3)
+    await _exercise_server_metrics_backend(backend)
+
+    counter_metrics = {name for name, _ in backend._counter_state.keys()}  # pyright: ignore[reportPrivateUsage]
+    hist_metrics = {name for name, _ in backend._hist_state.keys()}  # pyright: ignore[reportPrivateUsage]
+    assert "agl.http.total" in counter_metrics
+    assert "agl.store.total" in counter_metrics
+    assert "agl.http.latency" in hist_metrics
+
+
+@pytest.mark.asyncio
+@pytest.mark.prometheus
+async def test_prometheus_metrics_backend_tracks_http_metrics(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+    backend = PrometheusMetricsBackend()
+    await _exercise_server_metrics_backend(backend)
+
+    http_counter = next(inst for inst in stub.counter_instances if inst.name == "agl_http_total")
+    http_histogram = next(inst for inst in stub.histogram_instances if inst.name == "agl_http_latency")
+    assert any(child.value > 0 for child in http_counter.children.values())
+    assert any(child.values for child in http_histogram.children.values())
+
+
+@pytest.mark.asyncio
+@pytest.mark.prometheus
+async def test_multi_metrics_backend_updates_all_children(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+    console_backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=3600.0, group_level=3)
+    prom_backend = PrometheusMetricsBackend()
+    backend = MultiMetricsBackend([console_backend, prom_backend])
+    await _exercise_server_metrics_backend(backend)
+
+    console_counters = {
+        name for name, _ in console_backend._counter_state.keys()  # pyright: ignore[reportPrivateUsage]
+    }
+    assert "agl.http.total" in console_counters
+
+    prom_counter = next(inst for inst in stub.counter_instances if inst.name == "agl_http_total")
+    assert any(child.value > 0 for child in prom_counter.children.values())
+
+
 @pytest.mark.asyncio
 async def test_client_enqueue_many_rollouts_uses_batch_payload(monkeypatch: MonkeyPatch) -> None:
     client = LightningStoreClient("http://localhost:9000")

tests/utils/test_metrics.py

@@ -6,7 +6,6 @@ from __future__ import annotations
 
 import logging
 import sys
-import types
 from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import pytest
@@ -14,6 +13,8 @@ import pytest
 import agentlightning.utils.metrics as metrics_module
 from agentlightning.utils.metrics import ConsoleMetricsBackend, MetricsBackend, MultiMetricsBackend
 
+from ..common.prometheus_stub import make_prometheus_stub
+
 
 def test_validate_labels_reports_missing_label_with_metric_name() -> None:
     labels = {"method": "GET", "status": "200"}
@@ -28,19 +29,19 @@ def test_validate_labels_reports_missing_label_with_metric_name() -> None:
     assert "'status' is required" in message
 
 
-def test_normalize_label_names_is_sorted() -> None:
-    assert metrics_module._normalize_label_names(["b", "a", "b"]) == ("a", "b", "b")
+def test_normalize_label_names_preserves_original_order() -> None:
+    assert metrics_module._normalize_label_names(["b", "a", "b"]) == ("b", "a", "b")
     assert metrics_module._normalize_label_names(None) == ()
 
 
-def test_console_backend_logs_counters_with_sorted_labels() -> None:
-    backend = ConsoleMetricsBackend()
+def test_console_backend_logs_counters_with_registration_order() -> None:
+    backend = ConsoleMetricsBackend(log_interval_seconds=0.0)
     line = backend._log_counter(
         "rate",
-        {"group2": "b", "group1": "a"},
+        {"group1": "a", "group2": "b"},
         timestamps=[0.0, 1.0],
         amounts=[1.0, 1.0],
-        snapshot_time=2.0,
+        snapshot_time=5.0,
     )
 
     assert line == "rate{group1=a,group2=b}=0.40/s"
@@ -50,7 +51,7 @@ def test_console_backend_logs_histograms_with_human_units() -> None:
     backend = ConsoleMetricsBackend()
     line = backend._log_histogram(
         "latency",
-        {"group2": "b", "group1": "a"},
+        {"group1": "a", "group2": "b"},
         values=[0.00395, 0.0168, 3.5],
         buckets=(0.5,),
         snapshot_time=1.0,
@@ -65,9 +66,9 @@ def test_console_backend_logs_histograms_with_human_units() -> None:
 
 
 def test_console_backend_respects_group_level_limit():
-    backend = ConsoleMetricsBackend(group_level=2)
-    truncated = backend._truncate_labels_for_logging({"group3": "c", "group1": "a", "group2": "b"})
-    line = backend._log_counter("metric", truncated, [0.0, 2.0], [1.0, 1.0], snapshot_time=3.0)
+    backend = ConsoleMetricsBackend(group_level=2, log_interval_seconds=0.0)
+    truncated = backend._truncate_labels_for_logging({"group1": "a", "group2": "b", "group3": "c"}, backend.group_level)
+    line = backend._log_counter("metric", truncated, [0.0, 2.0], [1.0, 1.0], snapshot_time=5.0)
 
     assert line == "metric{group1=a,group2=b}=0.40/s"
 
@@ -83,37 +84,44 @@ class _RecordingBackend(MetricsBackend):
     def __init__(self) -> None:
         self.calls: List[Tuple[str, Tuple[Any, ...]]] = []
 
-    def register_counter(self, name: str, label_names: Optional[Sequence[str]] = None) -> None:
-        self.calls.append(("register_counter", (name, tuple(label_names or ()))))
+    def register_counter(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
+    ) -> None:
+        self.calls.append(("register_counter", (name, tuple(label_names or ()), group_level)))
 
     def register_histogram(
         self,
         name: str,
         label_names: Optional[Sequence[str]] = None,
         buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
     ) -> None:
-        self.calls.append(("register_histogram", (name, tuple(label_names or ()), tuple(buckets or ()))))
+        self.calls.append(("register_histogram", (name, tuple(label_names or ()), tuple(buckets or ()), group_level)))
 
-    def inc_counter(self, name: str, amount: float = 1.0, labels: Optional[Dict[str, str]] = None) -> None:
+    async def inc_counter(self, name: str, amount: float = 1.0, labels: Optional[Dict[str, str]] = None) -> None:
         self.calls.append(("inc_counter", (name, amount, labels or {})))
 
-    def observe_histogram(self, name: str, value: float, labels: Optional[Dict[str, str]] = None) -> None:
+    async def observe_histogram(self, name: str, value: float, labels: Optional[Dict[str, str]] = None) -> None:
         self.calls.append(("observe_histogram", (name, value, labels or {})))
 
 
-def test_multi_metrics_backend_fans_out_calls() -> None:
+@pytest.mark.asyncio
+async def test_multi_metrics_backend_fans_out_calls() -> None:
     backend_a = _RecordingBackend()
     backend_b = _RecordingBackend()
     multi = MultiMetricsBackend([backend_a, backend_b])
 
     multi.register_counter("hits", label_names=["method"])
     multi.register_histogram("latency", label_names=["method"], buckets=[0.1, 1.0])
-    multi.inc_counter("hits", amount=2.5, labels={"method": "GET"})
-    multi.observe_histogram("latency", value=0.4, labels={"method": "GET"})
+    await multi.inc_counter("hits", amount=2.5, labels={"method": "GET"})
+    await multi.observe_histogram("latency", value=0.4, labels={"method": "GET"})
 
     expected_calls: List[Tuple[str, Tuple[Any, ...]]] = [
-        ("register_counter", ("hits", ("method",))),
-        ("register_histogram", ("latency", ("method",), (0.1, 1.0))),
+        ("register_counter", ("hits", ("method",), None)),
+        ("register_histogram", ("latency", ("method",), (0.1, 1.0), None)),
         ("inc_counter", ("hits", 2.5, {"method": "GET"})),
         ("observe_histogram", ("latency", 0.4, {"method": "GET"})),
     ]
@@ -122,100 +130,17 @@ def test_multi_metrics_backend_fans_out_calls() -> None:
     assert backend_b.calls == expected_calls
 
 
-class _PrometheusStub(types.ModuleType):
-    def __init__(self) -> None:
-        super().__init__("prometheus_client")
-        self.counter_instances: List["_PromCounter"] = []
-        self.histogram_instances: List["_PromHistogram"] = []
-
-        self_owner = self
-
-        class _CounterChild:
-            def __init__(self) -> None:
-                self.value = 0.0
-
-            def inc(self, amount: float = 1.0) -> None:
-                self.value += amount
-
-        class _HistogramChild:
-            def __init__(self) -> None:
-                self.values: List[float] = []
-
-            def observe(self, value: float) -> None:
-                self.values.append(value)
-
-        class _PromCounter:
-            def __init__(self, name: str, doc: str, labelnames: Sequence[str]) -> None:
-                self.name = name
-                self.doc = doc
-                self.labelnames = tuple(labelnames)
-                self.default = _CounterChild()
-                self.children: Dict[Tuple[Tuple[str, str], ...], _CounterChild] = {}
-                self._register()
-
-            def _register(self) -> None:
-                self_owner.counter_instances.append(self)
-
-            def labels(self, **kwargs: str) -> _CounterChild:
-                key = tuple(sorted(kwargs.items()))
-                return self.children.setdefault(key, _CounterChild())
-
-            def inc(self, amount: float = 1.0) -> None:
-                self.default.inc(amount)
-
-        class _PromHistogram:
-            def __init__(self, name: str, doc: str, labelnames: Sequence[str], buckets: Sequence[float]) -> None:
-                self.name = name
-                self.doc = doc
-                self.labelnames = tuple(labelnames)
-                self.buckets = tuple(buckets)
-                self.default = _HistogramChild()
-                self.children: Dict[Tuple[Tuple[str, str], ...], _HistogramChild] = {}
-                self._register()
-
-            def _register(self) -> None:
-                self_owner.histogram_instances.append(self)
-
-            def labels(self, **kwargs: str) -> _HistogramChild:
-                key = tuple(sorted(kwargs.items()))
-                return self.children.setdefault(key, _HistogramChild())
-
-            def observe(self, value: float) -> None:
-                self.default.observe(value)
-
-        self.Counter = _PromCounter
-        self.Histogram = _PromHistogram
-
-        class CollectorRegistry:
-            pass
-
-        self.CollectorRegistry = CollectorRegistry
-        self.REGISTRY = CollectorRegistry()
-
-        class _Multiprocess:
-            def __init__(self) -> None:
-                self.registry: Optional[CollectorRegistry] = None
-
-            def MultiProcessCollector(self, registry: CollectorRegistry) -> None:
-                self.registry = registry
-
-        self.multiprocess = _Multiprocess()
-
-
-def _make_prometheus_stub() -> _PrometheusStub:
-    return _PrometheusStub()
-
-
-def test_prometheus_backend_binds_stubbed_prometheus(monkeypatch: pytest.MonkeyPatch) -> None:
-    stub = _make_prometheus_stub()
+@pytest.mark.asyncio
+async def test_prometheus_backend_binds_stubbed_prometheus(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
     monkeypatch.setitem(sys.modules, "prometheus_client", stub)
 
     backend = metrics_module.PrometheusMetricsBackend()
     backend.register_counter("hits", ["method"])
     backend.register_histogram("latency", ["method"], buckets=[0.5])
 
-    backend.inc_counter("hits", amount=2.0, labels={"method": "GET"})
-    backend.observe_histogram("latency", value=0.1, labels={"method": "GET"})
+    await backend.inc_counter("hits", amount=2.0, labels={"method": "GET"})
+    await backend.observe_histogram("latency", value=0.1, labels={"method": "GET"})
 
     counter_instance = stub.counter_instances[0]
     histogram_instance = stub.histogram_instances[0]
@@ -223,6 +148,28 @@ def test_prometheus_backend_binds_stubbed_prometheus(monkeypatch: pytest.MonkeyP
     assert histogram_instance.children[(("method", "GET"),)].values == [0.1]
 
 
+def test_prometheus_backend_normalizes_metric_names(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+
+    backend = metrics_module.PrometheusMetricsBackend()
+    backend.register_counter("api.v1.hits")
+    backend.register_histogram("latency.v1")
+
+    assert stub.counter_instances[0].name == "api_v1_hits"
+    assert stub.histogram_instances[0].name == "latency_v1"
+
+
+def test_prometheus_backend_detects_normalized_name_conflicts(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+
+    backend = metrics_module.PrometheusMetricsBackend()
+    backend.register_counter("api.v1.hits")
+    with pytest.raises(ValueError):
+        backend.register_histogram("api_v1_hits")
+
+
 def _split_segments(log_lines: Sequence[str]) -> List[str]:
     segments: List[str] = []
     for line in log_lines:
@@ -230,8 +177,9 @@ def _split_segments(log_lines: Sequence[str]) -> List[str]:
     return [seg for seg in segments if seg]
 
 
-def test_console_backend_sliding_window_rate_and_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
-    backend = ConsoleMetricsBackend(window_seconds=5.0, log_interval_seconds=0.0)
+@pytest.mark.asyncio
+async def test_console_backend_sliding_window_rate_and_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=5.0, log_interval_seconds=0.0, group_level=2)
     backend.register_counter("requests", ["group", "path"])
 
     logged: List[str] = []
@@ -241,9 +189,9 @@ def test_console_backend_sliding_window_rate_and_eviction(monkeypatch: pytest.Mo
     monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
 
     labels = {"group": "api", "path": "/list"}
-    backend.inc_counter("requests", labels=labels)
-    backend.inc_counter("requests", labels=labels)
-    backend.inc_counter("requests", labels=labels)
+    await backend.inc_counter("requests", labels=labels)
+    await backend.inc_counter("requests", labels=labels)
+    await backend.inc_counter("requests", labels=labels)
 
     segments = [seg for seg in _split_segments(logged) if seg.startswith("requests{group=api,path=/list}")]
     assert len(segments) == 3
@@ -266,7 +214,8 @@ def _duration_to_seconds(payload: str) -> float:
     raise AssertionError(f"Unknown duration format: {payload}")
 
 
-def test_console_backend_histogram_quantiles_and_group_depth(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.asyncio
+async def test_console_backend_histogram_quantiles_and_group_depth(monkeypatch: pytest.MonkeyPatch) -> None:
     backend = ConsoleMetricsBackend(window_seconds=10.0, log_interval_seconds=0.0, group_level=2)
     backend.register_histogram("latency", ["service", "endpoint", "status"])
 
@@ -279,13 +228,13 @@ def test_console_backend_histogram_quantiles_and_group_depth(monkeypatch: pytest
     svc_a_labels = {"service": "svcA", "endpoint": "/search", "status": "200"}
     svc_b_labels = {"service": "svcB", "endpoint": "/chat", "status": "500"}
 
-    backend.observe_histogram("latency", value=0.01, labels=svc_a_labels)
-    backend.observe_histogram("latency", value=0.02, labels=svc_a_labels)
-    backend.observe_histogram("latency", value=0.03, labels=svc_a_labels)
-    backend.observe_histogram("latency", value=2.0, labels=svc_b_labels)
+    await backend.observe_histogram("latency", value=0.01, labels=svc_a_labels)
+    await backend.observe_histogram("latency", value=0.02, labels=svc_a_labels)
+    await backend.observe_histogram("latency", value=0.03, labels=svc_a_labels)
+    await backend.observe_histogram("latency", value=2.0, labels=svc_b_labels)
 
     svc_a_segments = [
-        seg for seg in _split_segments(logged) if seg.startswith("latency{endpoint=/search,service=svcA}")
+        seg for seg in _split_segments(logged) if seg.startswith("latency{service=svcA,endpoint=/search}")
     ]
     assert svc_a_segments, "expected log entries for service A"
     latest = svc_a_segments[-1]
@@ -302,7 +251,8 @@ def test_console_backend_histogram_quantiles_and_group_depth(monkeypatch: pytest
     assert abs(p99 - 0.0298) < 5e-3
 
 
-def test_console_backend_logs_all_metric_groups(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.asyncio
+async def test_console_backend_logs_all_metric_groups(monkeypatch: pytest.MonkeyPatch) -> None:
     backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0)
     backend.register_counter("requests", ["group"])
     backend.register_counter("errors", ["group"])
@@ -313,13 +263,13 @@ def test_console_backend_logs_all_metric_groups(monkeypatch: pytest.MonkeyPatch)
     times = iter([0.0, 1.0])
     monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
 
-    backend.inc_counter("requests", labels={"group": "api"})
-    backend.inc_counter("errors", labels={"group": "api"})
+    await backend.inc_counter("requests", labels={"group": "api"})
+    await backend.inc_counter("errors", labels={"group": "api"})
 
     assert logged, "expected log output"
     last_line = logged[-1]
-    assert "requests{group=api}" in last_line
-    assert "errors{group=api}" in last_line
+    assert "requests=" in last_line
+    assert "errors=" in last_line
 
 
 def test_console_backend_snapshot_logs_single_line() -> None:
@@ -337,4 +287,145 @@ def test_console_backend_snapshot_logs_single_line() -> None:
         snapshot_time=1.5,
     )
 
-    assert logged and "counter{g=1}" in logged[0] and "latency{g=1}" in logged[0]
+    assert logged and "counter=" in logged[0] and "latency=" in logged[0]
+
+
+def test_console_backend_snapshot_entries_are_sorted() -> None:
+    backend = ConsoleMetricsBackend(group_level=1)
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("zeta", {"a": "1"}, [0.0], [1.0]),
+            ("alpha", {"b": "2"}, [0.1], [2.0]),
+        ],
+        [],
+        snapshot_time=1.0,
+    )
+
+    assert logged, "expected log output"
+    line = logged[-1]
+    assert line.startswith("alpha{b=2}")
+    assert "  zeta{a=1}" in line
+
+
+def test_console_backend_group_level_none_aggregates_all_label_groups() -> None:
+    backend = ConsoleMetricsBackend(group_level=None)
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("requests", {"group": "api", "path": "/a"}, [0.0], [1.0]),
+            ("requests", {"group": "api", "path": "/b"}, [0.5], [2.0]),
+        ],
+        [
+            ("latency", {"group": "api", "path": "/a"}, [0.1], (0.5,)),
+            ("latency", {"group": "api", "path": "/b"}, [0.2], (0.5,)),
+        ],
+        snapshot_time=1.0,
+    )
+
+    assert logged, "expected log output"
+    line = logged[-1]
+    assert line.count("requests=") == 1
+    assert line.count("latency=") == 1
+
+
+def test_console_backend_group_level_positive_only_aggregates_prefix() -> None:
+    backend = ConsoleMetricsBackend(group_level=1)
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("requests", {"a": "x", "b": "y"}, [0.0], [1.0]),
+            ("requests", {"a": "x", "b": "z"}, [0.1], [2.0]),
+            ("requests", {"a": "w", "b": "y"}, [0.2], [3.0]),
+        ],
+        [],
+        snapshot_time=1.0,
+    )
+
+    assert logged, "expected log output"
+    line = logged[-1]
+    # With group_level=1 we should see aggregation by the first declared label key.
+    assert line.count("requests{a=w}") == 1
+    assert line.count("requests{a=x}") == 1
+
+
+@pytest.mark.asyncio
+async def test_console_backend_logs_preserve_registration_label_order(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0, group_level=3)
+    backend.register_counter("requests", ["path", "method", "status"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.inc_counter("requests", labels={"status": "200", "path": "/search", "method": "GET"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests{path=/search,method=GET,status=200}" in last_line
+
+
+@pytest.mark.asyncio
+async def test_console_backend_metric_specific_group_level_applies(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0)
+    backend.register_counter("requests", ["service", "endpoint", "status"], group_level=2)
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.inc_counter("requests", labels={"status": "200", "endpoint": "/search", "service": "svcA"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests{service=svcA,endpoint=/search}" in last_line
+    assert "status" not in last_line
+
+
+@pytest.mark.asyncio
+async def test_console_backend_global_group_level_overrides_metric(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0, group_level=1)
+    backend.register_counter("requests", ["service", "endpoint"], group_level=2)
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.inc_counter("requests", labels={"service": "svcA", "endpoint": "/search"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests{service=svcA}" in last_line
+    assert "endpoint" not in last_line
+
+
+@pytest.mark.asyncio
+async def test_console_backend_histogram_metric_group_level(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0)
+    backend.register_histogram("latency", ["service", "endpoint"], group_level=1)
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0, 2.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.observe_histogram("latency", value=0.01, labels={"service": "svcA", "endpoint": "/search"})
+    await backend.observe_histogram("latency", value=0.02, labels={"service": "svcA", "endpoint": "/chat"})
+
+    assert logged, "expected log output"
+    latest = logged[-1]
+    assert "latency{service=svcA}" in latest
+    assert "endpoint" not in latest