.

.github/workflows/badge-claude-code.yml

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

.github/workflows/badge-examples.yml

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

.github/workflows/benchmark.yml

@@ -1,3 +1,6 @@
+# This workflow is used to benchmark the performance of the project.
+# It's kept as a placeholder for now.
+
 name: Benchmark
 permissions:
   contents: read
@@ -6,182 +9,11 @@ on:
 
 jobs:
   benchmark:
-    name: Benchmark (${{ matrix.backend.id }}, ${{ matrix.scenario.display }})
-    runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
+    name: Benchmark
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
     timeout-minutes: 60
-    strategy:
-      fail-fast: false
-      matrix:
-        backend:
-          - id: memory
-            compose_file: compose.prometheus-memory-store.yml
-          - id: mongo
-            compose_file: compose.prometheus-mongo-store.yml
-        scenario:
-          - id: minimal-production
-            display: Minimal production scale
-            store_workers: 4
-            args: >-
-              --mode batch
-              --total-tasks 4096
-              --batch-size 256
-              --n-runners 32
-              --max-rounds 6
-              --sleep-seconds 0.5
-          - id: medium-production
-            display: Medium production scale
-            store_workers: 16
-            args: >-
-              --mode batch
-              --total-tasks 10000
-              --batch-size 1000
-              --n-runners 100
-              --max-rounds 10
-              --sleep-seconds 0.1
-          - id: large-batch
-            display: Large batch waves
-            store_workers: 32
-            args: >-
-              --mode batch
-              --total-tasks 100000
-              --batch-size 8192
-              --n-runners 256
-              --max-rounds 6
-              --sleep-seconds 0.1
-          - id: long-queues
-            display: Long rollout queues
-            store_workers: 32
-            args: >-
-              --mode batch_partial
-              --total-tasks 100000
-              --batch-size 1024
-              --n-runners 256
-              --remaining-tasks 4096
-              --max-rounds 4
-              --sleep-seconds 0.1
-          - id: high-concurrency
-            display: High-throughput concurrent requests
-            store_workers: 32
-            args: >-
-              --mode single
-              --total-tasks 100000
-              --concurrency 2048
-              --n-runners 256
-              --max-rounds 2
-              --sleep-seconds 0.1
-          - id: heavy-traces
-            display: Heavy rollouts with deep traces
-            store_workers: 64
-            args: >-
-              --mode batch_partial
-              --total-tasks 10000
-              --batch-size 1024
-              --remaining-tasks 256
-              --n-runners 512
-              --max-rounds 20
-              --sleep-seconds 1.0
-    env:
-      STORE_URL: http://localhost:4747
-      STORE_API_URL: http://localhost:4747/v1/agl
-      PROM_URL: http://localhost:9090
-      SCENARIO_ID: ${{ matrix.scenario.id }}
-      BACKEND_ID: ${{ matrix.backend.id }}
-      ARTIFACT_DIR: artifacts/${{ matrix.scenario.id }}-${{ matrix.backend.id }}
-      COMPOSE_FILE: ${{ matrix.backend.compose_file }}
-      AGL_STORE_N_WORKERS: ${{ matrix.scenario.store_workers }}
     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: Check GPU status
+        run: nvidia-smi
       - name: Check disk space
         run: df -h
-
-      - 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
-          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
-
-      - 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

.github/workflows/examples-claude-code.yml

@@ -1,151 +0,0 @@
-name: Examples - Claude Code
-permissions:
-  contents: read
-on:
-  schedule:
-    # Every day at 4 AM UTC+8
-    - cron: "0 20 * * *"
-
-  workflow_dispatch:
-
-  repository_dispatch:
-    types: [ci-claude-code, ci-all]
-
-run-name: >-
-  ${{ github.event_name == 'repository_dispatch'
-      && format(
-        'Claude Code - PR #{0} - {1} - {2}',
-        github.event.client_payload.pull_number,
-        github.event.client_payload.ci_label,
-        github.event.client_payload.correlation_id
-      )
-      || format('Claude Code - {0}', github.event_name) }}
-
-jobs:
-  claude-code:
-    if: >
-      github.event_name != 'repository_dispatch' ||
-      github.event.action == 'ci-claude-code' ||
-      github.event.action == 'ci-all'
-    name: Claude Code (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
-    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
-    timeout-minutes: 60
-    strategy:
-      matrix:
-        include:
-          - python-version: "3.12"
-            setup-script: "stable"
-          - python-version: "3.13"
-            setup-script: "latest"
-      fail-fast: false
-    steps:
-      - name: Check GPU status
-        run: nvidia-smi
-      - name: Check disk space
-        run: df -h
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
-      - uses: astral-sh/setup-uv@v7
-        with:
-          enable-cache: true
-          python-version: ${{ matrix.python-version }}
-      - name: Upgrade dependencies (latest)
-        run: uv lock --upgrade
-        if: matrix.setup-script == 'latest'
-      - name: Sync dependencies
-        run: |
-          uv sync --frozen --no-default-groups \
-            --group dev --group experiment --group agents --group torch-gpu-stable
-      - name: Freeze dependencies
-        run: |
-          set -ex
-          uv pip freeze | tee requirements-freeze.txt
-          echo "UV_LOCKED=1" >> $GITHUB_ENV
-          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
-      - name: Upload dependencies artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: dependencies-claude-code-${{ matrix.python-version }}-${{ matrix.setup-script }}
-          path: requirements-freeze.txt
-          compression-level: 0
-
-      - name: Download model
-        run: |
-          source .venv/bin/activate
-          python -c "from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-Coder-30B-A3B-Instruct')"
-
-      - name: Launch vLLM server
-        run: |
-          set -euo pipefail
-          source .venv/bin/activate
-          vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
-            --max-model-len 131072 \
-            --enable-auto-tool-choice \
-            --tool-call-parser qwen3_coder \
-            --port 45993 &
-
-          VLLM_READY=0
-          for i in {1..100}; do
-            if curl -sSf http://localhost:45993/v1/models > /dev/null 2>&1; then
-              echo "vLLM server is ready!"
-              VLLM_READY=1
-              break
-            fi
-            echo "Waiting for vLLM server to be ready... (${i})"
-            sleep 5
-          done
-          if [[ "$VLLM_READY" != "1" ]]; then
-            echo "vLLM server failed to start!"
-            exit 1
-          fi
-
-      - name: Claude Code sanity check with vLLM models
-        run: |
-          source .venv/bin/activate
-          cd examples/claude_code
-          python claude_code_agent.py vllm --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct --base-url http://localhost:45993/v1 --debug
-        shell: bash
-
-      - name: Upload sanity check artifacts for vLLM
-        if: ${{ always() }}
-        uses: actions/upload-artifact@v4
-        with:
-          name: claude-code-sanity-check-vllm-${{ matrix.setup-script }}
-          path: |
-            examples/claude_code/data/
-            examples/claude_code/logs/
-          if-no-files-found: error
-
-      - name: Cleanup vLLM
-        run: |
-          set -euo pipefail
-          pkill -f vllm
-          for i in {1..60}; do
-            if ! pgrep -f vllm; then
-              break
-            fi
-            sleep 5
-          done
-          rm -rf examples/claude_code/data/
-          rm -rf examples/claude_code/logs/
-
-      - name: Claude Code sanity check with OpenAI models
-        run: |
-          source .venv/bin/activate
-          cd examples/claude_code
-          python claude_code_agent.py openai --backend-model-high gpt-5.1-codex-mini --backend-model-low gpt-4.1-mini --debug
-        shell: bash
-        env:
-          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
-          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
-
-      - name: Upload sanity check artifacts for OpenAI
-        if: ${{ always() }}
-        uses: actions/upload-artifact@v4
-        with:
-          name: claude-code-sanity-check-openai-${{ matrix.setup-script }}
-          path: |
-            examples/claude_code/data/
-            examples/claude_code/logs/
-          if-no-files-found: error

agentlightning/__init__.py

@@ -7,7 +7,6 @@ from .algorithm import *
 from .client import AgentLightningClient, DevTaskLoader  # deprecated  # type: ignore
 from .config import *
 from .emitter import *
-from .env_var import *
 from .execution import *
 from .litagent import *
 from .llm_proxy import *

agentlightning/adapter/triplet.py

@@ -11,8 +11,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel
 
-from agentlightning.emitter.reward import get_reward_value
-from agentlightning.types import Span, Triplet
+from agentlightning.types import Span, SpanNames, Triplet
 
 from .base import TraceAdapter
 
@@ -314,11 +313,24 @@ class TraceTree:
         Returns:
             Dictionary containing reward metadata, or an empty dictionary when no reward is found.
         """
-        reward_value = get_reward_value(self.span)
-        if reward_value is not None:
-            return {"type": "reward", "value": reward_value}
-        else:
-            return {}
+        for key in [
+            "agentops.task.output",  # newer versions of agentops
+            "agentops.entity.output",
+        ]:
+            output = self.span.attributes.get(key)  # type: ignore
+            if output:
+                if isinstance(output, dict):
+                    return output
+                elif isinstance(output, str):
+                    try:
+                        return json.loads(output)
+                    except json.JSONDecodeError:
+                        return {}
+
+        # Latest emit reward format
+        if self.span.name == SpanNames.REWARD.value and self.span.attributes:
+            return {"type": "reward", "value": self.span.attributes.get("reward", None)}
+        return {}
 
     def is_reward_span(self) -> bool:
         """Return whether the span explicitly encodes a reward.
@@ -764,7 +776,24 @@ class LlmProxyTraceToTriplet(TraceToTripletBase):
 
     def _maybe_reward_value(self, span: Span) -> Optional[float]:
         """Parse reward from typical AgentOps payloads or explicit reward spans."""
-        return get_reward_value(span)
+        attrs = span.attributes or {}
+
+        # AgentOps new/old keys
+        for k in ("agentops.task.output", "agentops.entity.output"):
+            v = attrs.get(k)
+            v = self._literal_eval_maybe(v)
+            if isinstance(v, dict) and cast(Dict[str, Any], v).get("type") == "reward":
+                rv = cast(Dict[str, Any], v).get("value", None)
+                if rv is None or isinstance(rv, (int, float)):
+                    return None if rv is None else float(rv)
+
+        # Explicit reward span
+        if span.name == SpanNames.REWARD.value:
+            rv = attrs.get("reward", None)
+            if rv is None or isinstance(rv, (int, float)):
+                return None if rv is None else float(rv)
+
+        return None
 
     def _request_id_from_attrs(self, attrs: Dict[str, Any]) -> Optional[str]:
         # Prefer OpenAI-like id if present, else proxy raw id.

agentlightning/cli/store.py

@@ -64,13 +64,11 @@ def main(argv: Iterable[str] | None = None) -> int:
     setup_logging(args.log_level)
 
     if args.backend == "memory":
-        store = InMemoryLightningStore(
-            prometheus=args.prometheus, thread_safe=True
-        )  # Using thread_safe store for server
+        store = InMemoryLightningStore()
     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)
     else:
         raise ValueError(f"Invalid backend: {args.backend}")
 

agentlightning/emitter/__init__.py

@@ -1,15 +1,13 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .annotation import emit_annotation
 from .exception import emit_exception
-from .message import emit_message, get_message_value
-from .object import emit_object, get_object_value
+from .message import emit_message
+from .object import emit_object
 from .reward import (
     emit_reward,
     find_final_reward,
     find_reward_spans,
     get_reward_value,
-    get_rewards_from_span,
     is_reward_span,
     reward,
 )
@@ -18,14 +16,10 @@ __all__ = [
     "reward",
     "emit_reward",
     "get_reward_value",
-    "get_rewards_from_span",
     "is_reward_span",
     "find_reward_spans",
     "find_final_reward",
     "emit_message",
     "emit_object",
     "emit_exception",
-    "emit_annotation",
-    "get_message_value",
-    "get_object_value",
 ]

agentlightning/emitter/annotation.py

@@ -1,48 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Helpers for emitting annotation spans."""
-
-import logging
-from typing import Any, Dict
-
-from opentelemetry.sdk.trace import ReadableSpan
-
-from agentlightning.semconv import AGL_ANNOTATION
-from agentlightning.utils.otel import flatten_attributes, get_tracer
-
-logger = logging.getLogger(__name__)
-
-
-def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> ReadableSpan:
-    """Emit a new annotation span.
-
-    This is the underlying implementation of [`emit_reward`][agentlightning.emit_reward].
-
-    Annotation spans are used to annotate a specific event or a part of rollout.
-    See [semconv][agentlightning.semconv] for conventional annotation keys in Agent-lightning.
-
-    If annotations contain nested dicts, they will be flattened before emitting.
-    Complex objects will lead to emitting failures.
-
-    Args:
-        annotation: Dictionary containing annotation key-value pairs.
-            Representatives are rewards, tags, and metadata.
-        propagate: Whether to propagate the span to exporters automatically.
-    """
-    annotation_attributes = flatten_attributes(annotation)
-    if any(not isinstance(v, (str, int, float, bool, bytes)) for v in annotation_attributes.values()):
-        raise TypeError("All annotation attributes must be primitive types (str, int, float, bool, bytes)")
-
-    # TODO: this should use a tracer from current context rather than the singleton
-    tracer = get_tracer(use_active_span_processor=propagate)
-    span = tracer.start_span(
-        AGL_ANNOTATION,
-        attributes=annotation_attributes,
-    )
-    logger.debug("Emitting annotation span with keys %s", annotation_attributes)
-    with span:
-        pass
-    if not isinstance(span, ReadableSpan):
-        raise ValueError(f"Span is not a ReadableSpan: {span}")
-
-    return span

agentlightning/emitter/exception.py

@@ -2,53 +2,43 @@
 
 import logging
 import traceback
-from typing import Any, Dict, Optional
 
 from opentelemetry.semconv.attributes import exception_attributes
 
-from agentlightning.semconv import AGL_EXCEPTION
-from agentlightning.utils.otel import get_tracer
+from agentlightning.types import SpanNames
+
+from .utils import get_tracer
 
 logger = logging.getLogger(__name__)
 
 
-def emit_exception(
-    exception: BaseException, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True
-) -> None:
+def emit_exception(exception: BaseException) -> None:
     """Record an exception with OpenTelemetry metadata.
 
-    Classic OpenTelemetry records exceptions in a dedicated logging service.
-    We simplify the model and use trace spans to record exceptions as well.
-
     Args:
         exception: Raised exception instance to serialize into telemetry attributes.
-        attributes: Additional attributes to attach to the exception span.
-        propagate: Whether to propagate the span to exporters automatically.
 
     !!! note
-
-        The helper validates its input. If a non-exception value is provided,
-        a TypeError is raised to indicate a programming mistake.
+        The helper validates its input. Non-exception values are ignored to prevent
+        noisy telemetry and indicate programming mistakes via the logger.
     """
     if not isinstance(exception, BaseException):  # type: ignore
-        raise TypeError(f"Expected a BaseException instance, got: {type(exception)}.")
+        logger.error(f"Expected an BaseException instance, got: {type(exception)}. Skip emit_exception.")
+        return
 
-    tracer = get_tracer(use_active_span_processor=propagate)
+    tracer = get_tracer()
     stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
-    span_attributes = {
+    attributes = {
         exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
         exception_attributes.EXCEPTION_MESSAGE: str(exception),
         exception_attributes.EXCEPTION_ESCAPED: True,
     }
     if stacktrace.strip():
-        span_attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
-
-    if attributes:
-        span_attributes.update(attributes)
+        attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
 
     span = tracer.start_span(
-        AGL_EXCEPTION,
-        attributes=span_attributes,
+        SpanNames.EXCEPTION.value,
+        attributes=attributes,
     )
     logger.debug("Emitting exception span for %s", type(exception).__name__)
     with span:

agentlightning/emitter/message.py

@@ -1,55 +1,33 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import logging
-from typing import Any, Dict, Optional
 
-from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
-from agentlightning.types import SpanLike
-from agentlightning.utils.otel import get_tracer
+from agentlightning.types import SpanAttributeNames, SpanNames
+
+from .utils import get_tracer
 
 logger = logging.getLogger(__name__)
 
 
-def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
+def emit_message(message: str) -> None:
     """Emit a textual message as an OpenTelemetry span.
 
-    Commonly used for sending debugging and logging messages.
-
     Args:
         message: Human readable message to attach as a span attribute.
-        attributes: Additional attributes to attach to the message span.
-        propagate: Whether to propagate the span to exporters automatically.
 
     !!! note
         OpenTelemetry distinguishes between logs and spans. Emitting the message as a
         span keeps all Agent Lightning telemetry in a single data store for analysis.
     """
     if not isinstance(message, str):  # type: ignore
-        raise TypeError(f"Message must be a string or list of strings, got: {type(message)}.")
+        logger.error(f"Message must be a string, got: {type(message)}. Skip emit_message.")
+        return
 
-    tracer = get_tracer(use_active_span_processor=propagate)
-    span_attributes = {LightningSpanAttributes.MESSAGE_BODY.value: message}
-    if attributes:
-        span_attributes.update(attributes)
+    tracer = get_tracer()
     span = tracer.start_span(
-        AGL_MESSAGE,
-        attributes=span_attributes,
+        SpanNames.MESSAGE.value,
+        attributes={SpanAttributeNames.MESSAGE.value: message},
     )
     logger.debug("Emitting message span with message: %s", message)
     with span:
         pass
-
-
-def get_message_value(span: SpanLike) -> Optional[str]:
-    """Extract the message string from a message span.
-
-    Args:
-        span: Span-like object to extract the message from.
-    """
-    span_attributes = span.attributes or {}
-    if LightningSpanAttributes.MESSAGE_BODY.value not in span_attributes:
-        return None
-    message = span_attributes[LightningSpanAttributes.MESSAGE_BODY.value]
-    if isinstance(message, str):
-        return message
-    raise TypeError(f"Message must be a string, got: {type(message)}.")

agentlightning/emitter/object.py

@@ -1,106 +1,37 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import base64
 import json
 import logging
-from typing import Any, Dict, Optional
+from typing import Any
 
-from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
-from agentlightning.types import SpanLike
-from agentlightning.utils.otel import full_qualified_name, get_tracer
+from agentlightning.types import SpanAttributeNames, SpanNames
+
+from .utils import get_tracer
 
 logger = logging.getLogger(__name__)
 
 
-def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
+def emit_object(object: Any) -> None:
     """Emit an object's serialized representation as an OpenTelemetry span.
 
     Args:
         object: Data structure to encode as JSON and attach to the span payload.
-        attributes: Additional attributes to attach to the object span.
-        propagate: Whether to propagate the span to exporters automatically.
 
     !!! note
-        The payload must be JSON serializable. Non-serializable objects will lead to a RuntimeError.
+        The payload must be JSON serializable. Non-serializable objects are ignored and
+        an error is logged to aid debugging.
     """
-    span_attributes = encode_object(object)
-    if attributes:
-        span_attributes.update(attributes)
-    tracer = get_tracer(use_active_span_processor=propagate)
+    try:
+        serialized = json.dumps(object)
+    except (TypeError, ValueError):
+        logger.error(f"Object must be JSON serializable, got: {type(object)}. Skip emit_object.")
+        return
+
+    tracer = get_tracer()
     span = tracer.start_span(
-        AGL_OBJECT,
-        attributes=span_attributes,
+        SpanNames.OBJECT.value,
+        attributes={SpanAttributeNames.OBJECT.value: serialized},
     )
-    attr_length = 0
-    if LightningSpanAttributes.OBJECT_JSON.value in span_attributes:
-        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_JSON.value])
-    elif LightningSpanAttributes.OBJECT_LITERAL.value in span_attributes:
-        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_LITERAL.value])
-    logger.debug("Emitting object span with payload size %d characters", attr_length)
+    logger.debug("Emitting object span with payload size %d characters", len(serialized))
     with span:
         pass
-
-
-def encode_object(object: Any) -> Dict[str, Any]:
-    """Encode an object as span attributes.
-
-    Args:
-        object: Data structure to encode as JSON.
-    """
-    span_attributes = {}
-    if isinstance(object, (str, int, float, bool)):
-        span_attributes = {
-            LightningSpanAttributes.OBJECT_TYPE.value: type(object).__name__,
-            LightningSpanAttributes.OBJECT_LITERAL.value: str(object),
-        }
-    elif isinstance(object, bytes):
-        b64_encoded = base64.b64encode(object).decode("utf-8")
-        span_attributes = {
-            LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
-            LightningSpanAttributes.OBJECT_LITERAL.value: b64_encoded,
-        }
-    else:
-        try:
-            serialized = json.dumps(object)
-        except (TypeError, ValueError) as exc:
-            raise RuntimeError(f"Object must be JSON serializable, got: {type(object)}.") from exc
-
-        span_attributes = {
-            LightningSpanAttributes.OBJECT_TYPE.value: full_qualified_name(type(object)),  # type: ignore
-            LightningSpanAttributes.OBJECT_JSON.value: serialized,
-        }
-
-    return span_attributes
-
-
-def get_object_value(span: SpanLike) -> Any:
-    """Extract the object payload from an object span.
-
-    Args:
-        span: Span object produced by Agent Lightning emitters.
-    """
-    attributes = span.attributes or {}
-    if LightningSpanAttributes.OBJECT_JSON.value in attributes:
-        serialized = attributes[LightningSpanAttributes.OBJECT_JSON.value]
-        try:
-            return json.loads(serialized)  # type: ignore
-        except (TypeError, ValueError) as exc:
-            raise RuntimeError("Failed to deserialize object JSON from span.") from exc
-    elif LightningSpanAttributes.OBJECT_LITERAL.value in attributes:
-        literal = attributes[LightningSpanAttributes.OBJECT_LITERAL.value]
-        obj_type = attributes.get(LightningSpanAttributes.OBJECT_TYPE.value, "str")
-        if obj_type == "str":
-            return literal
-        elif obj_type == "int":
-            # Let it raise errors if there are any
-            return int(literal)  # type: ignore
-        elif obj_type == "float":
-            return float(literal)  # type: ignore
-        elif obj_type == "bool":
-            return literal.lower() == "true"  # type: ignore
-        elif obj_type == "bytes":
-            return base64.b64decode(literal.encode("utf-8"))  # type: ignore
-        else:
-            raise RuntimeError(f"Unsupported object type for literal deserialization: {obj_type}")
-    else:
-        return None

agentlightning/emitter/reward.py

@@ -23,13 +23,10 @@ from typing import (
 import agentops
 from agentops.sdk.decorators import operation
 from opentelemetry.sdk.trace import ReadableSpan
-from pydantic import TypeAdapter
 
-from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
-from agentlightning.types import SpanLike
-from agentlightning.utils.otel import filter_and_unflatten_attributes
+from agentlightning.types import SpanLike, SpanNames
 
-from .annotation import emit_annotation
+from .utils import get_tracer
 
 logger = logging.getLogger(__name__)
 
@@ -37,26 +34,18 @@ __all__ = [
     "reward",
     "emit_reward",
     "get_reward_value",
-    "get_rewards_from_span",
     "is_reward_span",
     "find_reward_spans",
     "find_final_reward",
 ]
 
 
-class RewardDimension(TypedDict):
-    """Type representing a single dimension in a multi-dimensional reward."""
-
-    name: str
-    value: float
-
-
-class _RewardSpanData(TypedDict):
+class RewardSpanData(TypedDict):
     type: Literal["reward"]
     value: Optional[float]
 
 
-_FnType = TypeVar("_FnType", bound=Callable[..., Any])
+FnType = TypeVar("FnType", bound=Callable[..., Any])
 
 
 def _agentops_initialized() -> bool:
@@ -64,7 +53,7 @@ def _agentops_initialized() -> bool:
     return agentops.get_client().initialized
 
 
-def reward(fn: _FnType) -> _FnType:
+def reward(fn: FnType) -> FnType:
     """Decorate a reward function so its outputs are tracked as spans.
 
     The decorator integrates with AgentOps when it is available and falls back to
@@ -81,7 +70,7 @@ def reward(fn: _FnType) -> _FnType:
         Wrapped callable that preserves the original signature.
     """
 
-    def wrap_result(result: Optional[float]) -> _RewardSpanData:
+    def wrap_result(result: Optional[float]) -> RewardSpanData:
         """Normalize the reward value into the span payload format."""
         if result is None:
             return {"type": "reward", "value": None}
@@ -105,7 +94,7 @@ def reward(fn: _FnType) -> _FnType:
             result: Optional[float] = None
 
             @operation
-            async def agentops_reward_operation() -> _RewardSpanData:
+            async def agentops_reward_operation() -> RewardSpanData:
                 # The reward function we are interested in tracing
                 # It takes zero inputs and return a formatted dict
                 nonlocal result
@@ -129,7 +118,7 @@ def reward(fn: _FnType) -> _FnType:
             result: Optional[float] = None
 
             @operation
-            def agentops_reward_operation() -> _RewardSpanData:
+            def agentops_reward_operation() -> RewardSpanData:
                 nonlocal result
                 result = fn(*args, **kwargs)
                 return wrap_result(result)
@@ -140,36 +129,13 @@ def reward(fn: _FnType) -> _FnType:
         return wrapper  # type: ignore
 
 
-def emit_reward(
-    reward: float | Dict[str, Any],
-    *,
-    primary_key: str | None = None,
-    attributes: Dict[str, Any] | None = None,
-    propagate: bool = True,
-) -> ReadableSpan:
+def emit_reward(reward: float, auto_export: bool = True) -> ReadableSpan:
     """Emit a reward value as an OpenTelemetry span.
 
-    Examples:
-        Emit a single-dimensional reward:
-        >>> emit_reward(1.0)
-
-        Emit multi-dimensional rewards:
-        >>> emit_reward({"task_completion": 1.0, "efficiency": 0.8}, primary_key="task_completion")
-
-        Emit a reward with additional attributes (for example linking to another response span):
-        >>> from agentlightning.utils.otel import make_link_attributes
-        >>> emit_reward(0.5, attributes=make_link_attributes({"gen_ai.response.id": "response-123"}))
-
-        Or adding tags onto the reward span:
-        >>> from agentlightning.utils.otel import make_tag_attributes
-        >>> emit_reward(0.7, attributes=make_tag_attributes(["fast", "reliable"]))
-
     Args:
         reward: Numeric reward to record. Integers and booleans are converted to
             floating point numbers for consistency.
-            Use a dictionary to represent a multi-dimensional reward.
-        attributes: Other optional span attributes.
-        propagate: Whether to propagate the span to exporters automatically.
+        auto_export: Whether to export the span automatically.
 
     Returns:
         Readable span capturing the recorded reward.
@@ -179,34 +145,20 @@ def emit_reward(
             resulting span is not a [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) instance.
     """
     logger.debug(f"Emitting reward: {reward}")
-    reward_dimensions: List[RewardDimension] = []
-    if isinstance(reward, dict):
-        reward_dict: Dict[str, float] = {}
-        for k, v in reward.items():
-            if isinstance(v, (int, bool)):
-                reward_dict[k] = float(v)
-            elif isinstance(v, float):
-                reward_dict[k] = v
-            else:
-                raise ValueError(f"Reward value must be a number, got: {type(v)} for key {k}")
-        if primary_key is None:
-            raise ValueError("When emitting a multi-dimensional reward as a dict, primary_key must be provided.")
-        if primary_key not in reward_dict:
-            raise ValueError(f"Primary key '{primary_key}' not found in reward dict keys: {list(reward_dict.keys())}")
-        reward_dimensions.append(RewardDimension(name=primary_key, value=reward_dict[primary_key]))
-        for k, v in reward_dict.items():
-            if k != primary_key:
-                reward_dimensions.append(RewardDimension(name=k, value=v))
-    else:
-        if isinstance(reward, (int, bool)):
-            reward = float(reward)
-        elif not isinstance(reward, float):  # pyright: ignore[reportUnnecessaryIsInstance]
-            raise TypeError(f"Reward must be a number, got: {type(reward)}")
-        reward_dimensions.append(RewardDimension(name="primary", value=reward))
+    if isinstance(reward, (int, bool)):
+        reward = float(reward)
+    if not isinstance(reward, float):
+        raise ValueError(f"Reward must be a number, got: {type(reward)}")
 
-    return emit_annotation(
-        {LightningSpanAttributes.REWARD.value: reward_dimensions, **(attributes or {})}, propagate=propagate
-    )
+    # TODO: This should use the tracer from current context by tracer
+    tracer = get_tracer(use_active_span_processor=auto_export)
+    span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
+    # Do nothing; it's just a number
+    with span:
+        pass
+    if not isinstance(span, ReadableSpan):
+        raise ValueError(f"Span is not a ReadableSpan: {span}")
+    return span
 
 
 def get_reward_value(span: SpanLike) -> Optional[float]:
@@ -216,14 +168,8 @@ def get_reward_value(span: SpanLike) -> Optional[float]:
         span: Span object produced by AgentOps or Agent Lightning emitters.
 
     Returns:
-        The primary reward encoded in the span or `None` when the span does not represent a reward.
+        The reward encoded in the span or `None` when the span does not represent a reward.
     """
-    # v0.3+ emit reward format
-    reward_list = get_rewards_from_span(span)
-    if reward_list:
-        # Reward list is ordered and the first element is the primary reward
-        return reward_list[0].value
-
     for key in [
         "agentops.task.output",  # newer versions of agentops
         "agentops.entity.output",
@@ -246,45 +192,19 @@ def get_reward_value(span: SpanLike) -> Optional[float]:
                 return None
             if not isinstance(reward_value, float):
                 logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
-            logger.warning(
-                f"Extracted reward {reward_value} from AgentOps. This format is deprecated, please migrate to using `emit_reward`."
-            )
             return cast(float, reward_value)
 
-    # v0.2 emit reward format
-    if span.name == AGL_ANNOTATION and span.attributes:
+    # Latest emit reward format
+    if span.name == SpanNames.REWARD.value and span.attributes:
         reward_value = span.attributes.get("reward", None)
         if reward_value is None:
             return None
         if not isinstance(reward_value, float):
             logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
-        logger.warning(
-            f"Extracted reward {reward_value} from a legacy version of reward span. You might have inconsistent agent-lightning versions."
-        )
         return cast(float, reward_value)
-
     return None
 
 
-def get_rewards_from_span(span: SpanLike) -> List[RewardPydanticModel]:
-    """Extract the reward as a list from a span, if available.
-
-    Args:
-        span: Span object produced by AgentOps or Agent Lightning emitters.
-
-    Returns:
-        A list of reward dimensions encoded in the span or an empty list when the span does not represent a reward.
-    """
-    if span.attributes and any(key.startswith(LightningSpanAttributes.REWARD.value) for key in span.attributes):
-        reward_attr = filter_and_unflatten_attributes(
-            cast(Any, span.attributes or {}), LightningSpanAttributes.REWARD.value
-        )
-        recovered_rewards = TypeAdapter(List[RewardPydanticModel]).validate_python(reward_attr)
-        return recovered_rewards
-    else:
-        return []
-
-
 def is_reward_span(span: SpanLike) -> bool:
     """Return ``True`` when the provided span encodes a reward value."""
     maybe_reward = get_reward_value(span)

agentlightning/emitter/utils.py

@@ -0,0 +1,57 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utilities shared across emitter implementations."""
+
+from typing import cast
+from warnings import filterwarnings
+
+import opentelemetry.trace as trace_api
+from opentelemetry.sdk.trace import SpanLimits, SynchronousMultiSpanProcessor, Tracer
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
+from opentelemetry.trace import get_tracer_provider
+
+
+def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
+    """Resolve the OpenTelemetry tracer configured for Agent Lightning.
+
+    Args:
+        use_active_span_processor: Whether to use the active span processor.
+
+    Returns:
+        OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.
+
+    Raises:
+        RuntimeError: If OpenTelemetry was not initialized before calling this helper.
+    """
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+
+    tracer_provider = cast(TracerProviderImpl, get_tracer_provider())
+
+    if use_active_span_processor:
+        return tracer_provider.get_tracer("agentlightning")
+
+    else:
+        filterwarnings(
+            "ignore",
+            message=r"You should use InstrumentationScope. Deprecated since version 1.11.1.",
+            category=DeprecationWarning,
+            module="opentelemetry.sdk.trace",
+        )
+
+        return Tracer(
+            tracer_provider.sampler,
+            tracer_provider.resource,
+            # We use an empty span processor to avoid emitting spans to the tracer
+            SynchronousMultiSpanProcessor(),
+            tracer_provider.id_generator,
+            InstrumentationInfo("agentlightning", "", ""),  # type: ignore
+            SpanLimits(),
+            InstrumentationScope(
+                "agentlightning",
+                "",
+                "",
+                {},
+            ),
+        )

agentlightning/env_var.py

@@ -1,156 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Environment variable managements."""
-
-from __future__ import annotations
-
-import os
-from enum import Enum
-from typing import overload
-
-__all__ = [
-    "LightningEnvVar",
-    "resolve_bool_env_var",
-    "resolve_int_env_var",
-    "resolve_str_env_var",
-]
-
-
-class LightningEnvVar(Enum):
-    """Environment variables for Agent Lightning."""
-
-    AGL_EMITTER_DEBUG = "AGL_EMITTER_DEBUG"
-    """Enable debug logging for the emitter."""
-
-    AGL_MANAGED_STORE = "AGL_MANAGED_STORE"
-    """If yes, the [`ExecutionStrategy`][agentlightning.ExecutionStrategy]
-    constructs LightningStore wrappers automatically. When `False` the provided
-    `store` is passed directly to the bundles, allowing callers to manage
-    store wrappers manually."""
-
-    AGL_CURRENT_ROLE = "AGL_CURRENT_ROLE"
-    """Which side(s) to run in this process. Used in
-    [`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy]."""
-
-    AGL_SERVER_HOST = "AGL_SERVER_HOST"
-    """Interface the [`LightningStoreServer`][agentlightning.LightningStoreServer]
-    binds to when running the algorithm bundle locally."""
-
-    AGL_SERVER_PORT = "AGL_SERVER_PORT"
-    """Port the [`LightningStoreServer`][agentlightning.LightningStoreServer] listens to."""
-
-
-_TRUTHY_VALUES = {"1", "true", "yes", "on"}
-_FALSY_VALUES = {"0", "false", "no", "off"}
-
-
-@overload
-def resolve_bool_env_var(env_var: LightningEnvVar, override: bool, fallback: bool) -> bool: ...
-
-
-@overload
-def resolve_bool_env_var(env_var: LightningEnvVar, *, fallback: bool) -> bool: ...
-
-
-@overload
-def resolve_bool_env_var(
-    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
-) -> bool | None: ...
-
-
-def resolve_bool_env_var(
-    env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
-) -> bool | None:
-    """Resolve a boolean environment variable.
-
-    Args:
-        env_var: The environment variable to resolve.
-        override: Optional override supplied by the caller.
-        fallback: Default value if the environment variable is not set.
-    """
-
-    if override is not None:
-        return override
-
-    env_value = os.getenv(env_var.value)
-    if env_value is None:
-        return fallback
-
-    normalized = env_value.strip().lower()
-    if normalized in _TRUTHY_VALUES:
-        return True
-    if normalized in _FALSY_VALUES:
-        return False
-
-    raise ValueError(f"{env_var.value} must be one of {_TRUTHY_VALUES} or {_FALSY_VALUES}")
-
-
-@overload
-def resolve_int_env_var(env_var: LightningEnvVar, override: int, fallback: int) -> int: ...
-
-
-@overload
-def resolve_int_env_var(env_var: LightningEnvVar, *, fallback: int) -> int: ...
-
-
-@overload
-def resolve_int_env_var(
-    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
-) -> int | None: ...
-
-
-def resolve_int_env_var(
-    env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
-) -> int | None:
-    """Resolve an integer environment variable.
-
-    Args:
-        env_var: The environment variable to resolve.
-        override: Optional override supplied by the caller.
-        fallback: Default value if the environment variable is not set.
-    """
-    if override is not None:
-        return override
-
-    env_value = os.getenv(env_var.value)
-    if env_value is None:
-        return fallback
-
-    try:
-        return int(env_value)
-    except ValueError:
-        raise ValueError(f"{env_var.value} must be an integer")
-
-
-@overload
-def resolve_str_env_var(env_var: LightningEnvVar, override: str, fallback: str) -> str: ...
-
-
-@overload
-def resolve_str_env_var(env_var: LightningEnvVar, *, fallback: str) -> str: ...
-
-
-@overload
-def resolve_str_env_var(
-    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
-) -> str | None: ...
-
-
-def resolve_str_env_var(
-    env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
-) -> str | None:
-    """Resolve a string environment variable.
-
-    Args:
-        env_var: The environment variable to resolve.
-        override: Optional override supplied by the caller.
-        fallback: Default value if the environment variable is not set.
-    """
-    if override is not None:
-        return override
-
-    env_value = os.getenv(env_var.value)
-    if env_value is None:
-        return fallback
-
-    return env_value

agentlightning/execution/base.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from typing import Protocol
 
 from agentlightning.store.base import LightningStore
@@ -12,6 +13,47 @@ from .events import ExecutionEvent
 logger = logging.getLogger(__name__)
 
 
+_TRUTHY_VALUES = {"1", "true", "yes", "on"}
+_FALSY_VALUES = {"0", "false", "no", "off"}
+
+
+def resolve_managed_store_flag(value: bool | None) -> bool:
+    """Determine whether execution helpers should wrap the provided store.
+
+    The helper first honours an explicit `value`. When `None` it falls back
+    to the `AGL_MANAGED_STORE` environment variable, accepting a variety
+    of truthy and falsy spellings. Missing environment configuration defaults to
+    `True` so that higher-level strategies create the appropriate client or
+    server wrappers automatically.
+
+    Args:
+        value: Optional override supplied by the caller.
+
+    Returns:
+        `True` when a managed store should be created around the provided
+        instance, otherwise `False`.
+
+    Raises:
+        ValueError: If `AGL_MANAGED_STORE` is set to an unsupported
+            value.
+    """
+
+    if value is not None:
+        return value
+
+    env_value = os.getenv("AGL_MANAGED_STORE")
+    if env_value is None:
+        return True
+
+    normalized = env_value.strip().lower()
+    if normalized in _TRUTHY_VALUES:
+        return True
+    if normalized in _FALSY_VALUES:
+        return False
+
+    raise ValueError("AGL_MANAGED_STORE must be one of 1, 0, true, false, yes, no, on, or off")
+
+
 class AlgorithmBundle(Protocol):
     """Callable bundle produced by [`Trainer`][agentlightning.Trainer].
 

agentlightning/execution/client_server.py

@@ -9,11 +9,10 @@ import time
 from multiprocessing.context import BaseContext
 from typing import Callable, Iterable, Literal, cast
 
-from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var, resolve_int_env_var, resolve_str_env_var
 from agentlightning.store.base import LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 
-from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
 from .events import ExecutionEvent, MultiprocessingEvent
 
 logger = logging.getLogger(__name__)
@@ -100,28 +99,44 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
                 By default, runner can exit gracefully with code 0 or terminated
                 by SIGTERM (-15).
         """
-        resolved_role = resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE, override=role, fallback="both")
-        if resolved_role not in ("algorithm", "runner", "both"):
-            raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
-        self.role: Literal["algorithm", "runner", "both"] = resolved_role
+        if role is None:
+            role_env = os.getenv("AGL_CURRENT_ROLE")
+            if role_env is None:
+                # Use both if not specified via env var or argument
+                role = "both"
+            elif role_env not in ("algorithm", "runner", "both"):
+                raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
+            else:
+                role = role_env
+
+        if server_host is None:
+            server_host = os.getenv("AGL_SERVER_HOST", "localhost")
+
+        if server_port is None:
+            server_port_env = os.getenv("AGL_SERVER_PORT")
+            if server_port_env is None:
+                server_port = 4747
+            else:
+                try:
+                    server_port = int(server_port_env)
+                except ValueError as exc:
+                    raise ValueError("AGL_SERVER_PORT must be an integer") from exc
+
+        self.role = role
         self.n_runners = n_runners
-        self.server_host = resolve_str_env_var(
-            LightningEnvVar.AGL_SERVER_HOST, override=server_host, fallback="localhost"
-        )
-        self.server_port = resolve_int_env_var(LightningEnvVar.AGL_SERVER_PORT, override=server_port, fallback=4747)
+        self.server_host = server_host
+        self.server_port = server_port
         self.graceful_timeout = graceful_timeout
         self.terminate_timeout = terminate_timeout
         if main_process not in ("algorithm", "runner"):
             raise ValueError("main_process must be 'algorithm' or 'runner'")
         if main_process == "runner":
-            if self.role != "both":
+            if role != "both":
                 raise ValueError("main_process='runner' is only supported when role='both'")
             if n_runners != 1:
                 raise ValueError("main_process='runner' requires n_runners to be 1")
         self.main_process = main_process
-        self.managed_store = resolve_bool_env_var(
-            LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
-        )
+        self.managed_store = resolve_managed_store_flag(managed_store)
         self.allowed_exit_codes = tuple(allowed_exit_codes)
 
     async def _execute_algorithm(

agentlightning/execution/shared_memory.py

@@ -7,11 +7,10 @@ from contextlib import suppress
 from queue import SimpleQueue
 from typing import Any, Awaitable, Callable, List, Literal, Optional, Tuple
 
-from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
 from agentlightning.store.base import LightningStore
 from agentlightning.store.threading import LightningStoreThreaded
 
-from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle, resolve_managed_store_flag
 from .events import ExecutionEvent, ThreadingEvent
 
 logger = logging.getLogger(__name__)
@@ -63,9 +62,7 @@ class SharedMemoryExecutionStrategy(ExecutionStrategy):
         self.join_timeout = join_timeout
         self.graceful_delay = graceful_delay
         self.poll_interval = poll_interval
-        self.managed_store = resolve_bool_env_var(
-            LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
-        )
+        self.managed_store = resolve_managed_store_flag(managed_store)
 
     async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
         """Run `coro` until it finishes or a cooperative stop is requested.

agentlightning/llm_proxy.py

@@ -47,8 +47,7 @@ from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.types import Scope
 
-from agentlightning.semconv import LightningResourceAttributes
-from agentlightning.types import LLM, ProxyLLM
+from agentlightning.types import LLM, ProxyLLM, SpanNames
 from agentlightning.utils.server_launcher import (
     LaunchMode,
     PythonServerLauncher,
@@ -175,24 +174,6 @@ class AddReturnTokenIds(CustomLogger):
         return {**data, "return_token_ids": True}
 
 
-class AddLogprobs(CustomLogger):
-    """LiteLLM logger hook to request logprobs from vLLM.
-
-    This mutates the outgoing request payload to include `logprobs=1`
-    for backends that support logprobs return (e.g., vLLM).
-    """
-
-    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
-        """Async pre-call hook to adjust request payload."""
-        try:
-            data = _get_pre_call_data(args, kwargs)
-        except Exception as e:
-            return e
-
-        # Ensure logprobs are requested from the backend when supported.
-        return {**data, "logprobs": 1}
-
-
 class LightningSpanExporter(SpanExporter):
     """Buffered OTEL span exporter with subtree flushing and training-store sink.
 
@@ -415,9 +396,9 @@ class LightningSpanExporter(SpanExporter):
                     span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
                         Resource.create(
                             {
-                                LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
-                                LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
-                                LightningResourceAttributes.SPAN_SEQUENCE_ID.value: sequence_id_decimal,
+                                SpanNames.ROLLOUT_ID: rollout_id,
+                                SpanNames.ATTEMPT_ID: attempt_id,
+                                SpanNames.SPAN_SEQUENCE_ID: sequence_id_decimal,
                             }
                         )
                     )
@@ -999,7 +980,6 @@ _MIDDLEWARE_REGISTRY: Dict[str, Type[BaseHTTPMiddleware]] = {
 
 _CALLBACK_REGISTRY = {
     "return_token_ids": AddReturnTokenIds,
-    "logprobs": AddLogprobs,
     "opentelemetry": LightningOpenTelemetry,
 }
 
@@ -1058,7 +1038,7 @@ class LLMProxy:
             Middlewares are the **first layer** of request processing. They are applied to all requests before the LiteLLM proxy.
         callbacks: List of LiteLLM callback classes or strings to register. You can specify the class aliases or classes that have been imported.
             If not provided, the default callbacks (AddReturnTokenIds and LightningOpenTelemetry) will be used.
-            Available callback aliases are: "return_token_ids", "opentelemetry", "logprobs".
+            Available callback aliases are: "return_token_ids", "opentelemetry".
     """
 
     def __init__(
@@ -1072,8 +1052,8 @@ class LLMProxy:
         num_workers: int = 1,
         launch_mode: LaunchMode = "mp",
         launcher_args: PythonServerLauncherArgs | None = None,
-        middlewares: Sequence[Union[Type[BaseHTTPMiddleware], str]] | None = None,
-        callbacks: Sequence[Union[Type[CustomLogger], str]] | None = None,
+        middlewares: List[Union[Type[BaseHTTPMiddleware], str]] | None = None,
+        callbacks: List[Union[Type[CustomLogger], str]] | None = None,
     ):
         self.store = store
 
@@ -1179,9 +1159,6 @@ class LLMProxy:
         if _global_llm_proxy is not None:
             logger.warning("A global LLMProxy is already set. Overwriting it with the new instance.")
 
-        # Patch for LiteLLM v1.80.6+: https://github.com/BerriAI/litellm/issues/17243
-        os.environ["USE_OTEL_LITELLM_REQUEST_SPAN"] = "true"
-
         # Set the global LLMProxy reference for middleware/exporter access.
         set_active_llm_proxy(self)
 

agentlightning/runner/agent.py

@@ -73,7 +73,7 @@ class LitAgentRunner(Runner[T_task]):
         max_rollouts: Optional[int] = None,
         poll_interval: float = 5.0,
         heartbeat_interval: float = 10.0,
-        interval_jitter: float = 0.5,
+        interval_jitter: float = 0.1,
         heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
     ) -> None:
         """Initialize the agent runner.
@@ -287,7 +287,7 @@ class LitAgentRunner(Runner[T_task]):
             # Preserve the existing spans before another span is emitted
             trace_spans = list(self._tracer.get_last_trace())
             # This will NOT emit another span to the tracer
-            reward_span = emit_reward(raw_result, propagate=False)
+            reward_span = emit_reward(raw_result, auto_export=False)
             # We add it to the store manually
             await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
             trace_spans.append(reward_span)
@@ -577,6 +577,16 @@ class LitAgentRunner(Runner[T_task]):
                 if next_rollout is None:
                     return
 
+                try:
+                    # Claim the rollout but updating the current worker id
+                    await store.update_attempt(
+                        next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
+                    )
+                except Exception:
+                    # This exception could happen if the rollout is dequeued and the other end died for some reason
+                    logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
+                    continue
+
                 # Execute the step
                 await self._step_impl(next_rollout)
 
@@ -630,8 +640,12 @@ class LitAgentRunner(Runner[T_task]):
         else:
             resources_id = None
 
-        attempted_rollout = await self.get_store().start_rollout(
-            input=input, mode=mode, resources_id=resources_id, worker_id=self.get_worker_id()
+        attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
+        # Register the attempt as running by the current worker
+        await self.get_store().update_attempt(
+            attempted_rollout.rollout_id,
+            attempted_rollout.attempt.attempt_id,
+            worker_id=self.get_worker_id(),
         )
         rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
 

agentlightning/semconv.py

@@ -1,144 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Semantic conventions for Agent-lightning spans.
-
-Conventions in this file are added on demand. We generally DO NOT add
-new semantic conventions unless it's absolutely needed for certain algorithms or scenarios.
-"""
-
-from enum import Enum
-
-from pydantic import BaseModel
-
-AGL_ANNOTATION = "agentlightning.annotation"
-"""Agent-lightning's standard span name for annotations.
-
-Annotations are minimal span units for rewards, tags, and metadatas.
-They are used to "annotate" a specific event or a part of rollout.
-"""
-
-AGL_MESSAGE = "agentlightning.message"
-"""Agent-lightning's standard span name for messages and logs."""
-
-AGL_OBJECT = "agentlightning.object"
-"""Agent-lightning's standard span name for customized objects."""
-
-AGL_EXCEPTION = "agentlightning.exception"
-"""Agent-lightning's standard span name for exceptions.
-
-Used by the exception emitter to record exception details.
-"""
-
-AGL_VIRTUAL = "agentlightning.virtual"
-"""Agent-lightning's standard span name for virtual operations.
-
-Mostly used in adapter when needing to represent the root or intermediate operations.
-"""
-
-
-class LightningResourceAttributes(Enum):
-    """Resource attribute names used in Agent-lightning spans."""
-
-    ROLLOUT_ID = "agentlightning.rollout_id"
-    """Resource name for rollout ID in Agent-lightning spans."""
-
-    ATTEMPT_ID = "agentlightning.attempt_id"
-    """Resource name for attempt ID in Agent-lightning spans."""
-
-    SPAN_SEQUENCE_ID = "agentlightning.span_sequence_id"
-    """Resource name for span sequence ID in Agent-lightning spans."""
-
-
-class LightningSpanAttributes(Enum):
-    """Attribute names that commonly appear in Agent-lightning spans.
-
-    Exception types can't be found here because they are defined in OpenTelemetry's official semantic conventions.
-    """
-
-    REWARD = "agentlightning.reward"
-    """Attribute prefix for rewards-related data in reward spans.
-
-    It should be used as a prefix. For example, "agentlightning.reward.0.value" can
-    be used to track a specific metric. See [RewardAttributes][agentlightning.semconv.RewardAttributes].
-    """
-
-    LINK = "agentlightning.link"
-    """Attribute name for linking the current span to another span or other objects like requests/responses."""
-
-    TAG = "agentlightning.tag"
-    """Attribute name for tagging spans with customized strings."""
-
-    MESSAGE_BODY = "agentlightning.message.body"
-    """Attribute name for message text in message spans."""
-
-    OBJECT_TYPE = "agentlightning.object.type"
-    """Attribute name for object type (full qualified name) in object spans.
-
-    I think builtin types like str, int, bool, list, dict are self-explanatory and
-    should also be qualified to use here.
-    """
-
-    OBJECT_LITERAL = "agentlightning.object.literal"
-    """Attribute name for object literal value in object spans (for str, int, bool, ...)."""
-
-    OBJECT_JSON = "agentlightning.object.json"
-    """Attribute name for object serialized value (JSON) in object spans."""
-
-
-class RewardAttributes(Enum):
-    """Multi-dimensional reward attributes will look like:
-
-    ```json
-    {"agentlightning.reward.0.name": "efficiency", "agentlightning.reward.0.value": 0.75}
-    ```
-
-    The first reward in the reward list will automatically be the primary reward.
-    If the reward list has greater than 1, it shall be a multi-dimensional case.
-    """
-
-    REWARD_NAME = "name"
-    """Key for each dimension in multi-dimensional reward spans."""
-
-    REWARD_VALUE = "value"
-    """Value for each dimension in multi-dimensional reward spans."""
-
-
-class RewardPydanticModel(BaseModel):
-    """A stricter implementation of RewardAttributes used in otel helpers."""
-
-    name: str
-    """Name of the reward dimension."""
-
-    value: float
-    """Value of the reward dimension."""
-
-
-class LinkAttributes(Enum):
-    """Standard link types used in Agent-lightning spans.
-
-    The link is more powerful than [OpenTelemetry link](https://opentelemetry.io/docs/specs/otel/trace/api/#link)
-    in that it supports linking to a queryset of spans.
-    It can even link to span object that hasn't been emitted yet.
-    """
-
-    KEY_MATCH = "key_match"
-    """Linking to spans with matching attribute keys.
-
-    `trace_id` and `span_id` are reserved and will be used to link to specific spans directly.
-
-    For example, it can be `gen_ai.response.id` if intended to be link to a chat completion response span.
-    Or it can be `span_id` to link to a specific span by its ID.
-    """
-
-    VALUE_MATCH = "value_match"
-    """Linking to spans with corresponding attribute values on those keys."""
-
-
-class LinkPydanticModel(BaseModel):
-    """A stricter implementation of LinkAttributes used in otel helpers."""
-
-    key_match: str
-    """The attribute key to match on the target spans."""
-
-    value_match: str
-    """The attribute value to match on the target spans."""

agentlightning/store/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import LightningStore, LightningStoreCapabilities, LightningStoreStatistics
+from .base import LightningStore, LightningStoreCapabilities
 from .client_server import LightningStoreClient, LightningStoreServer
 from .collection_based import CollectionBasedLightningStore
 from .memory import InMemoryLightningStore
@@ -9,7 +9,6 @@ from .threading import LightningStoreThreaded
 __all__ = [
     "LightningStore",
     "LightningStoreCapabilities",
-    "LightningStoreStatistics",
     "LightningStoreClient",
     "LightningStoreServer",
     "InMemoryLightningStore",

agentlightning/store/base.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, TypedDict
+from typing import Any, Dict, List, Literal, Optional, Sequence, TypedDict
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -10,12 +10,10 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
-    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
     RolloutConfig,
-    RolloutMode,
     RolloutStatus,
     Span,
     TaskInput,
@@ -72,35 +70,6 @@ class LightningStoreCapabilities(TypedDict, total=False):
     """Whether the store supports OTLP/HTTP traces."""
 
 
-class LightningStoreStatistics(TypedDict, total=False):
-    """Statistics of a LightningStore implementation."""
-
-    name: str
-    """Name of the store implementation."""
-    total_rollouts: int
-    """Total number of rollouts in the store."""
-    total_attempts: int
-    """Total number of attempts in the store."""
-    total_spans: int
-    """Total number of spans in the store."""
-    total_resources: int
-    """Total number of resources in the store."""
-    total_workers: int
-    """Total number of workers in the store."""
-    uptime: float
-    """Uptime of since the store has been started."""
-
-    # Memory-related statistics
-    total_span_bytes: int
-    """Total number of bytes of spans in the store."""
-    eviction_threshold_bytes: int
-    """Eviction threshold for spans in bytes."""
-    safe_threshold_bytes: int
-    """Safe threshold for spans in bytes."""
-    memory_capacity_bytes: int
-    """Memory capacity of the store in bytes."""
-
-
 class LightningStore:
     """Contract for the persistent control-plane that coordinates training rollouts.
 
@@ -133,12 +102,6 @@ class LightningStore:
             otlp_traces=False,
         )
 
-    async def statistics(self) -> LightningStoreStatistics:
-        """Return the statistics of the store."""
-        return {
-            "name": self.__class__.__name__,
-        }
-
     def otlp_traces_endpoint(self) -> str:
         """Return the OTLP/HTTP traces endpoint of the store.
 
@@ -158,11 +121,10 @@ class LightningStore:
     async def start_rollout(
         self,
         input: TaskInput,
-        mode: RolloutMode | None = None,
+        mode: Literal["train", "val", "test"] | None = None,
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
-        worker_id: str | None = None,
     ) -> AttemptedRollout:
         """Register a rollout and immediately create its first attempt.
 
@@ -185,7 +147,6 @@ class LightningStore:
             resources_id: Concrete resource snapshot to execute against; defaults to the latest stored snapshot.
             config: Rollout retry/timeout policy. Should default to a fresh [`RolloutConfig`][agentlightning.RolloutConfig].
             metadata: Free-form metadata persisted verbatim with the rollout.
-            worker_id: Optional worker identifier to associate the new attempt with.
 
         Returns:
             The fully-populated [`AttemptedRollout`][agentlightning.AttemptedRollout] including
@@ -231,22 +192,6 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
-        """Persist multiple rollouts in `queuing` state.
-
-        The implementation can delegate to [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]
-        per request and preserves the input ordering. Subclasses can override to provide
-        more efficient bulk enqueue semantics.
-
-        Args:
-            rollouts: Rollout submission payloads mirroring [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]'s
-                parameters. Each entry requires `input` and can optionally include other fields.
-
-        Returns:
-            Rollouts enqueued in the same order as `rollouts`.
-        """
-        raise NotImplementedError()
-
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """Claim the oldest queued rollout and transition it to `preparing`.
 
@@ -263,9 +208,6 @@ class LightningStore:
         * Optionally refresh the caller's [`Worker`][agentlightning.Worker] telemetry
           (e.g., `last_dequeue_time`) when `worker_id` is provided.
 
-        Args:
-            worker_id: Optional worker identifier to associate the claimed attempt with.
-
         Returns:
             The next attempt to execute, or `None` when no eligible rollouts are queued.
 
@@ -274,30 +216,7 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def dequeue_many_rollouts(
-        self,
-        *,
-        limit: int = 1,
-        worker_id: Optional[str] = None,
-    ) -> Sequence[AttemptedRollout]:
-        """Claim up to `limit` queued rollouts without blocking.
-
-        The implementation can repeatedly invokes
-        [`dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] until reaching
-        the requested limit or the queue is empty. Subclasses can override it to fetch
-        multiple rollouts atomically.
-
-        Args:
-            limit: Maximum number of rollouts to claim. Non-positive values return an empty list.
-            worker_id: Optional worker identifier passed through to each dequeue call.
-
-        Returns:
-            Attempted rollouts claimed in FIFO order. May contain fewer than `limit` entries
-            when the queue is exhausted.
-        """
-        raise NotImplementedError()
-
-    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
         """Create a manual retry attempt for an existing rollout.
 
         This is typically invoked by runners that wish to retry outside of the
@@ -308,7 +227,6 @@ class LightningStore:
 
         Args:
             rollout_id: Unique identifier of the rollout receiving a new attempt.
-            worker_id: Optional worker identifier to associate the new attempt with.
 
         Returns:
             The rollout paired with its newly-created attempt.
@@ -319,15 +237,7 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
-        """Persist a sequence of pre-constructed spans emitted during rollout execution.
-
-        Implementations can simply delegate to [`add_span()`][agentlightning.LightningStore.add_span] for each span.
-        However, if the store supports bulk insertion, it can implement this method to improve performance.
-        """
-        raise NotImplementedError()
-
-    async def add_span(self, span: Span) -> Optional[Span]:
+    async def add_span(self, span: Span) -> Span:
         """Persist a pre-constructed span emitted during rollout execution.
 
         The provided [`Span`][agentlightning.Span] must already contain the `rollout_id`,
@@ -344,7 +254,6 @@ class LightningStore:
 
         Returns:
             The stored span record (implementations may return a copy).
-            Return `None` if the span was not added due to a duplicate.
 
         Raises:
             NotImplementedError: Subclasses must implement span persistence.
@@ -358,7 +267,7 @@ class LightningStore:
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Optional[Span]:
+    ) -> Span:
         """Convert and persist an OpenTelemetry span for a particular attempt.
 
         Implementations must transform the `readable_span` into a [`Span`][agentlightning.Span]
@@ -375,7 +284,7 @@ class LightningStore:
                 automatically.
 
         Returns:
-            The stored span record. Return `None` if the span was not added due to a duplicate.
+            The stored span record.
 
         Raises:
             NotImplementedError: Subclasses must implement span persistence.
@@ -576,20 +485,6 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
-        """Bulk allocate the next strictly increasing sequence number used to order spans.
-
-        Implementations may delegate to [`get_next_span_sequence_id()`][agentlightning.LightningStore.get_next_span_sequence_id]
-        for each rollout and attempt.
-
-        Args:
-            rollout_attempt_ids: List of tuples of rollout and attempt identifiers.
-
-        Returns:
-            List of sequence numbers.
-        """
-        raise NotImplementedError()
-
     async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         """Block until the targeted rollouts reach a terminal status or the timeout expires.
 
@@ -766,8 +661,7 @@ class LightningStore:
 
         When `attempt_id` is `"latest"` the update must target the attempt with the highest
         `sequence_id`; otherwise it must target the specific attempt. Implementations should
-        propagate status changes to the rollout (for example
-        via [`rollout_status_from_attempt()`][agentlightning.store.utils.rollout_status_from_attempt])
+        propagate status changes to the rollout (for example via [`propagate_status()`][agentlightning.store.utils.propagate_status])
         once the latest attempt transitions to a terminal state.
 
         Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],

agentlightning/store/client_server.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import asyncio
 import logging
 import os
-import re
 import threading
 import time
 import traceback
@@ -24,7 +23,6 @@ from typing import (
     Type,
     TypeVar,
     Union,
-    cast,
 )
 
 import aiohttp
@@ -47,7 +45,6 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
-    EnqueueRolloutRequest,
     NamedResources,
     PaginatedResult,
     ResourcesUpdate,
@@ -62,8 +59,7 @@ from agentlightning.types import (
 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 .utils import LATENCY_BUCKETS
+from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
 
 server_logger = logging.getLogger("agentlightning.store.server")
 client_logger = logging.getLogger("agentlightning.store.client")
@@ -82,26 +78,12 @@ class RolloutRequest(BaseModel):
     resources_id: Optional[str] = None
     config: Optional[RolloutConfig] = None
     metadata: Optional[Dict[str, Any]] = None
-    worker_id: Optional[str] = None
 
 
 class DequeueRolloutRequest(BaseModel):
     worker_id: Optional[str] = None
 
 
-class StartAttemptRequest(BaseModel):
-    worker_id: Optional[str] = None
-
-
-class EnqueueManyRolloutsRequest(BaseModel):
-    rollouts: List[EnqueueRolloutRequest]
-
-
-class DequeueManyRolloutsRequest(BaseModel):
-    limit: int = 1
-    worker_id: Optional[str] = None
-
-
 class QueryRolloutsRequest(BaseModel):
     status_in: Optional[List[RolloutStatus]] = Field(FastAPIQuery(default=None))
     rollout_id_in: Optional[List[str]] = Field(FastAPIQuery(default=None))
@@ -537,38 +519,22 @@ class LightningStoreServer(LightningStore):
         async def health():  # pyright: ignore[reportUnusedFunction]
             return {"status": "ok"}
 
-        @api.post(API_AGL_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=List[Rollout])
-        async def enqueue_rollouts(  # pyright: ignore[reportUnusedFunction]
-            request: EnqueueManyRolloutsRequest,
-        ) -> List[Rollout]:
-            enqueue_requests = request.rollouts
-            if not enqueue_requests:
-                return []
-            if len(enqueue_requests) == 1:
-                single = enqueue_requests[0]
-                rollout = await self.enqueue_rollout(
-                    input=single.input,
-                    mode=single.mode,
-                    resources_id=single.resources_id,
-                    config=single.config,
-                    metadata=single.metadata,
-                )
-                return [rollout]
-            rollouts = await self.enqueue_many_rollouts(enqueue_requests)
-            return list(rollouts)
+        @api.post(API_AGL_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=Rollout)
+        async def enqueue_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.enqueue_rollout(
+                input=request.input,
+                mode=request.mode,
+                resources_id=request.resources_id,
+                config=request.config,
+                metadata=request.metadata,
+            )
 
-        @api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=List[AttemptedRollout])
-        async def dequeue_rollouts(  # pyright: ignore[reportUnusedFunction]
-            request: DequeueManyRolloutsRequest | None = Body(None),
-        ) -> List[AttemptedRollout]:
-            payload = request or DequeueManyRolloutsRequest()
-            if payload.limit <= 0:
-                return []
-            if payload.limit == 1:
-                single = await self.dequeue_rollout(worker_id=payload.worker_id)
-                return [single] if single else []
-            rollouts = await self.dequeue_many_rollouts(limit=payload.limit, worker_id=payload.worker_id)
-            return list(rollouts)
+        @api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=Optional[AttemptedRollout])
+        async def dequeue_rollout(  # pyright: ignore[reportUnusedFunction]
+            request: DequeueRolloutRequest | None = Body(None),
+        ):
+            worker_id = request.worker_id if request else None
+            return await self.dequeue_rollout(worker_id=worker_id)
 
         @api.post(API_AGL_PREFIX + "/rollouts", status_code=201, response_model=AttemptedRollout)
         async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
@@ -578,7 +544,6 @@ class LightningStoreServer(LightningStore):
                 resources_id=request.resources_id,
                 config=request.config,
                 metadata=request.metadata,
-                worker_id=request.worker_id,
             )
 
         @api.get(API_AGL_PREFIX + "/rollouts", response_model=PaginatedResult[Union[AttemptedRollout, Rollout]])
@@ -597,24 +562,6 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(results, limit=params.limit, offset=params.offset)
 
-        @api.post(API_AGL_PREFIX + "/rollouts/search", response_model=PaginatedResult[Union[AttemptedRollout, Rollout]])
-        async def search_rollouts(request: QueryRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
-            _validate_paginated_request(request, Rollout)
-            status_in = request.status_in if "status_in" in request.model_fields_set else None
-            rollout_id_in = request.rollout_id_in if "rollout_id_in" in request.model_fields_set else None
-            # Get all rollouts from the underlying store
-            results = await self.query_rollouts(
-                status_in=status_in,
-                rollout_id_in=rollout_id_in,
-                rollout_id_contains=request.rollout_id_contains,
-                filter_logic=request.filter_logic,
-                sort_by=request.sort_by,
-                sort_order=request.sort_order,
-                limit=request.limit,
-                offset=request.offset,
-            )
-            return _build_paginated_response(results, limit=request.limit, offset=request.offset)
-
         @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}", response_model=Union[AttemptedRollout, Rollout])
         async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_rollout_by_id(rollout_id)
@@ -647,25 +594,8 @@ class LightningStoreServer(LightningStore):
             )
 
         @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", status_code=201, response_model=AttemptedRollout)
-        async def start_attempt(  # pyright: ignore[reportUnusedFunction]
-            rollout_id: str, request: StartAttemptRequest | None = Body(None)
-        ):
-            worker_id = request.worker_id if request else None
-            return await self.start_attempt(rollout_id, worker_id=worker_id)
-
-        @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/search", response_model=PaginatedResult[Attempt])
-        async def search_attempts(  # pyright: ignore[reportUnusedFunction]
-            rollout_id: str, request: QueryAttemptsRequest
-        ):
-            _validate_paginated_request(request, Attempt)
-            attempts = await self.query_attempts(
-                rollout_id,
-                sort_by=request.sort_by,
-                sort_order=request.sort_order,
-                limit=request.limit,
-                offset=request.offset,
-            )
-            return _build_paginated_response(attempts, limit=request.limit, offset=request.offset)
+        async def start_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.start_attempt(rollout_id)
 
         @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/{attempt_id}", response_model=Attempt)
         async def update_attempt(  # pyright: ignore[reportUnusedFunction]
@@ -694,21 +624,6 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(workers, limit=params.limit, offset=params.offset)
 
-        @api.post(API_AGL_PREFIX + "/workers/search", response_model=PaginatedResult[Worker])
-        async def search_workers(request: QueryWorkersRequest):  # pyright: ignore[reportUnusedFunction]
-            _validate_paginated_request(request, Worker)
-            status_in = request.status_in if "status_in" in request.model_fields_set else None
-            workers = await self.query_workers(
-                status_in=status_in,
-                worker_id_contains=request.worker_id_contains,
-                filter_logic=request.filter_logic,
-                sort_by=request.sort_by,
-                sort_order=request.sort_order,
-                limit=request.limit,
-                offset=request.offset,
-            )
-            return _build_paginated_response(workers, limit=request.limit, offset=request.offset)
-
         @api.get(API_AGL_PREFIX + "/workers/{worker_id}", response_model=Optional[Worker])
         async def get_worker(worker_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_worker_by_id(worker_id)
@@ -722,10 +637,6 @@ class LightningStoreServer(LightningStore):
                 heartbeat_stats=_get_mandatory_field_or_unset(request, "heartbeat_stats"),
             )
 
-        @api.get(API_AGL_PREFIX + "/statistics", response_model=Dict[str, Any])
-        async def get_statistics():  # pyright: ignore[reportUnusedFunction]
-            return await self.statistics()
-
         @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", response_model=PaginatedResult[Attempt])
         async def query_attempts(  # pyright: ignore[reportUnusedFunction]
             rollout_id: str, params: QueryAttemptsRequest = Depends()
@@ -775,7 +686,7 @@ class LightningStoreServer(LightningStore):
         async def get_resources_by_id(resources_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_resources_by_id(resources_id)
 
-        @api.post(API_AGL_PREFIX + "/spans", status_code=201, response_model=Optional[Span])
+        @api.post(API_AGL_PREFIX + "/spans", status_code=201, response_model=Span)
         async def add_span(span: Span):  # pyright: ignore[reportUnusedFunction]
             return await self.add_span(span)
 
@@ -801,28 +712,6 @@ class LightningStoreServer(LightningStore):
             )
             return _build_paginated_response(spans, limit=params.limit, offset=params.offset)
 
-        @api.post(API_AGL_PREFIX + "/spans/search", response_model=PaginatedResult[Span])
-        async def search_spans(request: QuerySpansRequest):  # pyright: ignore[reportUnusedFunction]
-            _validate_paginated_request(request, Span)
-            spans = await self.query_spans(
-                request.rollout_id,
-                request.attempt_id,
-                trace_id=request.trace_id,
-                trace_id_contains=request.trace_id_contains,
-                span_id=request.span_id,
-                span_id_contains=request.span_id_contains,
-                parent_id=request.parent_id,
-                parent_id_contains=request.parent_id_contains,
-                name=request.name,
-                name_contains=request.name_contains,
-                filter_logic=request.filter_logic,
-                sort_by=request.sort_by,
-                sort_order=request.sort_order,
-                limit=request.limit,
-                offset=request.offset,
-            )
-            return _build_paginated_response(spans, limit=request.limit, offset=request.offset)
-
         @api.post(API_AGL_PREFIX + "/spans/next", response_model=NextSequenceIdResponse)
         async def get_next_span_sequence_id(request: NextSequenceIdRequest):  # pyright: ignore[reportUnusedFunction]
             sequence_id = await self.get_next_span_sequence_id(request.rollout_id, request.attempt_id)
@@ -844,62 +733,31 @@ class LightningStoreServer(LightningStore):
     def _setup_prometheus(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,
+                CONTENT_TYPE_LATEST,
                 Counter,
                 Histogram,
-                multiprocess,
+                generate_latest,
             )
         except ImportError:
             raise ImportError(
                 "Prometheus client is not installed. Please either install it or set prometheus to False."
             )
 
-        # 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"],
         )
 
+        # TODO: For multi-process scenarios, should use prometheus_client.multiprocess mode.
         HTTP_LATENCY = Histogram(
             "http_request_duration_seconds",
             "Latency of HTTP requests",
             ["method", "path"],
-            buckets=LATENCY_BUCKETS,
+            buckets=[0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10],
         )
 
-        def get_template_path(path: str) -> str:
-            # Handle "latest" keywords BEFORE generic IDs
-            if path.endswith("/attempts/latest") and "/rollouts/" in path:
-                return re.sub(r"rollouts/[^/]+/attempts/latest$", "rollouts/{rollout_id}/attempts/latest", path)
-            if path.endswith("/attempts/search") and "/rollouts/" in path:
-                return re.sub(r"rollouts/[^/]+/attempts/search$", "rollouts/{rollout_id}/attempts/search", path)
-            if path.endswith("/resources/latest"):
-                return path
-            if path.endswith("/search"):
-                return path
-            if "enqueue" in path or "dequeue" in path:
-                return path
-
-            # Handle generic IDs
-            # (Order matters: longest paths first or lookaheads)
-            path = re.sub(r"/attempts/[^/]+$", "/attempts/{attempt_id}", path)
-            path = re.sub(r"/rollouts/[^/]+", "/rollouts/{rollout_id}", path)  # Handles root and middle
-            path = re.sub(r"/resources/[^/]+$", "/resources/{resources_id}", path)
-            path = re.sub(r"/workers/[^/]+$", "/workers/{worker_id}", path)
-
-            return path
-
         @app.middleware("http")
         async def prometheus_http_middleware(  # pyright: ignore[reportUnusedFunction]
             request: Request, call_next: Callable[[Request], Awaitable[Response]]
@@ -908,8 +766,7 @@ class LightningStoreServer(LightningStore):
             response = await call_next(request)
             elapsed = time.perf_counter() - start
 
-            # Strip the ID-specific URL parts
-            path = get_template_path(request.url.path)
+            path = request.url.path
             method = request.method
             status = response.status_code
 
@@ -918,22 +775,24 @@ class LightningStoreServer(LightningStore):
 
             return response
 
-        metrics_app = make_asgi_app(registry=registry)  # type: ignore
-
-        # This App would need to be accessed via /v1/prometheus/ (note the trailing slash)
-        app.mount(api.prefix + "/prometheus", metrics_app)  # pyright: ignore[reportUnknownArgumentType]
+        @api.get("/prometheus")
+        async def prometheus_metrics():  # pyright: ignore[reportUnusedFunction]
+            return Response(
+                content=generate_latest(),
+                media_type=CONTENT_TYPE_LATEST,
+            )
 
     def _setup_otlp(self, api: APIRouter):
         """Setup OTLP endpoints."""
 
         async def _trace_handler(request: PbExportTraceServiceRequest) -> None:
-            spans = await spans_from_proto(request, self.get_many_span_sequence_ids)
+            spans = await spans_from_proto(request, self)
             server_logger.debug(f"Received {len(spans)} OTLP spans: {', '.join([span.name for span in spans])}")
-            await self.add_many_spans(spans)
+            for span in spans:
+                await self.add_span(span)
 
         # Reserved methods for OTEL traces
         # https://opentelemetry.io/docs/specs/otlp/#otlphttp-request
-        # This is currently the recommended path for Otel compatibility and bulk-insertion support.
         @api.post("/traces")
         async def otlp_traces(request: Request):  # pyright: ignore[reportUnusedFunction]
             return await handle_otlp_export(
@@ -985,8 +844,6 @@ class LightningStoreServer(LightningStore):
 
         @self.app.get("/{full_path:path}", include_in_schema=False)
         def spa_fallback(full_path: str):  # pyright: ignore[reportUnusedFunction]
-            if full_path.startswith("v1/"):
-                raise HTTPException(status_code=404, detail="Not Found")
             # Let the frontend router handle it
             return FileResponse(index_file)
 
@@ -1032,9 +889,6 @@ class LightningStoreServer(LightningStore):
             self._client = LightningStoreClient(self.endpoint)
         return await getattr(self._client, method_name)(*args, **kwargs)
 
-    async def statistics(self) -> LightningStoreStatistics:
-        return await self._call_store_method("statistics")
-
     async def start_rollout(
         self,
         input: TaskInput,
@@ -1042,7 +896,6 @@ class LightningStoreServer(LightningStore):
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
-        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
         return await self._call_store_method(
             "start_rollout",
@@ -1051,7 +904,6 @@ class LightningStoreServer(LightningStore):
             resources_id,
             config,
             metadata,
-            worker_id,
         )
 
     async def enqueue_rollout(
@@ -1071,22 +923,11 @@ class LightningStoreServer(LightningStore):
             metadata,
         )
 
-    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
-        return await self._call_store_method("enqueue_many_rollouts", rollouts)
-
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         return await self._call_store_method("dequeue_rollout", worker_id)
 
-    async def dequeue_many_rollouts(
-        self,
-        *,
-        limit: int = 1,
-        worker_id: Optional[str] = None,
-    ) -> Sequence[AttemptedRollout]:
-        return await self._call_store_method("dequeue_many_rollouts", limit=limit, worker_id=worker_id)
-
-    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
-        return await self._call_store_method("start_attempt", rollout_id, worker_id)
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        return await self._call_store_method("start_attempt", rollout_id)
 
     async def query_rollouts(
         self,
@@ -1172,25 +1013,19 @@ class LightningStoreServer(LightningStore):
     async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
         return await self._call_store_method("get_latest_resources")
 
-    async def add_span(self, span: Span) -> Optional[Span]:
+    async def add_span(self, span: Span) -> Span:
         return await self._call_store_method("add_span", span)
 
-    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
-        return await self._call_store_method("add_many_spans", spans)
-
     async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
         return await self._call_store_method("get_next_span_sequence_id", rollout_id, attempt_id)
 
-    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
-        return await self._call_store_method("get_many_span_sequence_ids", rollout_attempt_ids)
-
     async def add_otel_span(
         self,
         rollout_id: str,
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Optional[Span]:
+    ) -> Span:
         return await self._call_store_method(
             "add_otel_span",
             rollout_id,
@@ -1373,10 +1208,6 @@ class LightningStoreClient(LightningStore):
         """Return the OTLP/HTTP traces endpoint of the store."""
         return f"{self.server_address_root}/v1/traces"
 
-    async def statistics(self) -> LightningStoreStatistics:
-        payload = await self._request_json("get", "/statistics")
-        return cast(LightningStoreStatistics, payload)
-
     def __getstate__(self):
         """
         When LightningStoreClient is pickled (e.g., passed to a subprocess), we only
@@ -1560,7 +1391,6 @@ class LightningStoreClient(LightningStore):
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
-        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
         data = await self._request_json(
             "post",
@@ -1571,7 +1401,6 @@ class LightningStoreClient(LightningStore):
                 resources_id=resources_id,
                 config=config,
                 metadata=metadata,
-                worker_id=worker_id,
             ).model_dump(exclude_none=False),
         )
         return AttemptedRollout.model_validate(data)
@@ -1584,64 +1413,18 @@ class LightningStoreClient(LightningStore):
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
     ) -> Rollout:
-        request_body = EnqueueManyRolloutsRequest(
-            rollouts=[
-                EnqueueRolloutRequest(
-                    input=input,
-                    mode=mode,
-                    resources_id=resources_id,
-                    config=config,
-                    metadata=metadata,
-                )
-            ]
-        ).model_dump(exclude_none=False)
         data = await self._request_json(
             "post",
             "/queues/rollouts/enqueue",
-            json=request_body,
+            json=RolloutRequest(
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                config=config,
+                metadata=metadata,
+            ).model_dump(exclude_none=False),
         )
-        if not data:
-            raise RuntimeError("enqueue_rollout returned no rollouts")
-        return Rollout.model_validate(data[0])
-
-    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
-        if not rollouts:
-            return []
-        request_body = EnqueueManyRolloutsRequest(rollouts=list(rollouts)).model_dump(exclude_none=False)
-        data = await self._request_json(
-            "post",
-            "/queues/rollouts/enqueue",
-            json=request_body,
-        )
-        return [Rollout.model_validate(entry) for entry in data]
-
-    async def _dequeue_batch(
-        self,
-        *,
-        limit: int,
-        worker_id: Optional[str],
-    ) -> List[AttemptedRollout]:
-        if limit <= 0:
-            return []
-        session = await self._get_session()
-        url = f"{self.server_address}/queues/rollouts/dequeue"
-        payload: Dict[str, Any] = {"limit": limit}
-        if worker_id is not None:
-            payload["worker_id"] = worker_id
-        try:
-            async with session.post(url, json=payload) as resp:
-                resp.raise_for_status()
-                data = await resp.json()
-                self._dequeue_was_successful = True
-                return [AttemptedRollout.model_validate(item) for item in data]
-        except Exception as e:
-            if self._dequeue_was_successful:
-                if self._dequeue_first_unsuccessful:
-                    client_logger.warning(f"dequeue_rollout failed with exception: {e}")
-                    self._dequeue_first_unsuccessful = False
-            client_logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
-            # Else ignore the exception because the server is not ready yet
-            return []
+        return Rollout.model_validate(data)
 
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """
@@ -1654,23 +1437,30 @@ class LightningStoreClient(LightningStore):
             This method does NOT retry on failures. If any exception occurs (network error,
             server error, etc.), it logs the error and returns None immediately.
         """
-        attempts = await self._dequeue_batch(limit=1, worker_id=worker_id)
-        return attempts[0] if attempts else None
+        session = await self._get_session()
+        url = f"{self.server_address}/queues/rollouts/dequeue"
+        request_kwargs: Dict[str, Any] = {}
+        if worker_id is not None:
+            request_kwargs["json"] = {"worker_id": worker_id}
+        try:
+            async with session.post(url, **request_kwargs) as resp:
+                resp.raise_for_status()
+                data = await resp.json()
+                self._dequeue_was_successful = True
+                return AttemptedRollout.model_validate(data) if data else None
+        except Exception as e:
+            if self._dequeue_was_successful:
+                if self._dequeue_first_unsuccessful:
+                    client_logger.warning(f"dequeue_rollout failed with exception: {e}")
+                    self._dequeue_first_unsuccessful = False
+            client_logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
+            # Else ignore the exception because the server is not ready yet
+            return None
 
-    async def dequeue_many_rollouts(
-        self,
-        *,
-        limit: int = 1,
-        worker_id: Optional[str] = None,
-    ) -> Sequence[AttemptedRollout]:
-        return await self._dequeue_batch(limit=limit, worker_id=worker_id)
-
-    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
-        payload = {"worker_id": worker_id} if worker_id is not None else None
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
         data = await self._request_json(
             "post",
             f"/rollouts/{rollout_id}/attempts",
-            json=payload,
         )
         return AttemptedRollout.model_validate(data)
 
@@ -1688,25 +1478,29 @@ class LightningStoreClient(LightningStore):
         status: Optional[Sequence[RolloutStatus]] = None,
         rollout_ids: Optional[Sequence[str]] = None,
     ) -> PaginatedResult[Union[AttemptedRollout, Rollout]]:
+        params_list: List[Tuple[str, Any]] = []
+
+        def _extend(key: str, values: Sequence[Any]) -> None:
+            for value in values:
+                params_list.append((key, value))
+
         resolved_status = status_in if status_in is not None else status
         resolved_rollout_ids = rollout_id_in if rollout_id_in is not None else rollout_ids
 
-        payload: Dict[str, Any] = {
-            "limit": limit,
-            "offset": offset,
-        }
         if resolved_status is not None:
-            payload["status_in"] = resolved_status
+            _extend("status_in", resolved_status)
         if resolved_rollout_ids is not None:
-            payload["rollout_id_in"] = resolved_rollout_ids
+            _extend("rollout_id_in", resolved_rollout_ids)
         if rollout_id_contains is not None:
-            payload["rollout_id_contains"] = rollout_id_contains
-        payload["filter_logic"] = filter_logic
+            params_list.append(("rollout_id_contains", rollout_id_contains))
+        params_list.append(("filter_logic", filter_logic))
         if sort_by is not None:
-            payload["sort_by"] = sort_by
-            payload["sort_order"] = sort_order
+            params_list.append(("sort_by", sort_by))
+            params_list.append(("sort_order", sort_order))
+        params_list.append(("limit", limit))
+        params_list.append(("offset", offset))
 
-        data = await self._request_json("post", "/rollouts/search", json=payload)
+        data = await self._request_json("get", "/rollouts", params=params_list or None)
         items = [
             (
                 AttemptedRollout.model_validate(item)
@@ -1726,14 +1520,14 @@ class LightningStoreClient(LightningStore):
         limit: int = -1,
         offset: int = 0,
     ) -> PaginatedResult[Attempt]:
-        payload: Dict[str, Any] = {
-            "limit": limit,
-            "offset": offset,
-        }
+        params: List[Tuple[str, Any]] = [
+            ("limit", limit),
+            ("offset", offset),
+        ]
         if sort_by is not None:
-            payload["sort_by"] = sort_by
-            payload["sort_order"] = sort_order
-        data = await self._request_json("post", f"/rollouts/{rollout_id}/attempts/search", json=payload)
+            params.append(("sort_by", sort_by))
+            params.append(("sort_order", sort_order))
+        data = await self._request_json("get", f"/rollouts/{rollout_id}/attempts", params=params)
         items = [Attempt.model_validate(item) for item in data["items"]]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 
@@ -1776,9 +1570,7 @@ class LightningStoreClient(LightningStore):
         """
         try:
             data = await self._request_json("get", f"/rollouts/{rollout_id}")
-            if data is None:
-                return None
-            elif isinstance(data, dict) and "attempt" in data:
+            if isinstance(data, dict) and "attempt" in data:
                 return AttemptedRollout.model_validate(data)
             else:
                 return Rollout.model_validate(data)
@@ -1868,17 +1660,9 @@ class LightningStoreClient(LightningStore):
             client_logger.error(f"get_latest_resources failed after all retries: {e}", exc_info=True)
             return None
 
-    async def add_span(self, span: Span) -> Optional[Span]:
+    async def add_span(self, span: Span) -> Span:
         data = await self._request_json("post", "/spans", json=span.model_dump(mode="json"))
-        return Span.model_validate(data) if data is not None else None
-
-    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
-        result: List[Span] = []
-        for span in spans:
-            ret = await self.add_span(span)
-            if ret is not None:
-                result.append(ret)
-        return result
+        return Span.model_validate(data)
 
     async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
         data = await self._request_json(
@@ -1889,19 +1673,13 @@ class LightningStoreClient(LightningStore):
         response = NextSequenceIdResponse.model_validate(data)
         return response.sequence_id
 
-    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
-        return [
-            await self.get_next_span_sequence_id(rollout_id, attempt_id)
-            for rollout_id, attempt_id in rollout_attempt_ids
-        ]
-
     async def add_otel_span(
         self,
         rollout_id: str,
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Optional[Span]:
+    ) -> Span:
         # unchanged logic, now benefits from retries inside add_span/get_next_span_sequence_id
         if sequence_id is None:
             sequence_id = await self.get_next_span_sequence_id(rollout_id, attempt_id)
@@ -1911,7 +1689,9 @@ class LightningStoreClient(LightningStore):
             attempt_id=attempt_id,
             sequence_id=sequence_id,
         )
-        return await self.add_span(span)
+        print("created span", span)
+        await self.add_span(span)
+        return span
 
     async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
         """Wait for rollouts to complete.
@@ -1953,30 +1733,32 @@ class LightningStoreClient(LightningStore):
         sort_by: Optional[str] = "sequence_id",
         sort_order: Literal["asc", "desc"] = "asc",
     ) -> PaginatedResult[Span]:
-        payload: Dict[str, Any] = {"rollout_id": rollout_id, "limit": limit, "offset": offset}
+        params: List[Tuple[str, Any]] = [("rollout_id", rollout_id)]
         if attempt_id is not None:
-            payload["attempt_id"] = attempt_id
+            params.append(("attempt_id", attempt_id))
         if trace_id is not None:
-            payload["trace_id"] = trace_id
+            params.append(("trace_id", trace_id))
         if trace_id_contains is not None:
-            payload["trace_id_contains"] = trace_id_contains
+            params.append(("trace_id_contains", trace_id_contains))
         if span_id is not None:
-            payload["span_id"] = span_id
+            params.append(("span_id", span_id))
         if span_id_contains is not None:
-            payload["span_id_contains"] = span_id_contains
+            params.append(("span_id_contains", span_id_contains))
         if parent_id is not None:
-            payload["parent_id"] = parent_id
+            params.append(("parent_id", parent_id))
         if parent_id_contains is not None:
-            payload["parent_id_contains"] = parent_id_contains
+            params.append(("parent_id_contains", parent_id_contains))
         if name is not None:
-            payload["name"] = name
+            params.append(("name", name))
         if name_contains is not None:
-            payload["name_contains"] = name_contains
-        payload["filter_logic"] = filter_logic
+            params.append(("name_contains", name_contains))
+        params.append(("filter_logic", filter_logic))
         if sort_by is not None:
-            payload["sort_by"] = sort_by
-            payload["sort_order"] = sort_order
-        data = await self._request_json("post", "/spans/search", json=payload)
+            params.append(("sort_by", sort_by))
+            params.append(("sort_order", sort_order))
+        params.append(("limit", limit))
+        params.append(("offset", offset))
+        data = await self._request_json("get", "/spans", params=params)
         items = [Span.model_validate(item) for item in data["items"]]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 
@@ -2044,17 +1826,21 @@ class LightningStoreClient(LightningStore):
         limit: int = -1,
         offset: int = 0,
     ) -> PaginatedResult[Worker]:
-        payload: Dict[str, Any] = {}
+        params: List[Tuple[str, Any]] = [
+            ("limit", limit),
+            ("offset", offset),
+        ]
         if status_in is not None:
-            payload["status_in"] = status_in
+            for value in status_in:
+                params.append(("status_in", value))
         if worker_id_contains is not None:
-            payload["worker_id_contains"] = worker_id_contains
-        payload["filter_logic"] = filter_logic
+            params.append(("worker_id_contains", worker_id_contains))
+        params.append(("filter_logic", filter_logic))
         if sort_by is not None:
-            payload["sort_by"] = sort_by
-            payload["sort_order"] = sort_order
+            params.append(("sort_by", sort_by))
+            params.append(("sort_order", sort_order))
 
-        data = await self._request_json("post", "/workers/search", json=payload)
+        data = await self._request_json("get", "/workers", params=params)
         items = [Worker.model_validate(item) for item in data.get("items", [])]
         return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
 

agentlightning/store/collection/__init__.py

@@ -1,21 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import (
-    AtomicLabels,
-    AtomicMode,
-    Collection,
-    FilterOptions,
-    KeyValue,
-    LightningCollections,
-    PaginatedResult,
-    Queue,
-    SortOptions,
-)
+from .base import Collection, FilterOptions, KeyValue, LightningCollections, PaginatedResult, Queue, SortOptions
 from .memory import DequeBasedQueue, DictBasedKeyValue, InMemoryLightningCollections, ListBasedCollection
 
 __all__ = [
-    "AtomicLabels",
-    "AtomicMode",
     "Collection",
     "Queue",
     "KeyValue",

agentlightning/store/collection/base.py

@@ -41,15 +41,6 @@ T = TypeVar("T")  # Recommended to be a BaseModel
 K = TypeVar("K")
 V = TypeVar("V")
 
-AtomicMode = Literal["r", "w", "rw"]
-"""What is expected within the atomic context. Can be "read", "write", or "read-write"."""
-
-AtomicLabels = Literal["rollouts", "attempts", "spans", "resources", "workers", "rollout_queue", "span_sequence_ids"]
-"""Labels for atomic operations.
-
-These labels are used to identify the collections that are affected by the atomic operation.
-"""
-
 
 class Collection(Generic[T]):
     """Behaves like a list of items. Supporting addition, updating, and deletion of items."""
@@ -123,42 +114,19 @@ class Collection(Generic[T]):
         """
         raise NotImplementedError()
 
-    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+    async def update(self, items: Sequence[T]) -> None:
         """Update the given items in the collection.
 
-        Args:
-            items: The items to update in the collection.
-            update_fields: The fields to update. If not provided, all fields in the type will be updated.
-                Only applicable if the item type is a Pydantic BaseModel.
-
         Raises:
             ValueError: If an item with the primary keys does not exist.
-
-        Returns:
-            The items that were updated.
         """
         raise NotImplementedError()
 
-    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+    async def upsert(self, items: Sequence[T]) -> None:
         """Upsert the given items into the collection.
 
         If the items with the same primary keys already exist, they will be updated.
         Otherwise, they will be inserted.
-
-        The operation has three semantics configurable via `update_fields`:
-
-        - `update_or_insert` via `collection.upsert(items, update_fields=["status", "updated_at"])`.
-          If the item with the same primary keys already exists, only the specified fields will be updated.
-          Otherwise, the item will be inserted.
-        - `get_or_insert` via `collection.upsert(items, update_fields=[])`.
-          If the item with the same primary keys already exists, the item will be left unchanged.
-          Otherwise, the item will be inserted.
-        - `replace_ish` via `collection.upsert(items)`.
-          If the item with the same primary keys already exists, all fields from the item will be set.
-          Otherwise, the item will be inserted.
-
-        Returns:
-            The items that were upserted.
         """
         raise NotImplementedError()
 
@@ -297,46 +265,20 @@ class LightningCollections:
         """Dictionary (counter) of span sequence IDs."""
         raise NotImplementedError()
 
-    def atomic(
-        self,
-        *,
-        mode: AtomicMode = "rw",
-        snapshot: bool = False,
-        commit: bool = False,
-        labels: Optional[Sequence[AtomicLabels]] = None,
-        **kwargs: Any,
-    ) -> AsyncContextManager[Self]:
+    def atomic(self, *args: Any, **kwargs: Any) -> AsyncContextManager[Self]:
         """Perform a atomic operation on the collections.
 
         Subclass may use args and kwargs to support multiple levels of atomicity.
-        The arguments can be seen as tags. They only imply the behavior of the operation, not the implementation.
 
         Args:
-            mode: The mode of atomicity. See [`AtomicMode`][agentlightning.store.collection.AtomicMode].
-            snapshot: Enable read snapshot for repeatable reads. Data consistency is guaranteed. The real behavior is implementation-dependent.
-            commit: Enable commitment for write operations. Unsuccessful operations will be rolled back depending on the implementation.
-                Recommend to use [`execute()`][agentlightning.store.collection.LightningCollections.execute] for this level to enable automatic retries.
-                Remember that the real behavior is implementation-dependent.
-            labels: Labels to add to the atomic operation (commonly used as lock names or collection names).
+            *args: Arguments to pass to the operation.
             **kwargs: Keyword arguments to pass to the operation.
         """
         raise NotImplementedError()
 
-    async def execute(
-        self,
-        callback: Callable[[Self], Awaitable[T]],
-        *,
-        mode: AtomicMode = "rw",
-        snapshot: bool = False,
-        commit: bool = False,
-        labels: Optional[Sequence[AtomicLabels]] = None,
-        **kwargs: Any,
-    ) -> T:
-        """Execute the given callback within an atomic operation. Retry on transient errors is implied.
-
-        See [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] for more details.
-        """
-        async with self.atomic(mode=mode, snapshot=snapshot, commit=commit, labels=labels, **kwargs) as collections:
+    async def execute(self, callback: Callable[[Self], Awaitable[T]]) -> T:
+        """Execute the given callback within an atomic operation."""
+        async with self.atomic() as collections:
             return await callback(collections)
 
 

agentlightning/store/collection/memory.py

@@ -4,10 +4,9 @@ from __future__ import annotations
 
 import asyncio
 import logging
-import threading
 import weakref
 from collections import deque
-from contextlib import AsyncExitStack, asynccontextmanager
+from contextlib import asynccontextmanager
 from typing import (
     Any,
     Deque,
@@ -25,8 +24,6 @@ from typing import (
     Union,
 )
 
-from pydantic import BaseModel
-
 from agentlightning.types import (
     Attempt,
     FilterField,
@@ -40,7 +37,6 @@ from agentlightning.types import (
 )
 
 from .base import (
-    AtomicMode,
     Collection,
     FilterMap,
     KeyValue,
@@ -286,7 +282,7 @@ class ListBasedCollection(Collection[T]):
         # We should always return inside the loop.
         raise RuntimeError("Unreachable")
 
-    def _mutate_single(self, item: T, mode: MutationMode, update_fields: Sequence[str] | None = None) -> Optional[T]:
+    def _mutate_single(self, item: T, mode: MutationMode) -> None:
         """Core mutation logic shared by insert, update, upsert, and delete."""
         self._ensure_item_type(item)
         key_values = self._extract_primary_key_values(item)
@@ -303,35 +299,7 @@ class ListBasedCollection(Collection[T]):
             else:  # upsert
                 if not exists:
                     self._size += 1
-                    parent[final_key] = item
-
-                elif update_fields is None:
-                    # update_or_insert: update all fields
-                    parent[final_key] = item
-
-                else:
-                    if not issubclass(self._item_type, BaseModel):
-                        raise TypeError(
-                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
-                        )
-
-                    # Try to fetch the existing item
-                    existing = parent[final_key]
-                    if not isinstance(existing, self._item_type):
-                        raise ValueError(
-                            f"Internal structure corrupted: expected {self._item_type.__name__}, got {type(existing)!r}"
-                        )
-
-                    if not isinstance(item, self._item_type):
-                        raise TypeError(
-                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
-                        )
-
-                    parent[final_key] = parent[final_key].model_copy(
-                        update={field: getattr(item, field) for field in update_fields}
-                    )
-
-            return parent[final_key]
+                parent[final_key] = item
 
         elif mode in ("update", "delete"):
             # For update/delete we must not create missing paths.
@@ -346,22 +314,7 @@ class ListBasedCollection(Collection[T]):
                 raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
 
             if mode == "update":
-                if update_fields is None:
-                    # replace the entire item
-                    parent[final_key] = item
-                else:
-                    if not issubclass(self._item_type, BaseModel):
-                        raise TypeError(
-                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
-                        )
-                    if not isinstance(item, self._item_type):
-                        raise TypeError(
-                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
-                        )
-                    parent[final_key] = parent[final_key].model_copy(
-                        update={field: getattr(item, field) for field in update_fields}
-                    )
-                return parent[final_key]
+                parent[final_key] = item
             else:  # delete
                 del parent[final_key]
                 self._size -= 1
@@ -586,44 +539,22 @@ class ListBasedCollection(Collection[T]):
         Raises:
             ValueError: If any item with the same primary keys already exists.
         """
-        seen_keys: set[Tuple[Any, ...]] = set()
-        prepared: List[T] = []
         for item in items:
-            self._ensure_item_type(item)
-            key_values = self._extract_primary_key_values(item)
-            if key_values in seen_keys:
-                raise ValueError(
-                    f"Insert payload contains duplicate primary key(s): {self._render_key_values(key_values)}"
-                )
-            seen_keys.add(key_values)
-            prepared.append(item)
-
-        for item in prepared:
             self._mutate_single(item, mode="insert")
 
-    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+    async def update(self, items: Sequence[T]) -> None:
         """Update the given items.
 
         Raises:
             ValueError: If any item with the given primary keys does not exist.
         """
-        updated_items: List[T] = []
         for item in items:
-            updated = self._mutate_single(item, mode="update", update_fields=update_fields)
-            if updated is None:
-                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
-            updated_items.append(updated)
-        return updated_items
+            self._mutate_single(item, mode="update")
 
-    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+    async def upsert(self, items: Sequence[T]) -> None:
         """Upsert the given items (insert if missing, otherwise update)."""
-        upserted_items: List[T] = []
         for item in items:
-            upserted = self._mutate_single(item, mode="upsert", update_fields=update_fields)
-            if upserted is None:
-                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
-            upserted_items.append(upserted)
-        return upserted_items
+            self._mutate_single(item, mode="upsert")
 
     async def delete(self, items: Sequence[T]) -> None:
         """Delete the given items.
@@ -719,16 +650,8 @@ class InMemoryLightningCollections(LightningCollections):
     Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
     """
 
-    def __init__(self, lock_type: Literal["thread", "asyncio"]):
-        self._lock = {
-            "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
-            "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
-            "spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
-            "resources": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
-            "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
-            "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
-            "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
-        }
+    def __init__(self):
+        self._lock = _LoopAwareAsyncLock()
         self._rollouts = ListBasedCollection(items=[], item_type=Rollout, primary_keys=["rollout_id"])
         self._attempts = ListBasedCollection(items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"])
         self._spans = ListBasedCollection(
@@ -768,25 +691,9 @@ class InMemoryLightningCollections(LightningCollections):
         return self._span_sequence_ids
 
     @asynccontextmanager
-    async def atomic(
-        self, *, mode: AtomicMode = "rw", snapshot: bool = False, labels: Optional[Sequence[str]] = None, **kwargs: Any
-    ):
-        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.
-
-        Skip the locking if mode is "r" and snapshot is False.
-
-        This collection implementation does NOT support rollback / commit.
-        """
-        if mode == "r" and not snapshot:
-            yield self
-            return
-        if not labels:
-            # If no labels are provided, use all locks.
-            labels = list(self._lock.keys())
-        managers = [self._lock[label] for label in labels]
-        async with AsyncExitStack() as stack:
-            for manager in managers:
-                await stack.enter_async_context(manager)
+    async def atomic(self, *args: Any, **kwargs: Any):
+        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside."""
+        async with self._lock:
             yield self
 
     async def evict_spans_for_rollout(self, rollout_id: str) -> None:
@@ -835,30 +742,3 @@ class _LoopAwareAsyncLock:
         if lock is None or not lock.locked():
             raise RuntimeError("Lock released without being acquired")
         lock.release()
-
-
-class _ThreadSafeAsyncLock:
-    """A threading.Lock that can be used in both async and sync contexts."""
-
-    def __init__(self):
-        self._lock = threading.Lock()
-
-    def __enter__(self):
-        self._lock.acquire()
-        return self
-
-    def __exit__(self, *args: Any, **kwargs: Any):
-        self._lock.release()
-
-    async def __aenter__(self):
-        # We run the blocking .acquire() in a thread pool so we don't block the event loop
-        loop = asyncio.get_running_loop()
-        # NOTE: If this fails to acquire, it will block the executor thread that
-        # is running it. That thread will not auto-terminate when asyncio is cancelled.
-        # Therefore, zombie thread is possible if the lock is held for a long time.
-        await loop.run_in_executor(None, self._lock.acquire)
-        return self
-
-    async def __aexit__(self, *args: Any, **kwargs: Any):
-        # .release() is non-blocking, so we can call it directly
-        self._lock.release()

agentlightning/store/memory.py

@@ -17,9 +17,7 @@ from typing import (
     Literal,
     Mapping,
     Optional,
-    Sequence,
     Set,
-    Tuple,
     TypeVar,
     Union,
     cast,
@@ -27,11 +25,11 @@ from typing import (
 
 from pydantic import BaseModel
 
-from agentlightning.types import AttemptedRollout, NamedResources, PaginatedResult, ResourcesUpdate, Rollout, Span
+from agentlightning.types import AttemptedRollout, PaginatedResult, Rollout, Span
 
-from .base import UNSET, LightningStoreCapabilities, LightningStoreStatistics, Unset, is_finished, is_running
+from .base import UNSET, LightningStoreCapabilities, Unset, is_finished, is_running
 from .collection import InMemoryLightningCollections
-from .collection_based import CollectionBasedLightningStore, tracked
+from .collection_based import CollectionBasedLightningStore
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 
@@ -83,18 +81,12 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     def __init__(
         self,
         *,
-        thread_safe: bool = False,
         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,
     ):
-        super().__init__(
-            collections=InMemoryLightningCollections(lock_type="thread" if thread_safe else "asyncio"),
-            prometheus=prometheus,
-        )
+        super().__init__(collections=InMemoryLightningCollections())
 
-        self._thread_safe = thread_safe
         self._start_time_by_rollout: Dict[str, float] = {}
         self._span_bytes_by_rollout: Dict[str, int] = Counter()
         self._total_span_bytes: int = 0
@@ -140,26 +132,15 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     def capabilities(self) -> LightningStoreCapabilities:
         """Return the capabilities of the store."""
         return LightningStoreCapabilities(
-            thread_safe=self._thread_safe,
+            thread_safe=False,
             async_safe=True,
             zero_copy=False,
             otlp_traces=False,
         )
 
-    async def statistics(self) -> LightningStoreStatistics:
-        """Return the statistics of the store."""
-        return {
-            **(await super().statistics()),
-            "total_span_bytes": self._total_span_bytes,
-            "eviction_threshold_bytes": self._eviction_threshold_bytes,
-            "safe_threshold_bytes": self._safe_threshold_bytes,
-            "memory_capacity_bytes": self._memory_capacity_bytes,
-        }
-
-    @tracked("wait_for_rollout")
     async def wait_for_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
         """Wait for a specific rollout to complete with a timeout."""
-        async with self.collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["rollouts"]) as collections:
+        async with self.collections.atomic() as collections:
             rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
             if rollout and is_finished(rollout):
                 return rollout
@@ -187,74 +168,47 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
 
             # If event was set (not timeout), check if rollout is finished
             if result:
-                async with self.collections.atomic(
-                    mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
-                ) as collections:
+                async with self.collections.atomic() as collections:
                     rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
                     if rollout and is_finished(rollout):
                         return rollout
 
         return None
 
-    @tracked("add_resources_inmemory")
-    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
-        ret = await super().add_resources(resources)
-        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
-            self._latest_resources_id = ret.resources_id
-        return ret
-
-    @tracked("update_resources_inmemory")
-    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
-        ret = await super().update_resources(resources_id, resources)
-        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
-            self._latest_resources_id = ret.resources_id
-        return ret
-
-    @tracked("_post_update_rollout_inmemory")
-    async def _post_update_rollout(self, rollouts: Sequence[Tuple[Rollout, Sequence[str]]]) -> None:
+    async def on_rollout_update(self, rollout: Rollout) -> None:
         """Update the running rollout ids set when the rollout updates."""
-        await super()._post_update_rollout(rollouts)
-        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["rollouts"]):
-            for rollout, _ in rollouts:
-                if is_running(rollout):
-                    self._running_rollout_ids.add(rollout.rollout_id)
-                else:
-                    self._running_rollout_ids.discard(rollout.rollout_id)
+        if is_running(rollout):
+            self._running_rollout_ids.add(rollout.rollout_id)
+        else:
+            self._running_rollout_ids.discard(rollout.rollout_id)
 
-                if is_finished(rollout):
-                    self._completion_events.setdefault(rollout.rollout_id, threading.Event())
-                    self._completion_events[rollout.rollout_id].set()
-                else:
-                    self._completion_events.setdefault(rollout.rollout_id, threading.Event())
-                # Rollout status can never transition from finished to running (unlike attempt)
-                # so we don't need to clear the completion event even in case of retrying.
+        if is_finished(rollout):
+            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
+            self._completion_events[rollout.rollout_id].set()
+        else:
+            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
+        # Rollout status can never transition from finished to running (unlike attempt)
+        # so we don't need to clear the completion event even in case of retrying.
 
-                if rollout.rollout_id not in self._start_time_by_rollout:
-                    self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
+        if rollout.rollout_id not in self._start_time_by_rollout:
+            self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
 
-    @tracked("_unlocked_get_running_rollouts")
-    async def _unlocked_get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
-        """Accelerated version of `_unlocked_get_running_rollouts` for in-memory store. Used for healthcheck."""
-        async with self.collections.atomic(
-            mode="r", snapshot=self._read_snapshot, labels=["rollouts", "attempts"]
-        ) as collections:
-            rollouts = await collections.rollouts.query(
-                filter={"rollout_id": {"within": list(self._running_rollout_ids)}}
+    async def get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
+        """Accelerated version of `get_running_rollouts` for in-memory store. Used for healthcheck."""
+        rollouts = await collections.rollouts.query(filter={"rollout_id": {"within": list(self._running_rollout_ids)}})
+        running_rollouts: List[AttemptedRollout] = []
+        for rollout in rollouts.items:
+            latest_attempt = await collections.attempts.get(
+                filter={"rollout_id": {"exact": rollout.rollout_id}},
+                sort={"name": "sequence_id", "order": "desc"},
             )
-            running_rollouts: List[AttemptedRollout] = []
-            for rollout in rollouts.items:
-                latest_attempt = await collections.attempts.get(
-                    filter={"rollout_id": {"exact": rollout.rollout_id}},
-                    sort={"name": "sequence_id", "order": "desc"},
-                )
-                if not latest_attempt:
-                    # The rollout is running but has no attempts, this should not happen
-                    logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
-                    continue
-                running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
+            if not latest_attempt:
+                # The rollout is running but has no attempts, this should not happen
+                logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
+                continue
+            running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
         return running_rollouts
 
-    @tracked("query_spans_inmemory")  # Since this method calls super, we need to track it separately
     async def query_spans(
         self,
         rollout_id: str,
@@ -265,28 +219,23 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
             raise RuntimeError(f"Spans for rollout {rollout_id} have been evicted")
         return await super().query_spans(rollout_id, attempt_id, **kwargs)
 
-    @tracked("_post_add_spans")
-    async def _post_add_spans(self, spans: Sequence[Span], rollout_id: str, attempt_id: str) -> None:
+    async def _add_span_unlocked(self, collections: InMemoryLightningCollections, span: Span) -> Span:
         """In-memory store needs to maintain the span data in memory, and evict spans when memory is low."""
 
-        await super()._post_add_spans(spans, rollout_id, attempt_id)
-        async with self.collections.atomic(
-            mode="rw", snapshot=self._read_snapshot, labels=["rollouts", "spans"]
-        ) as collections:
-            for span in spans:
-                await self._account_span_size(span)
-            await self._maybe_evict_spans(collections)
+        await super()._add_span_unlocked(collections, span)
+        self._account_span_size(span)
+        await self._maybe_evict_spans(collections)
 
-    @tracked("_get_latest_resources_inmemory")
-    async def _get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        return span
+
+    async def _get_latest_resources_id(self, collections: InMemoryLightningCollections) -> Optional[str]:
         if isinstance(self._latest_resources_id, Unset):
-            return await super()._get_latest_resources()
-        if self._latest_resources_id is not None:
-            async with self.collections.atomic(
-                mode="r", snapshot=self._read_snapshot, labels=["resources"]
-            ) as collections:
-                return await collections.resources.get(filter={"resources_id": {"exact": self._latest_resources_id}})
-        return None
+            latest_resources = await collections.resources.get(sort={"name": "update_time", "order": "desc"})
+            if latest_resources:
+                self._latest_resources_id = latest_resources.resources_id
+            else:
+                self._latest_resources_id = None
+        return self._latest_resources_id
 
     @staticmethod
     def _resolve_memory_threshold(
@@ -318,8 +267,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
 
         return resolved
 
-    @tracked("_account_span_size")
-    async def _account_span_size(self, span: Span) -> int:
+    def _account_span_size(self, span: Span) -> int:
         if self._custom_span_size_estimator is not None:
             size = max(int(self._custom_span_size_estimator(span)), 0)
         else:
@@ -329,7 +277,6 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
         self._total_span_bytes += size
         return size
 
-    @tracked("_maybe_evict_spans")
     async def _maybe_evict_spans(self, collections: InMemoryLightningCollections) -> None:
         if self._total_span_bytes <= self._eviction_threshold_bytes:
             return
@@ -352,7 +299,6 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
             await self._evict_spans_for_rollout(collections, rollout_id)
         logger.info(f"Freed up {memory_consumed_before - self._total_span_bytes} bytes of memory")
 
-    @tracked("_evict_spans_for_rollout")
     async def _evict_spans_for_rollout(self, collections: InMemoryLightningCollections, rollout_id: str) -> None:
         await collections.evict_spans_for_rollout(rollout_id)
         removed_bytes = self._span_bytes_by_rollout.pop(rollout_id, 0)

agentlightning/store/mongo.py

@@ -2,30 +2,21 @@
 
 from __future__ import annotations
 
-import asyncio
 import hashlib
 import logging
-import time
 import uuid
 from typing import (
     Any,
     Callable,
-    Dict,
-    List,
     Mapping,
-    Optional,
-    Sequence,
     TypeVar,
-    Union,
 )
 
 from pymongo import AsyncMongoClient
 
-from agentlightning.types import Attempt, AttemptedRollout, Rollout
-
-from .base import LightningStoreCapabilities, is_finished
-from .collection.mongo import MongoClientPool, MongoLightningCollections, MongoOperationPrometheusTracker
-from .collection_based import CollectionBasedLightningStore, healthcheck_before, tracked
+from .base import LightningStoreCapabilities
+from .collection.mongo import MongoClientPool, MongoLightningCollections
+from .collection_based import CollectionBasedLightningStore
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 
@@ -54,9 +45,7 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
         client: AsyncMongoClient[Mapping[str, Any]] | str,
         database_name: str | None = None,
         partition_id: str | None = None,
-        prometheus: bool = False,
     ) -> None:
-        self._enable_prometheus = prometheus
         self._auto_created_client = False
         if isinstance(client, str):
             self._client = AsyncMongoClient[Mapping[str, Any]](client)
@@ -73,15 +62,7 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
 
         self._client_pool = MongoClientPool(self._client)
 
-        super().__init__(
-            collections=MongoLightningCollections(
-                self._client_pool,
-                database_name,
-                partition_id,
-                prometheus_tracker=MongoOperationPrometheusTracker(enabled=self._enable_prometheus),
-            ),
-            prometheus=self._enable_prometheus,
-        )
+        super().__init__(collections=MongoLightningCollections(self._client_pool, database_name, partition_id))
 
     @property
     def capabilities(self) -> LightningStoreCapabilities:
@@ -99,67 +80,3 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
         # If I created the client, I should close it too.
         if self._auto_created_client:
             await self._client.close()
-
-    @tracked("wait_for_rollouts")
-    @healthcheck_before
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
-        """Wait for specified rollouts to complete with a timeout.
-
-        Concurrently wait for all rollouts to complete with a timeout.
-        """
-        start_time = time.time()
-        current_time = start_time
-        deadline = start_time + timeout if timeout is not None else None
-
-        finished_rollouts: Dict[str, Rollout] = {}
-        unfinished_rollout_ids = set(rollout_ids)
-
-        while deadline is None or current_time <= deadline:
-            async with self.collections.atomic(
-                mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
-            ) as collections:
-                # Query the rollouts that are not finished in a single query
-                rollouts = await collections.rollouts.query(
-                    filter={"rollout_id": {"within": list(unfinished_rollout_ids)}}
-                )
-            for rollout in rollouts.items:
-                if is_finished(rollout):
-                    finished_rollouts[rollout.rollout_id] = rollout
-                    unfinished_rollout_ids.remove(rollout.rollout_id)
-
-            if not unfinished_rollout_ids:
-                break
-
-            # Poll every 10 seconds by default
-            # Minus 0.1 to make sure the time is still sufficient for another call
-            rest_time = max(0.01, min(deadline - time.time() - 0.1, 10.0)) if deadline is not None else 10.0
-            await asyncio.sleep(rest_time)
-            current_time = time.time()
-
-        # 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]
-
-    @tracked("_unlocked_many_rollouts_to_attempted_rollouts")
-    async def _unlocked_many_rollouts_to_attempted_rollouts(
-        self, collections: MongoLightningCollections, rollouts: Sequence[Rollout]
-    ) -> List[Union[Rollout, AttemptedRollout]]:
-        """Query the latest attempts for the rollouts, and attach them to the rollout objects."""
-        async with collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["attempts"]) as collections:
-            attempts = await collections.attempts.query(
-                filter={"rollout_id": {"within": [rollout.rollout_id for rollout in rollouts]}},
-                sort={"name": "sequence_id", "order": "desc"},
-            )
-        latest_attempts: Dict[str, Attempt] = {}
-        for attempt in attempts:
-            if attempt.rollout_id not in latest_attempts:
-                latest_attempts[attempt.rollout_id] = attempt
-            # Otherwise we ignore the attempt because there's already a newer attempt
-
-        return [
-            (
-                AttemptedRollout(**rollout.model_dump(), attempt=latest_attempts[rollout.rollout_id])
-                if rollout.rollout_id in latest_attempts
-                else rollout
-            )
-            for rollout in rollouts
-        ]

agentlightning/store/threading.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import threading
-from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Literal, Optional, Sequence
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -11,7 +11,6 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
-    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
@@ -23,7 +22,7 @@ from agentlightning.types import (
     WorkerStatus,
 )
 
-from .base import UNSET, LightningStore, LightningStoreCapabilities, LightningStoreStatistics, Unset
+from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
 
 
 class LightningStoreThreaded(LightningStore):
@@ -48,11 +47,6 @@ class LightningStoreThreaded(LightningStore):
             "thread_safe": True,
         }
 
-    async def statistics(self) -> LightningStoreStatistics:
-        """Return the statistics of the store."""
-        with self._lock:
-            return await self.store.statistics()
-
     async def start_rollout(
         self,
         input: TaskInput,
@@ -60,17 +54,9 @@ class LightningStoreThreaded(LightningStore):
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
-        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
         with self._lock:
-            return await self.store.start_rollout(
-                input,
-                mode,
-                resources_id,
-                config,
-                metadata,
-                worker_id,
-            )
+            return await self.store.start_rollout(input, mode, resources_id, config, metadata)
 
     async def enqueue_rollout(
         self,
@@ -83,26 +69,13 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
 
-    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
-        with self._lock:
-            return await self.store.enqueue_many_rollouts(rollouts)
-
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         with self._lock:
             return await self.store.dequeue_rollout(worker_id=worker_id)
 
-    async def dequeue_many_rollouts(
-        self,
-        *,
-        limit: int = 1,
-        worker_id: Optional[str] = None,
-    ) -> Sequence[AttemptedRollout]:
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
         with self._lock:
-            return await self.store.dequeue_many_rollouts(limit=limit, worker_id=worker_id)
-
-    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
-        with self._lock:
-            return await self.store.start_attempt(rollout_id, worker_id)
+            return await self.store.start_attempt(rollout_id)
 
     async def query_rollouts(
         self,
@@ -194,11 +167,7 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.get_latest_resources()
 
-    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
-        with self._lock:
-            return await self.store.add_many_spans(spans)
-
-    async def add_span(self, span: Span) -> Optional[Span]:
+    async def add_span(self, span: Span) -> Span:
         with self._lock:
             return await self.store.add_span(span)
 
@@ -208,7 +177,7 @@ class LightningStoreThreaded(LightningStore):
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: int | None = None,
-    ) -> Optional[Span]:
+    ) -> Span:
         with self._lock:
             return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
 
@@ -220,10 +189,6 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
 
-    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
-        with self._lock:
-            return await self.store.get_many_span_sequence_ids(rollout_attempt_ids)
-
     async def query_spans(
         self,
         rollout_id: str,

agentlightning/store/utils.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import time
-from typing import Awaitable, Callable, Dict, List, Tuple
+from typing import Awaitable, Callable, List, cast
 
 from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
 
@@ -9,102 +9,66 @@ UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[Rollout]]
 UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]
 
 
-LATENCY_BUCKETS = [
-    0.000001,
-    0.000002,
-    0.000005,
-    0.00001,
-    0.00002,
-    0.00005,
-    0.0001,
-    0.0002,
-    0.0005,
-    0.001,
-    0.002,
-    0.003,
-    0.005,
-    0.007,
-    0.01,
-    0.015,
-    0.02,
-    0.03,
-    0.05,
-    0.07,
-    0.1,
-    0.2,
-    0.3,
-    0.5,
-    0.7,
-    1.0,
-    2.0,
-    3.0,
-    5.0,
-    7.0,
-    10.0,
-    12.0,
-    15.0,
-    20.0,
-    25.0,
-    30.0,
-    40.0,
-    50.0,
-    60.0,
-    90.0,
-    120.0,
-    180.0,
-    240.0,
-    300.0,
-]
-
-
-async def rollout_status_from_attempt(
+async def propagate_status(
+    update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
     attempt: Attempt,
     config: RolloutConfig,
-) -> RolloutStatus:
+) -> Rollout:
     """
     Propagate the status of an attempt to the rollout.
 
-    Returns:
-        The status of the rollout from the perspective of the attempt.
+    The rollout should be made sure in a state to be outdated.
+    Requeue the rollout if it should be retried.
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
     """
     # Propagate the status directly to the rollout
     if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
-        return attempt.status
+        return await update_rollout_status(
+            attempt.rollout_id,
+            attempt.status,
+        )
 
     if attempt.status == "failed" or attempt.status == "timeout" or attempt.status == "unresponsive":
         # Check if this status should trigger a retry
         if attempt.status in config.retry_condition:
             # If we haven't exceeded max attempts, retry
             if attempt.sequence_id < config.max_attempts:
-                return "requeuing"
+                return await update_rollout_status(
+                    attempt.rollout_id,
+                    "requeuing",
+                )
 
         # If we can't retry or shouldn't retry, mark as failed
-        return "failed"
+        return await update_rollout_status(
+            attempt.rollout_id,
+            "failed",
+        )
 
     raise ValueError(f"Invalid attempt status: {attempt.status}")
 
 
-async def scan_unhealthy_rollouts(
+async def healthcheck(
     rollouts: List[AttemptedRollout],
-) -> Dict[Tuple[str, str], AttemptStatus]:
+    update_rollout_status: UpdateRolloutStatus,
+    update_attempt_status: UpdateAttemptStatus,
+) -> None:
     """
     Perform health check on all running rollouts in the store.
 
     This method should be called periodically to:
 
-    1. Check for unresponsive attempts (no heartbeat or spans for a while)
-    2. Check for timed-out rollouts (running too long since start_time)
+    1. Update rollout status to failed to succeeded when the attempt is done
+    2. Check for unresponsive attempts (no heartbeat or spans for a while)
+    3. Check for timed-out rollouts (running too long since start_time)
+    4. Update attempt/rollout status accordingly
 
     This operation is completely unlocked. The caller is responsible for locking the store.
 
     Args:
-        rollouts: The list of running rollouts to check.
-
-    Returns:
-        A dictionary of updates to the rollouts.
+        store: The LightningStore instance to check rollouts from
     """
     current_time = time.time()
-    updates: Dict[Tuple[str, str], AttemptStatus] = {}
 
     for rollout in rollouts:
         config = rollout.config  # policy for retry and timeout
@@ -112,31 +76,52 @@ async def scan_unhealthy_rollouts(
         # Get the latest attempt for this rollout
         latest_attempt = rollout.attempt
         if not latest_attempt:
-            # This should not happen
+            continue
+
+        # Check if the attempt has already failed or succeeded
+        if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
+            await propagate_status(update_rollout_status, latest_attempt, config)
             continue
 
         # Check for timeout condition (based on attempt start_time, instead of rollout start_time)
         if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
-            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "timeout"
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "timeout",
+            )
             continue
 
         # Check for unresponsive condition (based on last heartbeat)
-        # (1) Haven't received heartbeat for a while
-        if (
-            latest_attempt.last_heartbeat_time
-            and config.unresponsive_seconds is not None
-            and current_time - latest_attempt.last_heartbeat_time > config.unresponsive_seconds
-        ):
-            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
-            continue
+        if latest_attempt.last_heartbeat_time:
+            if latest_attempt.status == "preparing":
+                # If still preparing, mark it as running
+                latest_attempt = await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "running",
+                )
 
-        # (2) Check if there's no last heartbeat (no spans) at all
+            # Haven't received heartbeat for a while
+            if (
+                config.unresponsive_seconds is not None
+                and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
+            ):
+                await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "unresponsive",
+                )
+                continue
+
+        # Check if there's no last heartbeat (no spans) at all
         if (
             latest_attempt.last_heartbeat_time is None
             and config.unresponsive_seconds is not None
             and current_time - latest_attempt.start_time > config.unresponsive_seconds
         ):
-            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
-            continue
-
-    return updates
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "unresponsive",
+            )

agentlightning/tracer/otel.py

@@ -18,9 +18,8 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 
-from agentlightning.semconv import LightningResourceAttributes
 from agentlightning.store.base import LightningStore
-from agentlightning.utils.otel import get_tracer_provider
+from agentlightning.types.tracer import SpanNames
 from agentlightning.utils.otlp import LightningStoreOTLPExporter
 
 from .base import Tracer
@@ -52,16 +51,7 @@ class OtelTracer(Tracer):
         logger.info(f"[Worker {worker_id}] Setting up OpenTelemetry tracer...")
 
         if self._initialized:
-            logger.info(f"[Worker {worker_id}] Tracer provider is already initialized. Skipping initialization.")
-            return
-
-        try:
-            get_tracer_provider()
-            logger.error(
-                f"[Worker {worker_id}] Tracer provider is already initialized but not by OtelTracer. OpenTelemetry may not work as expected."
-            )
-        except RuntimeError:
-            logger.debug(f"[Worker {worker_id}] Tracer provider is not initialized by OtelTracer. Initializing it now.")
+            logger.error("Tracer provider is already initialized. OpenTelemetry may not work as expected.")
 
         self._tracer_provider = TracerProvider()
         trace_api.set_tracer_provider(self._tracer_provider)
@@ -76,7 +66,8 @@ class OtelTracer(Tracer):
 
     def teardown_worker(self, worker_id: int):
         super().teardown_worker(worker_id)
-        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer does NOT remove the tracer provider.")
+        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
+        self._tracer_provider = None
 
     @asynccontextmanager
     async def trace_context(
@@ -153,8 +144,8 @@ class OtelTracer(Tracer):
         tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
             Resource.create(
                 {
-                    LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
-                    LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
+                    SpanNames.ROLLOUT_ID: rollout_id,
+                    SpanNames.ATTEMPT_ID: attempt_id,
                 }
             )
         )
@@ -191,8 +182,8 @@ class OtelTracer(Tracer):
         tracer_provider._resource = tracer_provider._resource.merge(  # pyright: ignore[reportPrivateUsage]
             Resource.create(
                 {
-                    LightningResourceAttributes.ROLLOUT_ID.value: "",
-                    LightningResourceAttributes.ATTEMPT_ID.value: "",
+                    SpanNames.ROLLOUT_ID: "",
+                    SpanNames.ATTEMPT_ID: "",
                 }
             )
         )  # reset resource
@@ -228,30 +219,6 @@ class LightningSpanProcessor(SpanProcessor):
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._loop_thread: Optional[threading.Thread] = None
 
-    def __repr__(self) -> str:
-        return (
-            f"{self.__class__.__name__}("
-            + f"disable_store_submission={self.disable_store_submission}, "
-            + f"store={self.store!r}, "
-            + f"rollout_id={self.rollout_id!r}, "
-            + f"attempt_id={self.attempt_id!r})"
-        )
-
-    @property
-    def store(self) -> Optional[LightningStore]:
-        """The store to submit the spans to."""
-        return self._store
-
-    @property
-    def rollout_id(self) -> Optional[str]:
-        """The rollout ID to submit the spans to."""
-        return self._rollout_id
-
-    @property
-    def attempt_id(self) -> Optional[str]:
-        """The attempt ID to submit the spans to."""
-        return self._attempt_id
-
     @property
     def disable_store_submission(self) -> bool:
         """Whether to disable submitting spans to the store."""

agentlightning/trainer/trainer.py

@@ -213,6 +213,10 @@ class Trainer(TrainerLegacy):
         # We might be able to support a list of resources in future.
         self.initial_resources = initial_resources
 
+        # The active store for the current execution context
+        self.store = self._make_store(store)
+        self.runner = self._make_runner(runner)
+
         self.port = port
 
         self.strategy = self._make_strategy(
@@ -220,11 +224,6 @@ class Trainer(TrainerLegacy):
             n_runners=self.n_runners,
             port=port,
         )
-
-        # The active store for the current execution context
-        self.store = self._make_store(store, self.strategy)
-        self.runner = self._make_runner(runner)
-
         if hasattr(self.strategy, "n_runners"):
             strategy_runners = getattr(self.strategy, "n_runners")
             if isinstance(strategy_runners, int) and strategy_runners > 0:
@@ -283,19 +282,13 @@ class Trainer(TrainerLegacy):
             type_error_fmt="Adapter factory returned {type_name}, which is not a TraceAdapter subclass.",
         )
 
-    def _make_store(self, store: ComponentSpec[LightningStore], strategy: ExecutionStrategy) -> LightningStore:
-        """Resolve the store implementation backing rollouts, attempts, spans, and resources.
-
-        By default, it's always a in-memory store. If using a client/server execution strategy,
-        the in-memory store will be initialized in a thread-safe manner.
-        """
-        is_client_server = isinstance(strategy, ClientServerExecutionStrategy)
-        default_store_factory = lambda: InMemoryLightningStore(thread_safe=is_client_server)
+    def _make_store(self, store: ComponentSpec[LightningStore]) -> LightningStore:
+        """Resolve the store implementation backing rollouts, attempts, spans, and resources."""
         return build_component(
             store,
             expected_type=LightningStore,
             spec_name="store",
-            default_factory=default_store_factory,
+            default_factory=InMemoryLightningStore,
             invalid_spec_error_fmt="Invalid store type: {actual_type}. Expected LightningStore, str, dict, or None.",
             type_error_fmt="Store factory returned {type_name}, which is not a LightningStore subclass.",
         )

agentlightning/types/core.py

@@ -53,7 +53,6 @@ __all__ = [
     "Rollout",
     "Attempt",
     "AttemptedRollout",
-    "EnqueueRolloutRequest",
     "Hook",
     "Worker",
     "WorkerStatus",
@@ -212,24 +211,6 @@ class AttemptedRollout(Rollout):
         return self
 
 
-class EnqueueRolloutRequest(BaseModel):
-    """Payload describing a rollout to be queued via [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout].
-
-    A subset of fields from [`Rollout`][agentlightning.Rollout] used for queuing new rollouts.
-    """
-
-    input: TaskInput
-    """Task input used to generate the rollout."""
-    mode: Optional[RolloutMode] = None
-    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][agentlightning.RolloutMode]."""
-    resources_id: Optional[str] = None
-    """Identifier of the resources required to execute the rollout."""
-    config: Optional[RolloutConfig] = None
-    """Retry and timeout configuration associated with the rollout."""
-    metadata: Optional[Dict[str, Any]] = None
-    """Additional metadata attached to the rollout."""
-
-
 WorkerStatus = Literal["idle", "busy", "unknown"]
 
 

agentlightning/types/tracer.py

@@ -16,8 +16,6 @@ from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 from opentelemetry.trace.status import Status as OtelStatus
 from pydantic import BaseModel, ConfigDict
 
-from agentlightning.semconv import AGL_VIRTUAL
-
 __all__ = [
     "AttributeValue",
     "Attributes",
@@ -381,7 +379,7 @@ class Span(BaseModel):
                 is_remote=False,
                 trace_state={},
             ),
-            name=name or AGL_VIRTUAL,
+            name=name or SpanNames.VIRTUAL.value,
             resource=resource or OtelResource(attributes={}, schema_url=""),
             attributes=attributes,
             status=TraceStatus(status_code="OK"),
@@ -401,7 +399,7 @@ class Span(BaseModel):
 
 
 class SpanNames(str, Enum):
-    """Enumerated span names recognised by Agent-lightning. Deprecated in favor of [semconv][agentlightning.semconv]."""
+    """Enumerated span names recognised by Agent-lightning."""
 
     REWARD = "agentlightning.reward"
     """The name of the reward span."""
@@ -422,7 +420,7 @@ class SpanNames(str, Enum):
 
 
 class SpanAttributeNames(str, Enum):
-    """Canonical attribute names written by Agent Lightning emitters. Deprecated in favor of [semconv][agentlightning.semconv]."""
+    """Canonical attribute names written by Agent Lightning emitters."""
 
     MESSAGE = "message"
     """The name of the message attribute."""

agentlightning/utils/otel.py

@@ -1,401 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Utilities shared for OpenTelemetry span (attributes) support."""
-
-import logging
-from typing import Any, Dict, List, Sequence, Union, cast
-from warnings import filterwarnings
-
-import opentelemetry.trace as trace_api
-from agentops.sdk.exporters import OTLPSpanExporter
-from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SynchronousMultiSpanProcessor, Tracer
-from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
-from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
-from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
-from opentelemetry.trace import get_tracer_provider as otel_get_tracer_provider
-from pydantic import TypeAdapter
-
-from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
-from agentlightning.semconv import LightningSpanAttributes, LinkAttributes, LinkPydanticModel
-from agentlightning.types import SpanLike
-from agentlightning.utils.otlp import LightningStoreOTLPExporter
-
-logger = logging.getLogger(__name__)
-
-__all__ = [
-    "full_qualified_name",
-    "get_tracer_provider",
-    "get_tracer",
-    "make_tag_attributes",
-    "extract_tags_from_attributes",
-    "make_link_attributes",
-    "query_linked_spans",
-    "extract_links_from_attributes",
-    "filter_attributes",
-    "filter_and_unflatten_attributes",
-    "flatten_attributes",
-    "unflatten_attributes",
-]
-
-
-def full_qualified_name(obj: type) -> str:
-    if str(obj.__module__) == "builtins":
-        return obj.__qualname__
-    return f"{obj.__module__}.{obj.__qualname__}"
-
-
-def get_tracer_provider(inspect: bool = True) -> TracerProviderImpl:
-    """Get the OpenTelemetry tracer provider configured for Agent Lightning.
-
-    Args:
-        inspect: Whether to inspect the tracer provider and log its configuration.
-            When it's on, make sure you also set the logger level to DEBUG to see the logs.
-    """
-    from agentlightning.tracer.otel import LightningSpanProcessor
-
-    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
-        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
-    tracer_provider = otel_get_tracer_provider()
-    if not isinstance(tracer_provider, TracerProviderImpl):
-        logger.error(
-            "Tracer provider is expected to be an instance of opentelemetry.sdk.trace.TracerProvider, found: %s",
-            full_qualified_name(type(tracer_provider)),
-        )
-        return cast(TracerProviderImpl, tracer_provider)
-
-    if not inspect:
-        return tracer_provider
-
-    emitter_debug = resolve_bool_env_var(LightningEnvVar.AGL_EMITTER_DEBUG, fallback=None)
-    logger_effective_level = logger.getEffectiveLevel()
-    if emitter_debug is True and logger_effective_level > logging.DEBUG:
-        logger.warning(
-            "Emitter debug logging is enabled but logging level is not set to DEBUG. Nothing will be logged."
-        )
-
-    if emitter_debug is None:
-        # Set to true by default if the logging level is lower than DEBUG
-        emitter_debug = logging.DEBUG >= logger_effective_level
-
-    if emitter_debug:
-        active_span_processor = tracer_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
-        processors: List[str] = []
-        active_span_processor_cls = active_span_processor.__class__.__name__
-        for processor in active_span_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
-            if isinstance(processor, LightningSpanProcessor):
-                # The legacy case for tracers without OTLP support.
-                processors.append(f"{active_span_processor_cls} - {processor!r}")
-            elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
-                processor_cls = processor.__class__.__name__
-                if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
-                    # This should be the main path now.
-                    processors.append(f"{active_span_processor_cls} - {processor_cls} - {processor.span_exporter!r}")
-                elif isinstance(processor.span_exporter, OTLPSpanExporter):
-                    # You need to be careful if the code goes into this path.
-                    endpoint = processor.span_exporter._endpoint  # pyright: ignore[reportPrivateUsage]
-                    processors.append(
-                        f"{active_span_processor_cls} - {processor_cls} - "
-                        f"{processor.span_exporter.__class__.__name__}(endpoint={endpoint!r})"
-                    )
-                else:
-                    # Other cases like Console Span Exporter.
-                    processors.append(
-                        f"{active_span_processor_cls} - {processor_cls} - {processor.span_exporter.__class__.__name__}"
-                    )
-            else:
-                processors.append(f"{active_span_processor_cls} - {processor.__class__.__name__}")
-
-        logger.debug(f"Tracer provider: {tracer_provider!r}. Active span processors:")
-        for processor in processors:
-            logger.debug("  * " + processor)
-
-    return tracer_provider
-
-
-def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
-    """Resolve the OpenTelemetry tracer configured for Agent Lightning.
-
-    Args:
-        use_active_span_processor: Whether to use the active span processor.
-
-    Returns:
-        OpenTelemetry tracer tagged with the `agentlightning` instrumentation name.
-
-    Raises:
-        RuntimeError: If OpenTelemetry was not initialized before calling this helper.
-    """
-    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
-        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
-
-    tracer_provider = get_tracer_provider(inspect=True)  # inspection is on by default
-
-    if use_active_span_processor:
-        return tracer_provider.get_tracer("agentlightning")
-
-    else:
-        filterwarnings(
-            "ignore",
-            message=r"You should use InstrumentationScope. Deprecated since version 1.11.1.",
-            category=DeprecationWarning,
-            module="opentelemetry.sdk.trace",
-        )
-
-        return Tracer(
-            tracer_provider.sampler,
-            tracer_provider.resource,
-            # We use an empty span processor to avoid emitting spans to the tracer
-            SynchronousMultiSpanProcessor(),
-            tracer_provider.id_generator,
-            InstrumentationInfo("agentlightning", "", ""),  # type: ignore
-            SpanLimits(),
-            InstrumentationScope(
-                "agentlightning",
-                "",
-                "",
-                {},
-            ),
-        )
-
-
-def make_tag_attributes(tags: List[str]) -> Dict[str, Any]:
-    """Convert a list of tags into flattened attributes for span tagging.
-
-    There is no syntax enforced for tags, they are just strings. For example:
-
-    ```python
-    ["gen_ai.model:gpt-4", "reward.extrinsic"]
-    ```
-    """
-    return flatten_attributes({LightningSpanAttributes.TAG.value: tags})
-
-
-def extract_tags_from_attributes(attributes: Dict[str, Any]) -> List[str]:
-    """Extract tag attributes from flattened span attributes.
-
-    Args:
-        attributes: A dictionary of flattened span attributes.
-    """
-    maybe_tag_list = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.TAG.value)
-    return TypeAdapter(List[str]).validate_python(maybe_tag_list)
-
-
-def make_link_attributes(links: Dict[str, str]) -> Dict[str, Any]:
-    """Convert a dictionary of links into flattened attributes for span linking.
-
-    Links example:
-
-    ```python
-    {
-        "gen_ai.response.id": "response-123",
-        "span_id": "abcd-efgh-ijkl",
-    }
-    ```
-    """
-    link_list: List[Dict[str, str]] = []
-    for key, value in links.items():
-        if not isinstance(value, str):  # pyright: ignore[reportUnnecessaryIsInstance]
-            raise ValueError(f"Link value must be a string, got {type(value)} for key '{key}'")
-        link_list.append({LinkAttributes.KEY_MATCH.value: key, LinkAttributes.VALUE_MATCH.value: value})
-    return flatten_attributes({LightningSpanAttributes.LINK.value: link_list})
-
-
-def query_linked_spans(spans: Sequence[SpanLike], links: List[LinkPydanticModel]) -> List[SpanLike]:
-    """Query spans that are linked by the given link attributes.
-
-    Args:
-        spans: A sequence of spans to search.
-        links: A list of link attributes to match.
-
-    Returns:
-        A list of spans that match the given link attributes.
-    """
-    matched_spans: List[SpanLike] = []
-
-    for span in spans:
-        span_attributes = span.attributes or {}
-        is_match = True
-        for link in links:
-            # trace_id and span_id must be full match.
-            if link.key_match == "trace_id":
-                if isinstance(span, ReadableSpan):
-                    trace_id = trace_api.format_trace_id(span.context.trace_id) if span.context else None
-                else:
-                    trace_id = span.trace_id
-                if trace_id != link.value_match:
-                    is_match = False
-                    break
-
-            elif link.key_match == "span_id":
-                if isinstance(span, ReadableSpan):
-                    span_id = trace_api.format_span_id(span.context.span_id) if span.context else None
-                else:
-                    span_id = span.span_id
-                if span_id != link.value_match:
-                    is_match = False
-                    break
-
-            else:
-                attribute = span_attributes.get(link.key_match)
-                # attributes must also be a full match currently.
-                if attribute != link.value_match:
-                    is_match = False
-                    break
-
-        if is_match:
-            matched_spans.append(span)
-
-    return matched_spans
-
-
-def extract_links_from_attributes(attributes: Dict[str, Any]) -> List[LinkPydanticModel]:
-    """Extract link attributes from flattened span attributes.
-
-    Args:
-        attributes: A dictionary of flattened span attributes.
-    """
-    maybe_link_list = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.LINK.value)
-    return TypeAdapter(List[LinkPydanticModel]).validate_python(maybe_link_list)
-
-
-def filter_attributes(attributes: Dict[str, Any], prefix: str) -> Dict[str, Any]:
-    """Filter attributes that start with the given prefix.
-
-    The attribute must start with `prefix.` or be exactly `prefix` to be included.
-
-    Args:
-        attributes: A dictionary of span attributes.
-        prefix: The prefix to filter by.
-
-    Returns:
-        A dictionary of attributes that start with the given prefix.
-    """
-    return {k: v for k, v in attributes.items() if k.startswith(prefix + ".") or k == prefix}
-
-
-def filter_and_unflatten_attributes(attributes: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
-    """Filter attributes that start with the given prefix and unflatten them.
-    The prefix will be removed during unflattening.
-
-    Args:
-        attributes: A dictionary of span attributes.
-        prefix: The prefix to filter by.
-
-    Returns:
-        A nested dictionary or list of attributes that start with the given prefix.
-    """
-    filtered_attributes = filter_attributes(attributes, prefix)
-    stripped_attributes: Dict[str, Any] = {}
-    for k, v in filtered_attributes.items():
-        if k == prefix:
-            raise ValueError(f"Cannot unflatten attribute with key exactly equal to prefix: {prefix}")
-        else:
-            stripped_key = k[len(prefix) + 1 :]  # +1 to remove the dot
-        stripped_attributes[stripped_key] = v
-    return unflatten_attributes(stripped_attributes)
-
-
-def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[str, Any]:
-    """Flatten a nested dictionary or list into a flat dictionary with dotted keys.
-
-    This function recursively traverses dictionaries and lists, producing a flat
-    key-value mapping where nested paths are represented via dot-separated keys.
-    Lists are indexed numerically.
-
-    Example:
-
-        >>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}})
-        {"a.b": 1, "a.c.0": 2, "a.c.1": 3}
-
-    Args:
-        nested_data: A nested structure composed of dictionaries, lists, or
-            primitive values.
-
-    Returns:
-        A flat dictionary mapping dotted-string paths to primitive values.
-    """
-
-    flat: Dict[str, Any] = {}
-
-    def _walk(value: Any, prefix: str = "") -> None:
-        if isinstance(value, dict):
-            for k, v in cast(Dict[Any, Any], value).items():
-                if not isinstance(k, str):
-                    raise ValueError(
-                        f"Only string keys are supported in dictionaries, got '{k}' of type {type(k)} in {prefix}"
-                    )
-                new_prefix = f"{prefix}.{k}" if prefix else k
-                _walk(v, new_prefix)
-        elif isinstance(value, list):
-            for idx, item in enumerate(cast(List[Any], value)):
-                new_prefix = f"{prefix}.{idx}" if prefix else str(idx)
-                _walk(item, new_prefix)
-        else:
-            flat[prefix] = value
-
-    _walk(nested_data)
-    return flat
-
-
-def unflatten_attributes(flat_data: Dict[str, Any]) -> Union[Dict[str, Any], List[Any]]:
-    """Reconstruct a nested dictionary/list structure from a flat dictionary.
-
-    Keys are dot-separated paths. Segments that are digit strings will only
-    become list indices if *all* keys in that dict form a consecutive
-    0..n-1 range. Otherwise they remain dict keys.
-
-    Example:
-
-        >>> unflatten_attributes({"a.b": 1, "a.c.0": 2, "a.c.1": 3})
-        {"a": {"b": 1, "c": [2, 3]}}
-
-    Args:
-        flat_data: A dictionary whose keys are dot-separated paths and whose
-            values are primitive data elements.
-
-    Returns:
-        A nested dictionary (and lists where appropriate) corresponding to
-        the flattened structure.
-    """
-    # 1) Build a pure dict tree first (no lists yet)
-    root: Dict[str, Any] = {}
-
-    for flat_key, value in flat_data.items():
-        parts = flat_key.split(".")
-        curr: Dict[str, Any] = root
-
-        for part in parts[:-1]:
-            # Ensure intermediate node is a dict
-            if part not in curr or not isinstance(curr[part], dict):
-                curr[part] = {}
-            curr = curr[part]  # type: ignore[assignment]
-
-        curr[parts[-1]] = value
-
-    # 2) Recursively convert dicts-with-consecutive-numeric-keys into lists
-    def convert(node: Union[Dict[str, Any], List[Any]]) -> Union[Dict[str, Any], List[Any]]:
-        if isinstance(node, dict):
-            # First convert children
-            for k, v in list(node.items()):
-                node[k] = convert(v)
-
-            if not node:
-                # empty dict stays dict
-                return node
-
-            # Check if keys are all numeric strings
-            keys = list(node.keys())
-            if all(isinstance(k, str) and k.isdigit() for k in keys):  # pyright: ignore[reportUnnecessaryIsInstance]
-                indices = sorted(int(k) for k in keys)
-                # Must be exactly 0..n-1
-                if indices == list(range(len(indices))):
-                    return [node[str(i)] for i in range(len(indices))]
-
-            return node
-
-        if isinstance(node, list):  # pyright: ignore[reportUnnecessaryIsInstance]
-            return [convert(v) for v in node]
-
-        # Keep as is
-        return node
-
-    return convert(root)

agentlightning/utils/otlp.py

@@ -1,7 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from __future__ import annotations
-
 import gzip
 import logging
 from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar
@@ -31,7 +29,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
 from opentelemetry.util.types import AttributeValue
 
-from agentlightning.semconv import LightningResourceAttributes
+from agentlightning.store.base import LightningStore
 from agentlightning.types.tracer import (
     Attributes,
     Event,
@@ -39,6 +37,7 @@ from agentlightning.types.tracer import (
     OtelResource,
     Span,
     SpanContext,
+    SpanNames,
     TraceStatus,
     convert_timestamp,
 )
@@ -109,10 +108,7 @@ async def handle_otlp_export(
     )
 
 
-async def spans_from_proto(
-    request: ExportTraceServiceRequest,
-    sequence_id_bulk_issuer: Callable[[Sequence[Tuple[str, str]]], Awaitable[Sequence[int]]],
-) -> List[Span]:
+async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningStore) -> List[Span]:
     """Parse an OTLP proto payload into List[Span].
 
     A store is needed here for generating a sequence ID for each span.
@@ -123,11 +119,11 @@ async def spans_from_proto(
         # Resource-level attributes & IDs
         resource_attrs = _kv_list_to_dict(resource_spans.resource.attributes)
         # rollout_id, attempt_id from resource attributes when present.
-        rollout_id_resource = resource_attrs.get(LightningResourceAttributes.ROLLOUT_ID.value)
-        attempt_id_resource = resource_attrs.get(LightningResourceAttributes.ATTEMPT_ID.value)
+        rollout_id_resource = resource_attrs.get(SpanNames.ROLLOUT_ID)
+        attempt_id_resource = resource_attrs.get(SpanNames.ATTEMPT_ID)
         # If sequence id is provided, all the spans will share the same sequence ID.
         # unless otherwise overridden by span-level attributes.
-        sequence_id_resource = resource_attrs.get(LightningResourceAttributes.SPAN_SEQUENCE_ID.value)
+        sequence_id_resource = resource_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
 
         otel_resource = _resource_from_proto(resource_spans.resource, getattr(resource_spans, "schema_url", ""))
 
@@ -158,9 +154,9 @@ async def spans_from_proto(
 
                 # Try to get if span attributes contain something like rollout_id or attempt_id
                 # Override the resource-level attributes with the span-level attributes if present.
-                rollout_id_span = span_attrs.get(LightningResourceAttributes.ROLLOUT_ID.value)
-                attempt_id_span = span_attrs.get(LightningResourceAttributes.ATTEMPT_ID.value)
-                sequence_id_span = span_attrs.get(LightningResourceAttributes.SPAN_SEQUENCE_ID.value)
+                rollout_id_span = span_attrs.get(SpanNames.ROLLOUT_ID)
+                attempt_id_span = span_attrs.get(SpanNames.ATTEMPT_ID)
+                sequence_id_span = span_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
 
                 # Normalize to regular strings and ints
                 rollout_id_raw = rollout_id_span if rollout_id_span is not None else rollout_id_resource
@@ -182,13 +178,9 @@ async def spans_from_proto(
 
                 # Generate a new sequence ID if not provided
                 if sequence_id is None:
-                    current_sequence_id = -1
-                elif sequence_id < 0:
-                    logger.error(
-                        "Invalid sequence_id value in resource attributes: %r. Must be a positive integer. Regenerating one.",
-                        sequence_id,
+                    current_sequence_id = await store.get_next_span_sequence_id(
+                        rollout_id=rollout_id, attempt_id=attempt_id
                     )
-                    current_sequence_id = -1
                 else:
                     current_sequence_id = sequence_id
 
@@ -214,14 +206,6 @@ async def spans_from_proto(
 
                 output_spans.append(span)
 
-    # Finalize the sequence IDs
-    bulk_issue_requests = [(span.rollout_id, span.attempt_id) for span in output_spans if span.sequence_id < 0]
-    bulk_sequence_ids = await sequence_id_bulk_issuer(bulk_issue_requests)
-    for span, sequence_id in zip(
-        [span for span in output_spans if span.sequence_id < 0], bulk_sequence_ids, strict=True
-    ):
-        span.sequence_id = sequence_id
-
     return output_spans
 
 
@@ -242,36 +226,6 @@ class LightningStoreOTLPExporter(OTLPSpanExporter):
     _rollout_id: Optional[str] = None
     _attempt_id: Optional[str] = None
 
-    def __repr__(self) -> str:
-        return (
-            f"{self.__class__.__name__}("
-            + f"endpoint={self.endpoint!r}, "
-            + f"rollout_id={self.rollout_id!r}, "
-            + f"attempt_id={self.attempt_id!r}, "
-            + f"should_bypass={self.should_bypass()!r})"
-        )
-
-    @property
-    def endpoint(self) -> Optional[str]:
-        """The endpoint to submit the spans to."""
-        if hasattr(self, "_endpoint"):
-            return self._endpoint
-        return None
-
-    @property
-    def rollout_id(self) -> Optional[str]:
-        """The rollout ID to submit the spans to."""
-        if hasattr(self, "_rollout_id"):
-            return self._rollout_id
-        return None
-
-    @property
-    def attempt_id(self) -> Optional[str]:
-        """The attempt ID to submit the spans to."""
-        if hasattr(self, "_attempt_id"):
-            return self._attempt_id
-        return None
-
     def enable_store_otlp(self, endpoint: str, rollout_id: str, attempt_id: str) -> None:
         """Enable storing OTLP data to a specific LightningStore rollout/attempt."""
         self._rollout_id = rollout_id
@@ -300,8 +254,8 @@ class LightningStoreOTLPExporter(OTLPSpanExporter):
                 span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
                     Resource.create(
                         {
-                            LightningResourceAttributes.ROLLOUT_ID.value: self._rollout_id,
-                            LightningResourceAttributes.ATTEMPT_ID.value: self._attempt_id,
+                            SpanNames.ROLLOUT_ID: self._rollout_id,
+                            SpanNames.ATTEMPT_ID: self._attempt_id,
                         }
                     )
                 )

agentlightning/utils/server_launcher.py

@@ -6,7 +6,6 @@ import asyncio
 import inspect
 import logging
 import multiprocessing
-import os
 import queue
 import signal
 import socket
@@ -939,11 +938,6 @@ class PythonServerLauncher:
                     self.args.process_join_timeout / 2
                 ),  # Allow half the timeout for graceful shutdown
             }
-            if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
-                from prometheus_client import multiprocess
-
-                options["child_exit"] = lambda server, worker: multiprocess.mark_process_dead(worker.pid)  # type: ignore
-
             self._gunicorn_app = GunicornApp(self.app, options)
 
             self._proc = ctx.Process(

agentlightning/verl/daemon.py

@@ -9,7 +9,7 @@ import time
 import uuid
 from collections import defaultdict
 from collections.abc import Mapping
-from typing import Any, Dict, List, Literal, Optional, Tuple, cast
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 import numpy as np
 import requests
@@ -22,7 +22,7 @@ from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLeg
 from agentlightning.adapter.triplet import TracerTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store.base import LightningStore
-from agentlightning.types import EnqueueRolloutRequest, Rollout, RolloutConfig, Task
+from agentlightning.types import Rollout, RolloutConfig, Task
 
 __all__ = [
     "AgentModeDaemon",
@@ -377,57 +377,42 @@ class AgentModeDaemon:
         num_samples = len(data[keys[0]])
         rollouts_per_sample = self.train_rollout_n if is_train else 1
 
-        enqueue_rollout_requests: List[EnqueueRolloutRequest] = []
-        data_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
-
         for i in range(num_samples):
             data_id = str(uuid.uuid4())
             original_sample = {key: data[key][i] for key in keys}
             original_sample["data_id"] = data_id
-            data_id_to_original_sample[data_id] = original_sample
 
             # For training, each sample is rolled out multiple times
-            # Data ID is different from Rollout ID, as one data can have multiple rollouts.
             for _ in range(rollouts_per_sample):
                 task_metadata = {"data_id": data_id, "is_train": is_train}
+
+                # Data ID is different from Rollout ID, as one data can have multiple rollouts.
                 if self.mode == "v0":
-                    # Queue immediately
                     rollout_id = await self.server.queue_task(
                         sample=_to_native(original_sample),
                         mode="train" if is_train else "val",
                         resources_id=resources_id,
                         metadata=task_metadata,
                     )
-
-                    # Store original sample data to reconstruct batch information later
-                    self._task_id_to_original_sample[rollout_id] = original_sample
-                    self._total_tasks_queued += 1
                 else:
-                    # Collect tasks to enqueue in batch and queue them later
-                    enqueue_rollout_requests.append(
-                        EnqueueRolloutRequest(
-                            input=_to_native(original_sample),
-                            mode="train" if is_train else "val",
-                            resources_id=resources_id,
-                            config=RolloutConfig(
-                                unresponsive_seconds=self.llm_timeout_seconds,
-                                timeout_seconds=self.llm_timeout_seconds,
-                            ),
-                            metadata=task_metadata,
-                        )
+                    rollout = await self.store.enqueue_rollout(
+                        input=_to_native(original_sample),
+                        mode="train" if is_train else "val",
+                        resources_id=resources_id,
+                        metadata=task_metadata,
                     )
+                    await self.store.update_rollout(
+                        rollout_id=rollout.rollout_id,
+                        config=RolloutConfig(
+                            unresponsive_seconds=self.llm_timeout_seconds,
+                            timeout_seconds=self.llm_timeout_seconds,
+                        ),
+                    )
+                    rollout_id = rollout.rollout_id
 
-        if self.mode == "v1":
-            # Enqueue all the tasks in a single batch
-            rollouts = await self.store.enqueue_many_rollouts(enqueue_rollout_requests)
-            self._task_id_to_original_sample.update(
-                {
-                    # Recover the original data and store it for later use.
-                    rollout.rollout_id: data_id_to_original_sample[cast(Dict[str, Any], rollout.metadata)["data_id"]]
-                    for rollout in rollouts
-                }
-            )
-            self._total_tasks_queued += len(rollouts)
+                # Store original sample data to reconstruct batch information later
+                self._task_id_to_original_sample[rollout_id] = original_sample
+                self._total_tasks_queued += 1
 
     def set_up_data_and_server(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
         """Synchronous wrapper for setting up data and server resources."""
@@ -576,7 +561,7 @@ class AgentModeDaemon:
             final_reward = self._fillna_reward(rollout)
             if not rollout.triplets:
                 print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
-                sample_stat_list.append({"reward": final_reward, "has_reward": final_reward_raw is not None})
+                sample_stat_list.append({"reward": final_reward})
                 continue
             response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
 

docker/compose.mongo.yml

@@ -2,8 +2,9 @@
 # It's used to test the MongoDB store implementation.
 
 services:
+
   mongo:
-    image: mongo:8.2
+    image: mongo:latest
     ulimits:
       nofile:
         soft: 65535

docker/compose.prometheus-memory-store.yml

@@ -1,4 +1,5 @@
 services:
+
   app:
     extends:
       file: compose.store.yml
@@ -10,17 +11,18 @@ services:
     image: prom/node-exporter:latest
     # In CI you might not have full /proc, but this is OK for container-level stats
     pid: "host"
+    network_mode: "service:app" # share network with app for simplicity
     command:
-      - "--path.rootfs=/host"
+      - '--path.rootfs=/host'
     volumes:
-      - "/:/host:ro,rslave"
+      - '/:/host:ro,rslave'
 
   prometheus:
     image: prom/prometheus:latest
     command:
-      - "--config.file=/etc/prometheus/prometheus.yml"
-      - "--storage.tsdb.path=/prometheus"
-      - "--storage.tsdb.retention.time=1h"
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--storage.tsdb.retention.time=1h'
     volumes:
       - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
       - ./data/prometheus:/prometheus
@@ -29,25 +31,3 @@ services:
       - node-exporter
     ports:
       - "9090:9090"
-
-  grafana:
-    image: grafana/grafana:latest
-    ports:
-      - "9091:3000"
-    depends_on:
-      - prometheus
-    volumes:
-      - ./data/grafana:/var/lib/grafana
-      # 1. Mount the Datasource Config
-      - ./grafana/datasource.yml:/etc/grafana/provisioning/datasources/datasource.yml
-      # 2. Mount the Dashboard Provider Config
-      - ./grafana/dashboard-provider.yml:/etc/grafana/provisioning/dashboards/provider.yml
-      # 3. Mount the folder containing the actual JSON files
-      - ./grafana/dashboards:/var/lib/grafana/dashboards
-
-    environment:
-      - GF_INSTALL_PLUGINS=grafana-piechart-panel
-      - GF_AUTH_ANONYMOUS_ENABLED=true
-      - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
-      - GF_AUTH_DISABLE_LOGIN_FORM=true
-      - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json

docker/compose.prometheus-mongo-store.yml

@@ -1,4 +1,5 @@
 services:
+
   mongo:
     extends:
       file: compose.mongo.yml
@@ -22,24 +23,14 @@ services:
     depends_on:
       - mongo
 
-    command:
-      - /bin/bash
-      - -c
-      - |
-          mkdir -p /tmp/prometheus &&
-          agl store --host 0.0.0.0 --port 4747 \
-            --prometheus --backend mongo \
-            --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 \
-            --n-workers ${AGL_STORE_N_WORKERS:-32}
-    environment:
-      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus
+    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend mongo --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 --n-workers 4
 
   mongodb-exporter:
     image: percona/mongodb_exporter:0.47.1
     command:
-      - "--mongodb.uri=mongodb://mongo:27017/"
-      - "--collect-all"
-      - "--mongodb.collstats-colls=agentlightning.rollouts,agentlightning.attempts,agentlightning.spans,agentlightning.resources,agentlightning.workers,agentlightning.rollout_queue,agentlightning.span_sequence_ids"
+      - '--mongodb.uri=mongodb://mongo:27017/'
+      - '--collect-all'
+      - '--mongodb.collstats-colls=agentlightning.rollouts,agentlightning.attempts,agentlightning.spans,agentlightning.resources,agentlightning.workers,agentlightning.rollout_queue,agentlightning.span_sequence_ids'
     depends_on:
       - mongo
     ports:
@@ -50,16 +41,16 @@ services:
     # In CI you might not have full /proc, but this is OK for container-level stats
     pid: "host"
     command:
-      - "--path.rootfs=/host"
+      - '--path.rootfs=/host'
     volumes:
-      - "/:/host:ro,rslave"
+      - '/:/host:ro,rslave'
 
   prometheus:
     image: prom/prometheus:latest
     command:
-      - "--config.file=/etc/prometheus/prometheus.yml"
-      - "--storage.tsdb.path=/prometheus"
-      - "--storage.tsdb.retention.time=1h"
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--storage.tsdb.retention.time=1h'
     volumes:
       - ./prometheus.mongo-store.yml:/etc/prometheus/prometheus.yml:ro
       - ./data/prometheus:/prometheus
@@ -69,25 +60,3 @@ services:
       - node-exporter
     ports:
       - "9090:9090"
-
-  grafana:
-    image: grafana/grafana:latest
-    ports:
-      - "9091:3000"
-    depends_on:
-      - prometheus
-    volumes:
-      - ./data/grafana:/var/lib/grafana
-      # 1. Mount the Datasource Config
-      - ./grafana/datasource.yml:/etc/grafana/provisioning/datasources/datasource.yml
-      # 2. Mount the Dashboard Provider Config
-      - ./grafana/dashboard-provider.yml:/etc/grafana/provisioning/dashboards/provider.yml
-      # 3. Mount the folder containing the actual JSON files
-      - ./grafana/dashboards:/var/lib/grafana/dashboards
-
-    environment:
-      - GF_INSTALL_PLUGINS=grafana-piechart-panel
-      - GF_AUTH_ANONYMOUS_ENABLED=true
-      - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
-      - GF_AUTH_DISABLE_LOGIN_FORM=true
-      - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json

docker/compose.store.yml

@@ -1,4 +1,5 @@
 services:
+
   app:
     build:
       context: ../

docker/grafana/dashboard-provider.yml

@@ -1,12 +0,0 @@
-apiVersion: 1
-
-providers:
-  - name: "default"
-    orgId: 1
-    folder: ""
-    type: file
-    disableDeletion: false
-    updateIntervalSeconds: 10
-    options:
-      # This tells Grafana to look for JSON files in this directory inside the container
-      path: /var/lib/grafana/dashboards

docker/grafana/datasource.yml

@@ -1,8 +0,0 @@
-apiVersion: 1
-
-datasources:
-  - name: Prometheus
-    type: prometheus
-    access: proxy
-    url: http://prometheus:9090
-    isDefault: true

docker/prometheus.memory-store.yml

@@ -6,7 +6,11 @@ scrape_configs:
   - job_name: app
     static_configs:
       - targets: ["app:4747"]
-    metrics_path: /v1/prometheus/
+    metrics_path: /v1/prometheus
+
+  - job_name: mongodb
+    static_configs:
+      - targets: ["mongodb-exporter:9216"]
 
   - job_name: node
     static_configs:

docker/prometheus.mongo-store.yml

@@ -6,12 +6,8 @@ scrape_configs:
   - job_name: app
     static_configs:
       - targets: ["app:4747"]
-    metrics_path: /v1/prometheus/
+    metrics_path: /v1/prometheus
 
   - job_name: node
     static_configs:
       - targets: ["node-exporter:9100"]
-
-  - job_name: mongodb
-    static_configs:
-      - targets: ["mongodb-exporter:9216"]

docker/setup.sh

@@ -3,7 +3,7 @@
 set -euo pipefail
 
 # Create data directories
-mkdir -p data/prometheus data/mongo-container data/mongo-host data/grafana
+mkdir -p data/prometheus data/mongo-container data/mongo-host
 
 # Change permissions
-chmod 777 data/prometheus data/mongo-container data/mongo-host data/grafana
+chmod 777 data/prometheus data/mongo-container data/mongo-host

docs/deep-dive/store.md

@@ -117,7 +117,7 @@ rollout = await store.enqueue_rollout(input, config=cfg)
 | ------------------------------------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------- |
 | N/A                                   | `queuing`                                 | Created by `enqueue_rollout()`.                                                                   |
 | `preparing`                           | `queuing/requeuing` → `preparing`         | Typically `dequeue_rollout()` or `start_rollout()`/`start_attempt()` creates a new attempt.       |
-| `running`                             | `preparing/queuing/requeuing` → `running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `rollout_status_from_attempt`. |
+| `running`                             | `preparing/queuing/requeuing` → `running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `propagate_status`. |
 | `succeeded`                           | `*` → `succeeded`                         | Terminal. Rollout `end_time` set.                                                                 |
 | `failed` / `timeout` / `unresponsive` | `*` → `requeuing`                         | **Only if** `status ∈ retry_condition ∧ sequence_id < max_attempts`.                              |
 | `failed` / `timeout` / `unresponsive` | `*` → `failed`                            | Otherwise (no retries left or retries disabled).                                                  |
@@ -125,7 +125,7 @@ rollout = await store.enqueue_rollout(input, config=cfg)
 
 !!! note "Why aggregation?"
 
-    In code, we use `rollout_status_from_attempt()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollout’s transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
+    In code, we use `propagate_status()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollout’s transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
 
 ## Spans
 

docs/how-to/examples-catalog.md

@@ -30,14 +30,6 @@
 
     [:octicons-repo-24: Browse source]({{ src("examples/calc_x") }})
 
--   :material-code-braces:{ .lg .middle } __Claude Code SWE-bench__
-
-    ---
-
-    Instrumented driver that runs Anthropic's Claude Code workflow on SWE-bench instances while streaming traces through Agent-lightning—supports hosted vLLM, official Anthropic, or any OpenAI-compatible backend and emits datasets for downstream tuning.
-
-    [:octicons-repo-24: Browse source]({{ src("examples/claude_code") }})
-
 -   :material-view-grid:{ .lg .middle } __Minimal building blocks__
 
     ---

docs/reference/agent.md

@@ -22,8 +22,6 @@
 
 ## Emitter
 
-::: agentlightning.emit_annotation
-
 ::: agentlightning.emit_reward
 
 ::: agentlightning.emit_message
@@ -32,11 +30,7 @@
 
 ::: agentlightning.emit_exception
 
-## Emitter Helpers
-
-::: agentlightning.get_message_value
-
-::: agentlightning.get_object_value
+## Reward Helpers
 
 ::: agentlightning.find_final_reward
 
@@ -44,6 +38,8 @@
 
 ::: agentlightning.get_reward_value
 
-::: agentlightning.get_rewards_from_span
-
 ::: agentlightning.is_reward_span
+
+## Legacy Emitter Decorators
+
+::: agentlightning.reward.reward

docs/reference/internal.md

@@ -4,8 +4,6 @@
 
     The following APIs should be used with extra caution because they are very likely to change in the future.
 
-## Algorithms and Adapters
-
 ::: agentlightning.adapter.messages.OpenAIMessages
 
 ::: agentlightning.adapter.triplet.TraceTree
@@ -16,16 +14,12 @@
 
 ::: agentlightning.algorithm.decorator.FunctionalAlgorithm
 
-## LitAgent
-
 ::: agentlightning.litagent.decorator.FunctionalLitAgent
 
 ::: agentlightning.litagent.decorator.llm_rollout
 
 ::: agentlightning.litagent.decorator.prompt_rollout
 
-## LLM Proxy
-
 ::: agentlightning.llm_proxy.ModelConfig
 
 ::: agentlightning.llm_proxy.LightningSpanExporter
@@ -40,58 +34,24 @@
 
 ::: agentlightning.llm_proxy.RolloutAttemptMiddleware
 
-## Store
-
 ::: agentlightning.store.base.UNSET
 
-::: agentlightning.store.utils.rollout_status_from_attempt
-
-::: agentlightning.store.utils.scan_unhealthy_rollouts
-
-## Tracing and OpenTelemetry
+::: agentlightning.store.utils.propagate_status
 
 ::: agentlightning.tracer.otel.LightningSpanProcessor
 
-## Utilities
-
 ::: agentlightning.utils.server_launcher.PythonServerLauncher
 
 ::: agentlightning.utils.server_launcher.PythonServerLauncherArgs
 
 ::: agentlightning.utils.server_launcher.LaunchMode
 
-::: agentlightning.utils.otel.full_qualified_name
-
-::: agentlightning.utils.otel.get_tracer_provider
-
-::: agentlightning.utils.otel.get_tracer
-
-::: agentlightning.utils.otel.make_tag_attributes
-
-::: agentlightning.utils.otel.extract_tags_from_attributes
-
-::: agentlightning.utils.otel.make_link_attributes
-
-::: agentlightning.utils.otel.query_linked_spans
-
-::: agentlightning.utils.otel.extract_links_from_attributes
-
-::: agentlightning.utils.otel.filter_attributes
-
-::: agentlightning.utils.otel.filter_and_unflatten_attributes
-
-::: agentlightning.utils.otel.flatten_attributes
-
-::: agentlightning.utils.otel.unflatten_attributes
-
 ::: agentlightning.utils.otlp.handle_otlp_export
 
 ::: agentlightning.utils.otlp.spans_from_proto
 
 ## Deprecated APIs
 
-::: agentlightning.emitter.reward.reward
-
 ::: agentlightning.server.AgentLightningServer
 
 ::: agentlightning.server.ServerDataStore

docs/reference/store.md

@@ -20,10 +20,6 @@
 
 ## Collections and Collection Implementations
 
-::: agentlightning.store.collection.AtomicMode
-
-::: agentlightning.store.collection.AtomicLabels
-
 ::: agentlightning.store.collection.Collection
 
 ::: agentlightning.store.collection.Queue

docs/reference/types.md

@@ -22,8 +22,6 @@
 
 ::: agentlightning.Rollout
 
-::: agentlightning.EnqueueRolloutRequest
-
 ::: agentlightning.Attempt
 
 ::: agentlightning.AttemptedRollout
@@ -78,20 +76,8 @@
 
 ::: agentlightning.Span
 
+::: agentlightning.SpanNames
+
 ::: agentlightning.SpanAttributeNames
 
 ::: agentlightning.SpanLike
-
-## Semantic Conventions
-
-::: agentlightning.semconv
-
-## Environment Variables
-
-::: agentlightning.LightningEnvVar
-
-::: agentlightning.resolve_bool_env_var
-
-::: agentlightning.resolve_int_env_var
-
-::: agentlightning.resolve_str_env_var

docs/stylesheets/extra.css

@@ -22,10 +22,6 @@
   font-size: 1.15em;
 }
 
-.md-typeset h5, .md-typeset h6 {
-  font-size: 1em;
-}
-
 /* Increase spacing between API references */
 .doc-class, .doc-function, .doc-attribute {
   padding-bottom: 2em;

docs/tutorials/debug.md

@@ -43,7 +43,7 @@ Example output (with a reward span captured):
 
 ```python
 [Rollout(rollout_id='ro-519769241af8', input='Explain why the sky appears blue using principles of light scattering in 100 words.', start_time=1760706315.6996238, ..., status='succeeded')]
-[Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=1, ..., name='agentlightning.annotation', attributes={'agentlightning.reward.0.value': 0.95}, ...)]
+[Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=1, ..., name='agentlightning.reward', attributes={'reward': 0.95}, ...)]
 ```
 
 Swap in an [`AgentOpsTracer`][agentlightning.AgentOpsTracer] instead of [`OtelTracer`][agentlightning.OtelTracer] to see the underlying LLM spans alongside reward information:
@@ -52,7 +52,7 @@ Swap in an [`AgentOpsTracer`][agentlightning.AgentOpsTracer] instead of [`OtelTr
 [
     Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=1, ..., name='openai.chat.completion', attributes={..., 'gen_ai.prompt.0.role': 'user', 'gen_ai.prompt.0.content': 'You are a helpful assistant. Explain why the sky appears blue using principles of light scattering in 100 words.', ...}),
     Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=2, ..., name='openai.chat.completion', attributes={..., 'gen_ai.prompt.0.role': 'user', 'gen_ai.prompt.0.content': 'Evaluate how well the output fulfills the task...', ...}),
-    Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=3, ..., name='agentlightning.annotation', attributes={'agentlightning.reward.0.value': 0.95}, ...)
+    Span(rollout_id='ro-519769241af8', attempt_id='at-a6b62caf', sequence_id=3, ..., name='agentlightning.reward', attributes={'reward': 0.95}, ...)
 ]
 ```
 
@@ -220,7 +220,7 @@ Just like [`Runner.run_context`][agentlightning.Runner.run_context], [`Trainer.d
 21:20:35 [Rollout ro-302fb202bd85 | Attempt 1] ID: at-f84ad21c. Status: succeeded. Worker: Worker-0
 21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span 3a286a856af6bea8] #1 (openai.chat.completion) ... 1.95 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
 21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span e2f44b775e058dd6] #2 (openai.chat.completion) ... 1.24 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
-21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span 45ee3c94fa1070ec] #3 (agentlightning.annotation) ... 0.00 seconds. Attribute keys: ['agentlightning.reward.0.value']
+21:20:35 [Rollout ro-302fb202bd85 | Attempt at-f84ad21c | Span 45ee3c94fa1070ec] #3 (agentlightning.reward) ... 0.00 seconds. Attribute keys: ['reward']
 21:20:35 [Rollout ro-302fb202bd85] Adapted data: [Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=None, metadata={'response_id': '...', 'agent_name': ''}), Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=0.95, metadata={'response_id': '...', 'agent_name': ''})]
 21:20:35 Finished 1 rollouts.
 21:20:35 [Rollout ro-e65a3ffaa540] Status changed to preparing.
@@ -228,7 +228,7 @@ Just like [`Runner.run_context`][agentlightning.Runner.run_context], [`Trainer.d
 21:20:40 [Rollout ro-e65a3ffaa540 | Attempt 1] ID: at-eaefa5d4. Status: succeeded. Worker: Worker-0
 21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 901dd6acc0f50147] #1 (openai.chat.completion) ... 1.30 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
 21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 52e0aa63e02be611] #2 (openai.chat.completion) ... 1.26 seconds. Attribute keys: ['gen_ai.request.type', 'gen_ai.system', ...]
-21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 6c452de193fbffd3] #3 (agentlightning.annotation) ... 0.00 seconds. Attribute keys: ['agentlightning.reward.0.value']
+21:20:40 [Rollout ro-e65a3ffaa540 | Attempt at-eaefa5d4 | Span 6c452de193fbffd3] #3 (agentlightning.reward) ... 0.00 seconds. Attribute keys: ['reward']
 21:20:40 [Rollout ro-e65a3ffaa540] Adapted data: [Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=None, metadata={'response_id': '...', 'agent_name': ''}), Triplet(prompt={'token_ids': []}, response={'token_ids': []}, reward=1.0, metadata={'response_id': '...', 'agent_name': ''})]
 21:20:40 Finished 2 rollouts.
 ```

docs/tutorials/traces.md

@@ -110,7 +110,7 @@ You can also customize an [`Adapter`][agentlightning.Adapter] by extending the i
 
 ### Reading Rewards
 
-Rewards are recorded as dedicated spans named [`agentlightning.annotation`][agentlightning.semconv.AGL_ANNOTATION]. Emitting a reward through [`emit_reward`][agentlightning.emit_reward] or [`emit_annotation`][agentlightning.emit_annotation] ensures the value is stored in the span’s `attributes`. To audit rewards, fetch spans from the store and use the helper utilities in [`agentlightning.emitter`](../reference/agent.md):
+Rewards are recorded as dedicated spans named [`agentlightning.reward`][agentlightning.SpanNames.REWARD]. Emitting a reward through [`emit_reward`][agentlightning.emit_reward] or the [`@reward` decorator][agentlightning.reward.reward] ensures the value is stored in the span’s `attributes["reward"]`. To audit rewards, fetch spans from the store and use the helper utilities in [`agentlightning.emitter`](../reference/agent.md):
 
 ```python
 from agentlightning.emitter import find_final_reward

docs/tutorials/write-agents.md

@@ -115,7 +115,7 @@ The value your agent function returns (i.e., the return value of the function de
 
 !!! important "Emitting the Final Reward"
 
-    When returning `None`, you must still ensure a final reward is logged. You can do this by using the [`emit_reward`][agentlightning.emit_reward] function (covered in the [Emitter section][using-emitter] below). Wrapping your reward calculation function with the `@reward` decorator is NOT the recommended approach any more.
+    When returning `None`, you must still ensure a final reward is logged. You can do this by using the [`emit_reward`][agentlightning.emit_reward] function (covered in the [Emitter section][using-emitter] below) or by wrapping your reward calculation function with the [`@reward`][agentlightning.reward.reward] decorator.
 
 * **`list[ReadableSpan]`** or **`list[Span]`**: For advanced use cases, you can manually construct and return a complete list of all spans for the rollout. This gives you full control over the trace data. You can return either a list of OpenTelemetry `ReadableSpan` objects or Agent-lightning's native `Span` objects.
 

examples/.gitignore

@@ -4,7 +4,6 @@ outputs/
 checkpoints/
 calc-x-data.zip
 spider-data.zip
-claude_code/logs/
 agentops.log
 unsloth/models/
 unsloth/unsloth_compiled_cache/

examples/README.md

@@ -7,7 +7,6 @@ This catalog highlights the examples shipped with Agent-lightning.
 | [apo](./apo) | Automatic Prompt Optimization tutorials covering built-in, custom, and debugging workflows. | [![apo workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-apo.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml) |
 | [azure](./azure) | Supervised fine-tuning with Azure OpenAI. | [![azure workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-azure.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-azure.yml) |
 | [calc_x](./calc_x) | VERL-powered math reasoning agent training that uses AutoGen with an MCP calculator tool. | [![calc_x workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-calc-x.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-calc-x.yml) |
-| [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 |

examples/azure/README.md

@@ -1,7 +1,5 @@
 # Supervised Fine-tuning with Azure OpenAI
 
-[![azure CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-azure.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-azure.yml)
-
 This example walks through an end-to-end supervised fine-tuning loop on Azure OpenAI. The trainer runs a toy capital-lookup agent, collects traces with rewards, submits fine-tuning jobs using those traces, and deploys every successful checkpoint as a new Azure OpenAI deployment.
 
 **NOTE: The example is tested and compatible with Agent-lightning v0.2.x, but it's not yet maintained on CI due to the difficulty of maintaining a logged-in status in the testing environment.**

examples/calc_x/train_calc_agent.py

@@ -38,7 +38,6 @@ from calc_agent import MathProblem, calc_agent
 from datasets import Dataset as HuggingFaceDataset
 
 import agentlightning as agl
-from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var, resolve_str_env_var
 
 
 def verl_default_config() -> Dict[str, Any]:
@@ -154,7 +153,7 @@ def train(
         PROJECT_NAME = "AgentLightningCI"
 
         # Skip this step if AGL_CURRENT_ROLE is runner
-        agl_current_role = resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE)
+        agl_current_role = os.getenv("AGL_CURRENT_ROLE")
 
         if agl_current_role != "runner":
             # Simulate writing to $GITHUB_OUTPUT if it’s set
@@ -223,7 +222,7 @@ def main():
 
     if args.external_store_address:
         print(f"Connecting to external store at: {args.external_store_address}")
-        if resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=True):
+        if not os.getenv("AGL_MANAGED_STORE"):
             raise ValueError(
                 "When using an external store, please set the environment variable AGL_MANAGED_STORE=0. "
                 "Otherwise the trainer will still try to manage the store lifecycle for you!"

examples/claude_code/README.md

@@ -1,112 +0,0 @@
-# Training Claude Code with Agent-lightning
-
-[![claude-code CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml)
-
-This example shows how to wrap Anthropic's Claude Code experience with Agent-lightning instrumentation to solve SWE-bench tasks, collect spans/logs, and optionally convert those traces into HuggingFace datasets.
-
-**NOTE:** This example only shows how to integrate Claude Code as an agent in Agent-lightning. The training part is still under development and welcoming contributions!
-
-## Overview
-
-`claude_code_agent.py` spins up a Lightning Store, an LLM proxy, and the Claude Code controller. Each SWE-bench instance is executed inside the official container image so you can either prompt-tune against Anthropic's hosted models or point Claude Code at a self-hosted OpenAI-compatible backend such as vLLM. When a backend surfaces token IDs/logprobs (e.g., vLLM), the traces are turned into triplets that downstream fine-tuning pipelines can consume.
-
-## Requirements
-
-First, install Agent-lightning following the [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/). Then install the SWE-bench harness plus utilities used by this example:
-
-```bash
-(uv) pip install swebench transformers datasets python-dotenv
-```
-
-Docker must be available because each SWE-bench instance is executed in a container via `swebench_utils`.
-
-Finally, set API credentials depending on backend:
-
-- `ANTHROPIC_API_KEY` for the official Claude Code path.
-- `OPENAI_API_KEY` (or another OpenAI-compatible key) for the `openai` backend.
-- A running OpenAI-compatible server (e.g., vLLM) when using the `vllm` backend.
-
-## Dataset
-
-`swebench_samples.jsonl` contains a handful of SWE-bench issues for smoke testing. For full-scale benchmarks load `princeton-nlp/SWE-bench` via `load_swebench_dataset` or point `--dataset-path` to your own JSONL file.
-
-## Included Files
-
-| File/Directory | Description |
-|----------------|-------------|
-| `claude_code_agent.py` | CLI entry point that launches the Lightning store, LLM proxy, and Claude Code agent |
-| `claude_code_controller.py` | Manages the SWE-bench Docker runtime and translates model outputs into git patches |
-| `extended_adapter.py` | Adapter that converts LLM proxy spans into triplets with token IDs, logprobs, and chat history |
-| `swebench_samples.jsonl` | Mini SWE-bench subset for quick validation |
-| `swebench_utils/` | Utilities for running/evaluating SWE-bench instances inside containers |
-| `templates/handle_hook.template.sh` | Helper script injected into containers for hook handling |
-| `templates/settings.template.json` | Base configuration consumed by Claude Code CLI |
-
-## Running the Example
-
-All commands are issued from `examples/claude_code`. Inspect the module-level docstring in `claude_code_agent.py` for the full CLI reference.
-
-### Hosted vLLM (open-source models)
-
-First, launch your model behind an OpenAI-compatible endpoint, for example:
-
-```bash
-vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
-    --max-model-len 131072 \
-    --enable-auto-tool-choice \
-    --tool-call-parser qwen3_coder
-```
-
-Run the Agent-lightning harness and point it at the server:
-
-```bash
-python claude_code_agent.py vllm \
-    --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct \
-    --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct \
-    --frontend-model-high claude-sonnet-4-5-20250929 \
-    --frontend-model-low claude-haiku-4-5-20251001 \
-    --base-url http://localhost:8000/v1 \
-    --dataset-path swebench_samples.jsonl \
-    --output-dir data_debug \
-    --max-turns 5 \
-    --limit 2
-```
-
-The backend model names must match what the server exposes. Because this mode surfaces token IDs/logprobs, the script saves both raw span logs and HuggingFace datasets per instance.
-
-### Official Claude Code (Anthropic API)
-
-```bash
-export ANTHROPIC_API_KEY=sk-...
-python claude_code_agent.py anthropic \
-    --dataset-path swebench_samples.jsonl \
-    --output-dir data_anthropic \
-    --frontend-model-high claude-sonnet-4-5-20250929 \
-    --frontend-model-low claude-haiku-4-5-20251001
-```
-
-Backend model flags are optional here because the Anthropic API strings match the frontend names. This path is ideal for validating prompts against the hosted experience (trace outputs do not contain token IDs or logprobs).
-
-### OpenAI-Compatible Providers
-
-```bash
-export OPENAI_API_KEY=sk-...
-python claude_code_agent.py openai \
-    --backend-model-high gpt-4.1 \
-    --backend-model-low gpt-4o-mini \
-    --dataset-path swebench_samples.jsonl \
-    --output-dir data_openai
-```
-
-Use this mode whenever Claude Code should talk to Azure OpenAI, OpenAI, or another compatible provider. `--base-url` is optional—pass it if your endpoint differs from the public OpenAI URL.
-
-Adjust `--max-turns`, `--cooldown-seconds`, and `--limit` to control runtime and rate limits regardless of backend.
-
-## Outputs and Trace Collection
-
-- `output_dir/stream_<instance_id>.json` contains the complete span stream captured from the Lightning Store for each rollout.
-- When running with `backend_type=vllm`, `output_dir/dataset-<instance_id>/` stores a HuggingFace dataset with token IDs, logprobs, prompts, and metadata produced by `ExtendedLlmProxyTraceToTriplet`.
-- `logs/<instance_id>/` is created by the SWE-bench runtime and mirrors the console output from the container.
-- Return values from the agent are also evaluated via `swebench_utils.evaluation.evaluate`, so `data_debug` (or your chosen folder) will contain evaluation reports alongside traces.
-
-Use these artifacts to fine-tune models, debug Claude Code behavior, or replay rollouts in downstream Agent-lightning workflows.

examples/claude_code/claude_code_agent.py

@@ -1,540 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Instrumented driver for running Claude Code on SWE-bench with Agent-lightning.
-
-This script wires together the Lightning Store, LLM proxy, and Claude Code controller so
-that every SWE-bench instance is executed inside the official Claude container while
-capturing full Agent-lightning traces. It supports three backend modes:
-
-- `vllm`: wrap an OpenAI-compatible endpoint (e.g., vLLM) for hosted OSS models while
-  collecting prompt/response token ids and logprobs.
-- `anthropic`: call the official Claude Code API via `ANTHROPIC_API_KEY` for prompt
-  tuning. Backend model defaults to the provided frontend names.
-- `openai`: route through any OpenAI-compatible provider using `OPENAI_API_KEY`.
-
-Typical usage: hosted vLLM (requires model paths and --base-url)
-
-```bash
-# Run vLLM in background
-vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
-  --max-model-len 131072 \
-  --enable-auto-tool-choice \
-  --tool-call-parser qwen3_coder \
-  --port 45993 &
-
-python claude_code_agent.py vllm \
-  --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct \
-  --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct \
-  --base-url http://localhost:45993/v1 \
-  --dataset-path swebench_samples.jsonl \
-```
-
-Official Claude Code via Anthropic:
-
-```bash
-export ANTHROPIC_API_KEY=sk-...
-python claude_code_agent.py anthropic \
-  --dataset-path swebench_samples.jsonl \
-  --output-dir data_anthropic
-```
-
-Any OpenAI-compatible backend:
-
-```bash
-export OPENAI_API_KEY=sk-...
-python claude_code_agent.py openai \
-  --backend-model-high gpt-5.1-codex-mini \
-  --backend-model-low gpt-4.1-mini \
-  --dataset-path swebench_samples.jsonl
-```
-
-Use `--debug` to enable debug loggings.
-"""
-
-import asyncio
-import json
-import logging
-import os
-import resource
-from argparse import ArgumentParser
-from typing import Any, Dict, List, Literal, Optional, Sequence, cast
-
-from claude_code_controller import ClaudeController
-from datasets import Dataset
-from extended_adapter import ExtendedLlmProxyTraceToTriplet
-from swebench.harness.constants import SWEbenchInstance
-from swebench.harness.utils import load_swebench_dataset  # pyright: ignore[reportUnknownVariableType]
-from swebench_utils.evaluation import evaluate
-from swebench_utils.logging import log_for_evaluation
-from transformers import AutoTokenizer, PreTrainedTokenizerBase
-
-from agentlightning import (
-    InMemoryLightningStore,
-    LightningStoreServer,
-    LitAgentRunner,
-    OtelTracer,
-    setup_logging,
-    setup_module_logging,
-)
-from agentlightning.litagent import LitAgent
-from agentlightning.llm_proxy import LLMProxy, ModelConfig
-from agentlightning.store import LightningStore
-from agentlightning.types import AttemptedRollout, NamedResources, ProxyLLM, Rollout, RolloutRawResult, Span
-
-logger = logging.getLogger("claude_code_agent")
-
-
-def _load_dataset(path: str, epoch: int = 0, limit: Optional[int] = None) -> List[SWEbenchInstance]:
-    instances: List[SWEbenchInstance] = []
-    with open(path) as f:
-        for line in f:
-            instance = json.loads(line)
-            instance["epoch"] = epoch
-            instances.append(instance)
-
-    if limit is not None:
-        instances = instances[:limit]
-    return instances
-
-
-def _flatten_messages(messages: List[Any]) -> List[Dict[str, str]]:
-    flattened: List[Dict[str, str]] = []
-    for msg in messages:
-        if msg["role"] in ["system", "user"] and isinstance(msg["content"], list):
-            msg_content: List[str] = []
-            for content in msg["content"]:
-                msg_content.append(content["text"])
-
-            msg["content"] = "".join(msg_content)
-        elif msg["role"] == "assistant" and "tool_calls" in msg:
-            # NOTE:
-            # Tool calls are list of dict, though in most case only one tool call is made per call
-            # We serialize it as json string here to avoid nested structure
-            msg["tool_calls"] = json.dumps(msg["tool_calls"])
-
-        for k in msg:
-            assert isinstance(msg[k], str), f"\n>>> {msg}"
-        flattened.append(msg)
-    return flattened
-
-
-class ClaudeCodeAgent(LitAgent[SWEbenchInstance]):
-    """Claude Code Agent implementation.
-
-    This agent is a wrapper of the Claude Code controller,
-    and it should be used to run the Claude Code agent on SWE-bench datasets.
-    """
-
-    def __init__(
-        self,
-        namespace: Literal["swebench", "starryzhang"] = "swebench",
-        max_turns: int = 5,
-        run_method: Literal["python", "cli"] = "cli",
-        open_file_limit: int = 4096,
-        cache_level: str = "env",  # ["none", "base", "env", "instance"]
-        clean: bool = False,
-        force_rebuild: bool = False,
-        timeout: int = 1_800,  # in sec
-        instance_image_tag: str = "latest",
-        rewrite_reports: bool = False,
-        swebench_full_dataset: Optional[List[SWEbenchInstance]] = None,
-    ) -> None:
-        super().__init__()
-        self.namespace = namespace
-        self.max_turns = max_turns
-        self.run_method = run_method
-
-        self.cache_level = cache_level
-        self.clean = clean
-        self.force_rebuild = force_rebuild
-        self.timeout = timeout
-        self.instance_image_tag = instance_image_tag
-        self.rewrite_reports = rewrite_reports
-
-        self.swebench_full_dataset = (
-            {each["instance_id"]: each for each in swebench_full_dataset} if swebench_full_dataset is not None else {}
-        )
-
-        # Set the maximum number of open files to the specified limit.
-        resource.setrlimit(resource.RLIMIT_NOFILE, (open_file_limit, open_file_limit))
-
-    async def rollout_async(
-        self, task: SWEbenchInstance, resources: NamedResources, rollout: Rollout
-    ) -> RolloutRawResult:
-        if not isinstance(rollout, AttemptedRollout):
-            # Technically, rollout should be an AttemptedRollout here.
-            # but the API is not stabilized yet.
-            raise ValueError("Rollout is not an AttemptedRollout.")
-
-        run_id = f"epoch_{task.get('epoch', 0)}"
-        image = f"{self.namespace}/sweb.eval.x86_64.{task['instance_id'].lower()}".replace("__", "_1776_")
-
-        llm = cast(ProxyLLM, resources["llm"])
-
-        try:
-            # 1. init container
-            controller = ClaudeController(
-                image,
-                task,
-                run_id,
-                llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
-                llm.api_key or os.environ.get("ANTHROPIC_AUTH_TOKEN", "dummy"),
-            )
-            # 2. execute task
-            prediction = controller.run_instance(
-                task, max_turns=self.max_turns, run_method=cast(Literal["python", "cli"], self.run_method)
-            )
-            del controller
-        except Exception as e:
-            log_for_evaluation(run_id, task["instance_id"], f"Exception during rollout: {e}")
-            return 0.0
-
-        # 3. obtain rewards (evaluation result)
-        reward = 0.0
-        # empty patch
-        if prediction["model_patch"] in ["", None]:
-            return reward
-
-        instance_id = prediction["instance_id"]
-
-        result = evaluate(
-            cast(Any, prediction),
-            self.swebench_full_dataset[instance_id],
-            self.cache_level,
-            self.clean,
-            self.force_rebuild,
-            run_id,
-            self.timeout,
-            namespace=self.namespace,
-            instance_image_tag=self.instance_image_tag,
-            rewrite_reports=self.rewrite_reports,
-        )
-
-        # error patch
-        if result is None:
-            return reward
-
-        report = result[1]
-        # resolved/unresolved patch
-        if report[instance_id]["resolved"]:
-            reward = 1.0
-        return reward
-
-
-def sanity_check_spans(spans: Sequence[Span]) -> None:
-    assert len(spans) > 1, f"At least two spans are expected for a valid rollout. Found {len(spans)} spans."
-    assert any(span.name == "raw_gen_ai_request" for span in spans), "raw_gen_ai_request span not found"
-    assert any(span.name == "agentlightning.annotation" for span in spans), "agentlightning.annotation span not found"
-
-
-async def run_instance_async(
-    instance: SWEbenchInstance,
-    agent: ClaudeCodeAgent,
-    runner: LitAgentRunner[SWEbenchInstance],
-    store: LightningStore,
-    output_dir: Optional[str],
-    adapter: Optional[ExtendedLlmProxyTraceToTriplet],
-    tokenizer: Optional[PreTrainedTokenizerBase],
-) -> None:
-    """Runs the agent on a specific SWE-bench instance.
-
-    Running on specific SWE-bench instance and queries the traced spans.
-    It then extracts the triplets and saves the dataset.
-    """
-
-    instance_id = instance["instance_id"]
-    logger.info(f"Starting to run instance: {instance_id}")
-
-    # Run the agent and query the traced spans.
-    with runner.run_context(agent=agent, store=store):
-        rollout = await runner.step(instance)
-
-    logger.info(f"Finished running instance: {instance_id}")
-
-    spans = await store.query_spans(rollout.rollout_id)
-
-    if output_dir is None:
-        logger.info(f"Generated {len(spans)} spans for {instance_id}")
-        return
-
-    # 1. Dump raw spans (Common for both types)
-    raw_path = os.path.join(output_dir, f"stream_{instance_id}.json")
-    with open(raw_path, "w") as f:
-        for span in spans:
-            f.write(json.dumps(span.model_dump()) + "\n")
-    logger.info(f"Dumped {len(spans)} spans to {raw_path}")
-
-    # 2. Extract Triplets and Save Dataset (vLLM specific)
-    if adapter is not None and tokenizer is not None:
-        try:
-            triplets = adapter.adapt(cast(List[Span], spans))
-            logger.info(f"Extracted {len(triplets)} triplets for {instance_id}")
-
-            all_triplets: List[Dict[str, Any]] = []
-            recent_reward: Optional[float] = None
-
-            # Process in reverse to propagate rewards if necessary/logic dictates
-            for triplet in reversed(triplets):
-                if triplet.reward is not None:
-                    recent_reward = triplet.reward
-
-                prompt_text = tokenizer.decode(triplet.prompt["token_ids"])  # type: ignore
-                all_triplets.append(
-                    {
-                        "repo": instance.get("repo", ""),
-                        "instance_id": instance_id,
-                        "turn": triplet.metadata["sequence_id"],
-                        "prompt_ids": triplet.prompt["token_ids"],
-                        "gold_completion_ids": triplet.response["token_ids"],
-                        "logprobs": triplet.response["logprobs"],
-                        "reward": recent_reward,
-                        "prompt": prompt_text,
-                        "messages": _flatten_messages(triplet.metadata["messages"]),
-                    }
-                )
-
-            if all_triplets:
-                ds = Dataset.from_list(all_triplets)  # type: ignore
-                save_path = os.path.join(output_dir, f"dataset-{instance_id}")
-                ds.save_to_disk(save_path)  # type: ignore
-                logger.info(f"Saved HuggingFace dataset to {save_path}")
-
-        except Exception as e:
-            logger.error(f"Failed to extract triplets for {instance_id}: {e}")
-
-    logger.info(f"Finished extracting spans and traces for instance: {instance_id}")
-
-    # Quickly sanity check the spans
-    sanity_check_spans(spans)
-    logger.info(f"Sanity check passed for instance: {instance_id}")
-
-
-async def dry_run_claude_code(
-    *,
-    dataset_path: str,
-    haiku_frontend_name: str,
-    haiku_backend_name: str,
-    sonnet_frontend_name: str,
-    sonnet_backend_name: str,
-    backend_type: Literal["vllm", "anthropic", "openai"],
-    api_base_url: Optional[str],
-    output_dir: Optional[str],
-    max_turns: int,
-    limit: Optional[int],
-    cooldown_seconds: float,
-) -> None:
-    """Executes a dry run of the Claude Code agent on a dataset.
-
-    This function handles both 'official' runs (interacting with Anthropic APIs)
-    and 'hosted' runs (interacting with vLLM or compatible servers). It manages
-    initialization of the Lightning Store, LLM Proxy, and the execution loop.
-
-    If running in 'vllm' mode, it will also attempt to extract triplets using
-    the provided backend name as the tokenizer path and save a HuggingFace Dataset.
-
-    Args:
-        dataset_path: Path to the JSONL dataset file.
-        haiku_frontend_name: The model name used in the code to request the 'fast' model.
-        haiku_backend_name: The actual model name/path on the backend.
-        sonnet_frontend_name: The model name used in the code to request the 'strong' model.
-        sonnet_backend_name: The actual model name/path on the backend.
-        backend_type: The type of backend to configure ("vllm", "anthropic" or "openai").
-        api_base_url: Base URL for the API. Required for "vllm" or "openai".
-        output_dir: Directory to save logs, spans, and datasets.
-        max_turns: Maximum number of steps the agent can take per instance.
-        limit: Optional limit on the number of instances to process.
-    """
-    dataset = _load_dataset(dataset_path, limit=limit)
-
-    # Initialize Infrastructure
-    tracer = OtelTracer()
-    runner = LitAgentRunner[SWEbenchInstance](tracer)
-    store = LightningStoreServer(InMemoryLightningStore(), host="0.0.0.0", port=7654)
-    await store.start()
-
-    # Enable callbacks for training data extraction if using vLLM
-    callbacks = ["return_token_ids", "opentelemetry", "logprobs"] if backend_type == "vllm" else ["opentelemetry"]
-    llm_proxy = LLMProxy(port=12358, store=store, callbacks=callbacks)
-
-    # Configure Models based on backend type
-    model_configs: List[ModelConfig] = []
-    model_params: Dict[str, Any] = {}
-
-    if backend_type == "vllm":
-        model_namespace = "hosted_vllm"
-        if api_base_url:
-            model_params["api_base"] = api_base_url
-        else:
-            raise ValueError("api_base_url is required for vllm backend")
-    elif backend_type == "anthropic":
-        model_namespace = "anthropic"
-        model_params["api_key"] = "os.environ/ANTHROPIC_API_KEY"
-        if api_base_url:
-            model_params["api_base"] = api_base_url
-    elif backend_type == "openai":
-        model_namespace = "openai"
-        model_params["api_key"] = "os.environ/OPENAI_API_KEY"
-        if api_base_url:
-            # Users can still override this via environment variables,
-            # even if they don't pass it in as an argument.
-            model_params["api_base"] = api_base_url
-
-    model_configs.extend(
-        [
-            ModelConfig(
-                model_name=sonnet_frontend_name,
-                litellm_params={
-                    "model": f"{model_namespace}/{sonnet_backend_name}",
-                    **model_params,
-                },
-            ),
-            ModelConfig(
-                model_name=haiku_frontend_name,
-                litellm_params={
-                    "model": f"{model_namespace}/{haiku_backend_name}",
-                    **model_params,
-                },
-            ),
-        ]
-    )
-
-    logger.info(f"Updating model list: {model_configs}")
-
-    llm_proxy.update_model_list(model_configs)
-    await llm_proxy.start()
-
-    try:
-        # Add the LLM proxy as a resource to the store
-        await store.add_resources({"llm": llm_proxy.as_resource(model="local")})
-
-        # Prepare for triplet extraction if vllm
-        adapter = ExtendedLlmProxyTraceToTriplet() if backend_type == "vllm" else None
-        tokenizer = None
-        if backend_type == "vllm":
-            try:
-                tokenizer = AutoTokenizer.from_pretrained(sonnet_backend_name)  # type: ignore
-            except Exception as e:
-                logger.warning(f"Could not load tokenizer for {sonnet_backend_name}: {e}")
-
-        # Load full swebench dataset. Mainly for evaluation purposes.
-        swebench_full_dataset = load_swebench_dataset("princeton-nlp/SWE-bench", split="test")
-
-        # Initialize Claude Code Agent
-        claude_code_agent = ClaudeCodeAgent(swebench_full_dataset=swebench_full_dataset, max_turns=max_turns)
-
-        # Execution Loop
-        for instance in dataset:
-            await run_instance_async(
-                instance,
-                claude_code_agent,
-                runner,
-                store,
-                output_dir,
-                adapter,
-                cast(PreTrainedTokenizerBase, tokenizer),
-            )
-            # Basic sleep to allow resource cleanup or rate limit cooling
-            await asyncio.sleep(cooldown_seconds)
-
-    finally:
-        await llm_proxy.stop()
-        await store.stop()
-
-
-if __name__ == "__main__":
-    parser = ArgumentParser(description="Run Claude Code Agent experiments.")
-
-    # Backend Selection
-    parser.add_argument(
-        "backend_type",
-        type=str,
-        choices=["vllm", "anthropic", "openai"],
-        help="Backend type: 'vllm' for hosted models, 'anthropic' for official API, 'openai' for OpenAI API.",
-    )
-
-    # Model Configuration
-    parser.add_argument(
-        "--backend-model-high",
-        type=str,
-        default=None,
-        help="Backend model path/name for expensive model usages (used as vLLM model name / OpenAI model name).",
-    )
-    parser.add_argument(
-        "--backend-model-low",
-        type=str,
-        default=None,
-        help="Backend model path/name for low-price model usages (used as vLLM model name / OpenAI model name).",
-    )
-    parser.add_argument(
-        "--base-url", type=str, default="http://localhost:8000/v1", help="LLM server address (required for vllm)."
-    )
-
-    # Frontend/Agent Configuration
-    parser.add_argument(
-        "--frontend-model-high",
-        type=str,
-        default="claude-sonnet-4-5-20250929",
-        help="The frontend high-price model name provided to Claude Code.",
-    )
-    parser.add_argument(
-        "--frontend-model-low",
-        type=str,
-        default="claude-haiku-4-5-20251001",
-        help="The frontend low-price model name provided to Claude Code.",
-    )
-
-    # Execution Configuration
-    parser.add_argument("--dataset-path", type=str, default="swebench_samples.jsonl", help="Path to the dataset.")
-    parser.add_argument("--max-turns", type=int, default=5, help="Maximum turns per instance.")
-    parser.add_argument("--output-dir", type=str, default="data", help="Directory to save output logs.")
-    parser.add_argument("--limit", type=int, default=None, help="Limit the number of instances to run (for debugging).")
-    parser.add_argument("--cooldown-seconds", type=float, default=2.0, help="Cooldown seconds between instances.")
-    parser.add_argument("--debug", action="store_true", help="Enable debug loggings.")
-
-    args = parser.parse_args()
-
-    if args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
-    if args.debug:
-        setup_logging()
-        setup_module_logging("DEBUG", name="claude_code_agent")
-    else:
-        setup_logging(apply_to=[logger.name])
-
-    # Map backend_type to the appropriate args
-    backend_mode = cast(Literal["vllm", "anthropic", "openai"], args.backend_type)
-
-    # If using anthropic, the backend name usually matches the frontend or is specific API string.
-    # Otherwise, the backend name is the model name/path (e.g., Qwen/...) and must be provided.
-    if args.backend_model_high is None:
-        if args.backend_type == "anthropic":
-            backend_model_high = args.frontend_model_high
-        else:
-            raise ValueError("--backend-model-high is required for non-anthropic backends")
-    else:
-        backend_model_high = args.backend_model_high
-
-    if args.backend_model_low is None:
-        if args.backend_type == "anthropic":
-            backend_model_low = args.frontend_model_low
-        else:
-            raise ValueError("--backend-model-low is required for non-anthropic backends")
-    else:
-        backend_model_low = args.backend_model_low
-
-    asyncio.run(
-        dry_run_claude_code(
-            dataset_path=args.dataset_path,
-            haiku_frontend_name=args.frontend_model_low,
-            haiku_backend_name=backend_model_low,
-            sonnet_frontend_name=args.frontend_model_high,
-            sonnet_backend_name=backend_model_high,
-            backend_type=backend_mode,
-            api_base_url=args.base_url if backend_mode == "vllm" else None,
-            output_dir=args.output_dir,
-            max_turns=args.max_turns,
-            limit=args.limit,
-            cooldown_seconds=args.cooldown_seconds,
-        )
-    )

examples/claude_code/claude_code_controller.py

@@ -1,227 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Controller module for managing Claude Code executions in containerized environments.
-
-This module provides the ClaudeController class that manages the execution of Claude Code
-within Docker containers. It handles container initialization, command execution, and
-patch application for SWE-bench evaluation tasks.
-"""
-
-import logging
-from functools import partial
-from typing import Literal, TypedDict
-
-import dotenv
-from swebench.harness.constants import SWEbenchInstance
-from swebench_utils.docker_runtime import Runtime
-from swebench_utils.logging import log_for_evaluation
-
-SWEBENCH_EXTRA_SYSTEM_PROMPT = """
- You are an expert software engineer solving swebench bug fixing tasks.
-"""
-
-SWEBENCH_USER_PROMPT = """
-You are given a code repository in the current directory (/testbed).
-The bug description is:
-{description}
-=================================================
-You task is to fix the bug with the following steps:
-(1) write test cases to reproduce the bug.
-(2) explore the source codes to locate the bug.
-(3) edit the source codes to fix the bug.
-(4) rerun your written test cases to validate that the bug is fixed. If not, go back to explore the source codes and fix the codes again.
-(5) remember to delete the test cases you write at last.
-Please do not commit your edits. We will do it later.
-"""
-
-logger = logging.getLogger("claude_code_agent")
-
-
-class RunInstanceResult(TypedDict):
-    instance_id: str
-    model_patch: str
-    model_name_or_path: str
-
-
-class ClaudeController:
-    """Manages the execution of Claude Code within a Docker runtime.
-
-    This controller handles the lifecycle of a SWE-bench task execution, including
-    environment setup, tool installation, agent execution (via CLI or Python SDK),
-    and result extraction.
-
-    Attributes:
-        container: The active Docker runtime session.
-    """
-
-    def __init__(self, image: str, instance: SWEbenchInstance, run_id: str, endpoint: str, api_key: str) -> None:
-        """Initialize the ClaudeController.
-
-        Args:
-            image: The Docker image tag.
-            instance: The dataset instance containing the problem statement and ID.
-            run_id: The identifier for the evaluation run.
-            endpoint: The API endpoint URL.
-            api_key: The API authentication key.
-        """
-        self.image = image
-        self.instance = instance
-        self.run_id = run_id
-        self.endpoint = endpoint
-        self.api_key = api_key
-        self.container: Runtime = self.init_container(self.image, self.instance)
-
-    def init_container(self, image: str, instance: SWEbenchInstance) -> Runtime:
-        """Initializes the Docker container and sets up the Claude Code environment.
-
-        This method starts the container session, installs the Claude CLI,
-        configures environment variables for authentication and sandbox mode.
-
-        Args:
-            image: The Docker image tag to start.
-            instance: The dataset instance to load into the environment.
-
-        Returns:
-            An initialized and configured Docker runtime object.
-        """
-        container = Runtime.start_session(
-            image,
-            instance,
-            log_function=partial(log_for_evaluation, run_id=self.run_id, instance_id=instance["instance_id"]),
-        )
-        # Install Claude CLI
-        container.send_command("curl -fsSL https://claude.ai/install.sh | bash")
-        container.send_command('alias claude="$HOME/.local/bin/claude"')
-
-        # Configure Environment
-        dotenv.load_dotenv()
-        container.send_command(f"export ANTHROPIC_BASE_URL={self.endpoint}")
-        container.send_command(f"export ANTHROPIC_AUTH_TOKEN={self.api_key}")
-        container.send_command("export IS_SANDBOX=1")
-
-        return container
-
-    def _run_cli(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
-        """Executes Claude Code using the Command Line Interface.
-
-        Constructs a safe heredoc for the prompt to avoid shell interpolation issues
-        and executes the `claude` binary directly.
-
-        Args:
-            instance: The problem instance containing the problem statement.
-            max_turns: The maximum number of interaction turns allowed.
-            time_limit: The execution time limit in minutes.
-        """
-        # Prepare prompt safely: write it to a file inside the container using a single-quoted heredoc
-        # directly applying prompt for heredoc may raise error for windows line ending \r\n
-        prompt_text = SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
-
-        # Choose a simple filename and a heredoc delimiter unlikely to collide
-        heredoc_cmd = "cat > /tmp/cc_prompt.txt <<'CC_PROMPT'\n" + prompt_text + "\nCC_PROMPT\n"
-        self.container.send_command(heredoc_cmd)
-
-        # Run claude reading the prompt from the file
-        claude_cmd = (
-            f'claude -p "$(cat /tmp/cc_prompt.txt)" '
-            f'--append-system-prompt "{SWEBENCH_EXTRA_SYSTEM_PROMPT}" '
-            f"--max-turns {max_turns} "
-            f"--dangerously-skip-permissions "
-            f"--output-format json --verbose"
-        )
-        logger.info(f"Running Claude Code CLI command: {claude_cmd}")
-        self.container.send_command(claude_cmd, time_limit * 60)
-        logger.info(f"Claude Code CLI command completed")
-
-    def _run_python_sdk(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
-        """Executes Claude Code using the Python SDK wrapper.
-
-        Installs the Python SDK if necessary, hydrates a template script with the
-        problem prompt, and executes the generated Python script.
-
-        Note:
-            This path is still under development and not yet stable.
-
-        Args:
-            instance: The problem instance containing the problem statement.
-            max_turns: The maximum number of interaction turns allowed.
-            time_limit: The execution time limit in minutes.
-        """
-        # Ensure Python 3.12 is available
-        self.container.send_command(
-            f"""
-if ! command -v python3 &> /dev/null; then
-    echo "Python is not installed. Installing Python 3.12..."
-    sudo apt-get update -qq && sudo apt-get install -y -qq python3.12
-else
-    echo "Python is already installed."
-fi
-"""
-        )
-        self.container.send_command("python3 -m pip install claude-code-sdk")
-
-        # Load and fill the execution template
-        with open("src/agent/cc/claude_code_main.py.template") as f:
-            entrance_template = f.read()
-
-        script_content = (
-            entrance_template.replace("SYS_PROMPT", SWEBENCH_EXTRA_SYSTEM_PROMPT)
-            .replace(
-                "PROMPT", SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
-            )
-            .replace("MAX_STEP", str(max_turns))
-        )
-
-        # Write the script to the container and execute
-        self.container.send_command(f"cat > /tmp/claude_code_main.py <<'CC_MAIN'\n{script_content}\nCC_MAIN\n")
-        self.container.send_command("python3 /tmp/claude_code_main.py", time_limit * 60)
-        return
-
-    def run_instance(
-        self,
-        instance: SWEbenchInstance,
-        max_turns: int = 40,
-        time_limit: int = 30,
-        run_method: Literal["python", "cli"] = "python",
-    ) -> RunInstanceResult:
-        """Runs the agent on a specific SWE-bench instance.
-
-        This method orchestrates the agent execution via the specified method (CLI or Python),
-        and extracts the generated git diff (patch) upon completion.
-
-        Args:
-            instance: The dataset instance dictionary.
-            max_turns: Maximum conversation turns allowed for the agent. Defaults to 40.
-            time_limit: Time limit for the execution in minutes. Defaults to 30.
-            run_method: The execution method, either "python" (SDK) or "cli". Defaults to "python".
-
-        Returns:
-            A dictionary containing the result:
-
-            - instance_id: The ID of the processed instance.
-            - model_patch: The git diff generated by the agent.
-            - model_name_or_path: Hardcoded to "cc" (Claude Code).
-
-        Raises:
-            ValueError: If `run_method` is not "python" or "cli".
-        """
-        if run_method == "python":
-            logger.warning("Running Claude Code using Python SDK is still under development and not yet stable.")
-            self._run_python_sdk(instance, max_turns, time_limit)
-        elif run_method == "cli":
-            self._run_cli(instance, max_turns, time_limit)
-        else:
-            raise ValueError(f"Wrong run_method '{run_method}', run_method should be in ['python', 'cli']")
-
-        result = self.container.send_command("git --no-pager diff HEAD")
-        git_diff = result.output.replace("git --no-pager diff HEAD\n", "")
-
-        return {
-            "instance_id": instance["instance_id"],
-            "model_patch": git_diff,
-            "model_name_or_path": "cc",
-        }
-
-    def __del__(self) -> None:
-        """Destructor to ensure container resources are cleaned up."""
-        if hasattr(self, "container"):
-            self.container.cleanup()

examples/claude_code/extended_adapter.py

@@ -1,163 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Custom adapter module for converting LLM proxy traces to augmented trajectories.
-
-This module provides an augmented LlmProxyTraceToTriplet adapter that converts
-LLM proxy spans into augmented trajectories for analysis and evaluation.
-It extends the base LlmProxyTraceToTriplet to include additional metadata like chat messages,
-log probabilities, and sequence IDs.
-"""
-
-import logging
-from typing import Any, Dict, List, Optional, Tuple, cast
-
-from agentlightning.adapter.triplet import LlmProxyTraceToTriplet
-from agentlightning.types import Span, Triplet
-
-logger = logging.getLogger(__name__)
-
-
-class ExtendedLlmProxyTraceToTriplet(LlmProxyTraceToTriplet):
-    """Convert LLM Proxy spans into trajectories with logprobs and customized metadata.
-
-    Augmented fields include:
-
-    - chat messages history from [`llm.hosted_vllm.messages`], saved to `Triplet.metadata['messages']`
-    - logprobs from [`llm.hosted_vllm.choices`], saved to `Triplet.response['logprobs']`
-    - sequence_id from [`Span.sequence_id`] to locate the order of the span (conversation turn), saved to `Triplet.metadata['sequence_id']`
-    """
-
-    def _extract_tokens_from_raw(self, attrs: Dict[str, Any]) -> Tuple[List[int], List[int], List[float]]:  # type: ignore
-        """Extract token ids from raw_gen_ai_request attributes.
-
-        - llm.hosted_vllm.prompt_token_ids: string -> List[int]
-        - llm.hosted_vllm.choices: string -> [{'token_ids': [...]}] -> take first
-        """
-        prompt_ids: List[int] = []
-        resp_ids: List[int] = []
-        logprobs: List[float] = []
-
-        # prompt
-        p = attrs.get("llm.hosted_vllm.prompt_token_ids")
-        p = self._literal_eval_maybe(p)
-        if isinstance(p, list) and all(isinstance(x, int) for x in p):  # type: ignore
-            prompt_ids = cast(List[int], p)
-
-        choices = attrs.get("llm.hosted_vllm.choices")
-        choices = self._literal_eval_maybe(choices)
-        if isinstance(choices, list) and choices:
-            cand = cast(Any, choices[0])
-            if isinstance(cand, dict):
-                tids = cast(Dict[str, Any], cand).get("token_ids")
-                if isinstance(tids, list) and all(isinstance(x, int) for x in tids):  # type: ignore
-                    resp_ids = cast(List[int], tids)
-
-                if "logprobs" in cand:
-                    logprobs_dict = cast(Dict[str, Any], cand).get("logprobs")
-                    if isinstance(logprobs_dict, dict) and "content" in logprobs_dict:
-                        content = cast(List[Dict[str, Any]], logprobs_dict["content"])
-                        logprobs = [float(item["logprob"]) for item in content if "logprob" in item]
-
-        return prompt_ids, resp_ids, logprobs
-
-    def adapt(self, source: List[Span], /) -> List[Triplet]:  # type: ignore
-        """Convert LLM Proxy spans into [`Triplet`][agentlightning.Triplet] trajectories.
-
-        Args:
-            source: Spans emitted by the LLM Proxy containing prompt, response, and reward data.
-
-        Returns:
-            Ordered trajectory transitions matched purely by `sequence_id`.
-        """
-        # 1) Sort deterministically by (sequence_id, start_time).
-        spans = sorted(
-            source,
-            key=lambda s: (s.sequence_id, s.start_time),
-        )
-
-        # 2) Collect LLM calls
-        llm_items: List[Dict[str, Any]] = []
-        seen_request_ids: set[str] = set()
-        for s in spans:
-            attrs = s.attributes or {}
-            prompt_ids: List[int] = []
-            resp_ids: List[int] = []
-            logprobs: List[float] = []
-
-            if s.name == "raw_gen_ai_request":
-                prompt_ids, resp_ids, logprobs = self._extract_tokens_from_raw(attrs)
-
-                if len(prompt_ids) == 0 or len(resp_ids) == 0:
-                    logger.warning(
-                        f"Span {s.span_id} is missing prompt (len={len(prompt_ids)}) or response (len={len(resp_ids)}) token ids. Ignoring this span."
-                    )
-                    continue
-                elif len(logprobs) == 0:
-                    logger.warning(f"Span {s.span_id} is missing logprobs. Ignoring logprobs for this span.")
-                    continue
-                elif len(resp_ids) != len(logprobs):
-                    logger.warning(
-                        f"Span {s.span_id} has mismatched response ids and logprobs lengths: "
-                        f"{len(resp_ids)} vs {len(logprobs)}. Ignoring this span."
-                    )
-                    continue
-
-            if prompt_ids and resp_ids and logprobs:
-                rid = self._request_id_from_attrs(attrs)
-                if rid:
-                    # Duplicated request ID. This request is already handled.
-                    if rid in seen_request_ids:
-                        continue
-                    seen_request_ids.add(rid)
-                llm_items.append(
-                    dict(
-                        span=s,
-                        seq=s.sequence_id,
-                        response_ids=resp_ids,
-                        prompt_ids=prompt_ids,
-                        request_id=rid,
-                        logprobs=logprobs,
-                    )
-                )
-
-        # Order LLM items by sequence only.
-        llm_items.sort(key=lambda x: x["seq"])
-
-        # Collect rewards by sequence only.
-        rewards: List[Tuple[int, Optional[float]]] = []
-        for s in spans:
-            val = self._maybe_reward_value(s)
-            if val is not None:
-                rewards.append((s.sequence_id, val))
-
-        # First-occurrence matching by sequence_id only:
-        # For reward at sequence R, assign to the most recent unmatched LLM with seq < R.
-        assigned: Dict[str, Optional[float]] = {}
-        for r_seq, r_val in sorted(rewards, key=lambda x: x[0]):
-            for item in reversed(llm_items):
-                sid = item["span"].span_id
-                if sid in assigned:
-                    continue
-                if item["seq"] < r_seq:
-                    assigned[sid] = r_val
-                    break
-
-        # Build triplets in LLM sequence order.
-        triplets: List[Triplet] = []
-        for item in llm_items:
-            s = item["span"]
-            triplets.append(
-                Triplet(
-                    prompt={"token_ids": item["prompt_ids"]},
-                    response={"token_ids": item["response_ids"], "logprobs": item["logprobs"]},
-                    reward=assigned.get(s.span_id, None),
-                    metadata=dict(
-                        # This is called response_id to align with the other adapters.
-                        response_id=item["request_id"],
-                        sequence_id=item["seq"],
-                        messages=self._literal_eval_maybe(s.attributes.get("llm.hosted_vllm.messages")),
-                    ),
-                )
-            )
-
-        return triplets

examples/claude_code/swebench_utils/__init__.py

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

examples/claude_code/swebench_utils/docker_runtime.py

@@ -1,430 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Docker runtime management for repository setup and command execution.
-
-Provides containerized environment for repository testing with command execution,
-file operations, and state management capabilities.
-"""
-
-from __future__ import annotations
-
-import json
-import logging
-import queue
-import re
-import threading
-import time
-import uuid
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional
-
-from docker.errors import DockerException, ImageNotFound
-from docker.models.containers import Container
-from swebench.harness.constants import SWEbenchInstance
-from typing_extensions import Self
-
-import docker
-
-# This will log to the console for debugging purposes.
-claude_code_logger = logging.getLogger("claude_code_agent.docker_runtime")
-
-CMD_OUTPUT_PS1_BEGIN = "\n###PS1JSON###\n"
-CMD_OUTPUT_PS1_END = "\n###PS1END###"
-CMD_OUTPUT_METADATA_PS1_REGEX = re.compile(
-    r"(?m)^\s*" + re.escape(CMD_OUTPUT_PS1_BEGIN.strip()) + r"\s*(.*?)\s*" + re.escape(CMD_OUTPUT_PS1_END.strip()),
-    re.DOTALL,
-)
-ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
-
-TIMEOUT_EXIT_CODE = 124
-
-MEM_LIMIT = "8g"
-CPU_CORES = 4
-
-
-VAR_PATTERNS = {
-    "exit_code": re.compile(r'"exit_code":\s*(-?\d+)\s*(?:,|\})'),
-    "username": re.compile(r'"username":\s*"([^"]*)"'),
-    "hostname": re.compile(r'"hostname":\s*"([^"]*)"'),
-    "working_dir": re.compile(r'"working_dir":\s*"([^"]*)"'),
-    "py_interpreter_path": re.compile(r'"py_interpreter_path":\s*"([^"]*)"'),
-}
-
-
-@dataclass
-class CmdOutputMetadata:
-    """
-    Additional metadata captured from PS1 shell prompt.
-
-    Provides context about command execution environment including
-    exit codes, user info, working directory, and Python interpreter.
-    """
-
-    exit_code: int = -1
-    username: str | None = None
-    hostname: str | None = None
-    working_dir: str | None = None
-    py_interpreter_path: str | None = None
-
-    @classmethod
-    def matches_ps1_metadata(cls, output: str) -> List[re.Match[str]]:
-        matches: List[re.Match[str]] = []
-        for match in CMD_OUTPUT_METADATA_PS1_REGEX.finditer(output):
-            scope = match.group(1).strip()
-            try:
-                d = json.loads(scope)  # Try to parse as JSON
-                matches.append(match)
-            except json.JSONDecodeError:
-                d = cls.best_effort_match(scope)
-                if len(d) > 0:
-                    matches.append(match)
-        return matches
-
-    @classmethod
-    def best_effort_match(cls, scope: str) -> Dict[str, Any]:
-        out: Dict[str, str] = {}
-        for field, pattern in VAR_PATTERNS.items():
-            m = pattern.search(scope)
-            if m:
-                out[field] = m.group(1)
-            else:
-                out[field] = ""
-        return out
-
-    @classmethod
-    def from_ps1_match(cls, match: re.Match[str]) -> Self:
-        """
-        Extract metadata from a PS1 prompt regex match.
-
-        Args:
-            match (re.Match[str]): Regex match containing JSON metadata
-
-        Returns:
-            Self: CmdOutputMetadata instance with parsed values
-        """
-        try:
-            metadata = json.loads(match.group(1))
-        except:
-            metadata = cls.best_effort_match(match.group(1))
-        # Create a copy of metadata to avoid modifying the original
-        processed = metadata.copy()
-        # Convert numeric fields
-        if "exit_code" in metadata:
-            try:
-                processed["exit_code"] = int(float(str(metadata["exit_code"])))
-            except (ValueError, TypeError):
-                processed["exit_code"] = -1
-        return cls(**processed)
-
-
-@dataclass
-class CommandResult:
-    """
-    Result of a command execution with output and metadata.
-
-    Attributes:
-        output (str): Command output text
-        metadata (Optional[CmdOutputMetadata]): Execution context metadata
-    """
-
-    output: str
-    metadata: Optional[CmdOutputMetadata]
-
-    def to_observation(self, strip: bool = True) -> str:
-        """
-        Convert command result to formatted observation string.
-
-        Args:
-            strip (bool): Whether to truncate long output
-
-        Returns:
-            str: Formatted observation with output and context
-        """
-        # compile regex once for efficiency
-        ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
-
-        output = ANSI_ESCAPE.sub("", self.output).replace("\r", "")
-
-        if len(output) > 1024 * 8 and strip:
-            output = output[: 1024 * 4] + "....stripped due to length....\n" + output[-1024 * 4 :]
-
-        if self.metadata is None:
-            return f"\n{output}\n"
-        return f"""{output}
-{self.metadata.username}@{self.metadata.hostname}:{self.metadata.working_dir} $
-
-exit code: {self.metadata.exit_code}
-"""
-
-
-class Runtime:
-    """
-    Docker container runtime for repository setup and testing.
-
-    Manages a Docker container with persistent bash session, command execution,
-    file operations, and container lifecycle management.
-    """
-
-    def __init__(self, container: Container, log_function: Callable[..., None]) -> None:
-        """
-        Initialize runtime with an existing Docker container.
-
-        Args:
-            container (Container): Docker container instance to manage
-        """
-        self.container = container
-        self.logger = log_function  # Set logger early so it's available even if init fails later
-        self.sock: Any = self.container.attach_socket(params={"stdin": 1, "stdout": 1, "stderr": 1, "stream": 1})  # type: ignore
-        self.output_queue: queue.Queue[bytes] = queue.Queue()
-        self._start_output_thread()
-        self._clear_initial_prompt()
-
-        json_str = json.dumps(
-            {
-                "exit_code": "$?",
-                "username": r"\u",
-                "hostname": r"\h",
-                "working_dir": r"$(pwd)",
-                "py_interpreter_path": r'$(which python 2>/dev/null || echo "")',
-            },
-            indent=2,
-        ).replace('"', r"\"")
-        ps1 = CMD_OUTPUT_PS1_BEGIN + json_str + CMD_OUTPUT_PS1_END + "\n"
-        self.send_command(f'export PROMPT_COMMAND=\'export PS1="{ps1}"\'; export PS2=""')
-        self.send_command("apt update -qq && apt install -y -qq git")
-        self.stopped = False
-
-    def _stream_output(self):
-        while True:
-            try:
-                output = self._recv_bytes(4096)
-                if not output:
-                    break
-                self.output_queue.put(output)
-            except (OSError, ConnectionError) as e:
-                print(f"Connection error in _stream_output: {e}")
-                break
-            except Exception as e:
-                # print(f"Unexpected error in _stream_output: {e}")
-                break
-
-    def _start_output_thread(self):
-        self.output_thread = threading.Thread(target=self._stream_output, daemon=True)
-        self.output_thread.start()
-        # TODO: kill the thread if main thread is stopped
-
-    def _clear_initial_prompt(self):
-        time.sleep(0.5)
-        while not self.output_queue.empty():
-            self.output_queue.get()
-
-    def _read_raw_output(self, timeout: float = 30) -> tuple[str, Optional[CmdOutputMetadata]]:
-        accumulated_output = ""
-        start_time = time.time()
-
-        while time.time() - start_time < timeout:
-            try:
-                chunk = self.output_queue.get(timeout=0.1)
-                accumulated_output += chunk.decode("utf-8", errors="ignore")
-                # PSReadLine injects ANSI + cursor control; normalize before matching
-                accumulated_clean = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
-                ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_clean)
-                if ps1_matches:
-                    break
-            except queue.Empty:
-                continue
-        accumulated_output = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
-        ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_output)
-        metadata = CmdOutputMetadata.from_ps1_match(ps1_matches[-1]) if ps1_matches else None
-        output = self._combine_outputs_between_matches(
-            accumulated_output,
-            ps1_matches,
-        )
-        return output, metadata
-
-    def _combine_outputs_between_matches(self, pane_content: str, ps1_matches: list[re.Match[str]]) -> str:
-        if len(ps1_matches) == 1:
-            return pane_content[: ps1_matches[0].start()]
-        elif len(ps1_matches) == 0:
-            return pane_content
-        output_segments: List[str] = []
-        for i in range(len(ps1_matches) - 1):
-            output_segment = pane_content[ps1_matches[i].end() + 1 : ps1_matches[i + 1].start()]
-            output_segments.append(output_segment)
-        return "\n".join(output_segments) + "\n" if output_segments else ""
-
-    def _recv_bytes(self, n: int = 4096) -> bytes:
-        # Prefer the public API on whatever object the SDK returns
-        for m in ("recv", "read"):
-            if hasattr(self.sock, m):
-                return getattr(self.sock, m)(n)
-        # Last-resort fallback for odd wrappers that still expose ._sock
-        if hasattr(self.sock, "_sock"):
-            for m in ("recv", "read"):
-                if hasattr(self.sock._sock, m):
-                    return getattr(self.sock._sock, m)(n)
-        raise TypeError(f"Don't know how to read from {type(self.sock).__name__}")
-
-    def _send_bytes(self, data: bytes) -> None:
-        if hasattr(self.sock, "_sock"):
-            for m in ("send", "sendall", "write"):
-                if hasattr(self.sock._sock, m):
-                    getattr(self.sock._sock, m)(data)
-                    return
-        for m in ("send", "sendall", "write"):
-            if hasattr(self.sock, m):
-                getattr(self.sock, m)(data)
-                return
-
-        raise TypeError(f"Don't know how to write to {type(self.sock).__name__}")
-
-    def _log_command_result(self, result: CommandResult) -> None:
-        claude_code_logger.debug("Docker runtime command finished with metadata: %s", result.metadata)
-        if len(result.output) > 2048:
-            logged_output = result.output[:1024] + "\n(... stripped due to length ...)\n" + result.output[-1024:]
-        else:
-            logged_output = result.output
-        claude_code_logger.debug(
-            "Docker runtime command finished with output (length = %d):\n%s", len(result.output), logged_output
-        )
-        # Output to the evaluation logger simultaneously
-        self.logger(text=logged_output)
-
-    def send_command(self, command: str, timeout: float = 20 * 60) -> CommandResult:
-        # Redact sensitive API keys from the command before logging
-        redacted_command = command
-        for sensitive_var in ["ANTHROPIC_AUTH_TOKEN", "API_KEY", "SECRET_KEY"]:
-            pattern = rf"(export (.*?){re.escape(sensitive_var)}(.*?)=)[^\s]+"
-            redacted_command = re.sub(pattern, rf"\1****REDACTED****", redacted_command)
-        claude_code_logger.info("Docker runtime receiving command: %s", redacted_command)
-        # Normalize newline semantics for interactive shells
-        if not command.endswith("\n"):
-            command += "\n"
-
-        while not self.output_queue.empty():
-            self.output_queue.get()
-
-        self._send_bytes(command.encode())
-
-        output, metadata = self._read_raw_output(timeout=timeout)
-        # TODO: Check exit code of the command (claude code download fail will not be caught by this)
-        if metadata is not None:
-            result = CommandResult(output=output, metadata=metadata)
-            self._log_command_result(result)
-            return result
-
-        # handle timeout
-        self._send_bytes(b"\x03")
-
-        kill_timeout = 5.0
-        kill_output, kill_metadata = self._read_raw_output(timeout=kill_timeout)
-
-        output = output + kill_output + "\n**Exited due to timeout**\n"
-        if kill_metadata is not None:
-            kill_metadata.exit_code = TIMEOUT_EXIT_CODE
-            result = CommandResult(output=output, metadata=kill_metadata)
-            self._log_command_result(result)
-            return result
-
-        fallback_metadata = CmdOutputMetadata(
-            exit_code=TIMEOUT_EXIT_CODE,
-        )
-        result = CommandResult(output=output, metadata=fallback_metadata)
-        self._log_command_result(result)
-        return result
-
-    def cleanup(self) -> None:
-        if self.stopped:
-            return
-        try:
-            claude_code_logger.info(f"Stopping container: {self.container.id}")
-            self.container.stop()
-            claude_code_logger.info(f"Removing container: {self.container.id}")
-            self.container.remove(force=True)
-            claude_code_logger.info(f"Container removed: {self.container.id}")
-            self.stopped = True
-        except Exception as e:
-            print(f"Failed to stop container: {e}")
-
-    def __del__(self):
-        self.cleanup()
-
-    @staticmethod
-    def pull_image(image_name: str) -> bool:
-        """
-        Pull Docker image from registry.
-
-        Args:
-            image_name (str): Name of the Docker image to pull
-
-        Returns:
-            bool: True if successful, False if image not found
-        """
-        client = docker.from_env()
-        try:
-            client.images.pull(image_name)
-            return True
-        except ImageNotFound:
-            return False
-
-    @classmethod
-    def start_session(
-        cls,
-        image_name: str,
-        instance: SWEbenchInstance,
-        log_function: Callable[..., None] = lambda: None,
-    ) -> Runtime:
-        """
-        Start a Docker container session for repository testing.
-
-        Args:
-            image_name (str): Base Docker image name
-            instance (dict): SWE-bench instance data with repo info
-
-        Returns:
-            SetupRuntime: Configured runtime session ready for command execution
-
-        Raises:
-            RuntimeError: If Docker is not available
-        """
-        try:
-            docker.from_env().ping()  # type: ignore
-        except DockerException:
-            raise RuntimeError("Docker is not installed or not running.")
-
-        _ = cls.pull_image(image_name)
-        client = docker.from_env(timeout=600)
-        container_id = instance["instance_id"]
-        container_name = f"git-launch-{container_id}-{str(uuid.uuid4())[:4]}"
-        info: Dict[str, str] = client.version()  # type: ignore
-        engine_os: str = (info.get("Os") or info.get("OSType") or "").lower()  # type: ignore
-        # which operating system this code is running on, note windows can run linux containers, so engine_os != (container) platform
-        extra_hosts = {"host.docker.internal": "host-gateway"} if "linux" in engine_os else None
-
-        shell_command = "/bin/bash"
-        working_dir = "/testbed"
-
-        claude_code_logger.info(
-            f"Starting container {container_name} with image {image_name}. Shell command: {shell_command}"
-        )
-        container = client.containers.run(
-            image_name,
-            name=container_name,
-            command=shell_command,
-            stdin_open=True,
-            tty=True,
-            detach=True,
-            environment={
-                "TERM": "xterm-mono",
-            },
-            working_dir=working_dir,
-            extra_hosts=extra_hosts,
-            network_mode="host",
-            cpu_quota=int(CPU_CORES * 100000),
-            mem_limit=MEM_LIMIT,
-        )
-        claude_code_logger.info(f"Container {container_name} started with ID: {container.id}")
-
-        session = cls(container, log_function=log_function)
-
-        return session

examples/claude_code/swebench_utils/evaluation.py

@@ -1,258 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Evaluation module for SWE-bench instance testing and grading.
-
-This module provides core functionality for evaluating model predictions on SWE-bench
-instances. It handles containerized execution of test scripts, patch application,
-and generation of evaluation reports. The module orchestrates the complete evaluation
-process including container management, patch application, test execution, and result
-grading.
-"""
-
-import json
-import traceback
-from pathlib import Path, PurePosixPath
-from typing import Any, Dict, Optional
-
-from docker.models.containers import ExecResult
-from swebench.harness.constants import (
-    APPLY_PATCH_FAIL,
-    APPLY_PATCH_PASS,
-    DOCKER_PATCH,
-    DOCKER_USER,
-    DOCKER_WORKDIR,
-    INSTANCE_IMAGE_BUILD_DIR,
-    KEY_MODEL,
-    KEY_PREDICTION,
-    LOG_INSTANCE,
-    LOG_REPORT,
-    LOG_TEST_OUTPUT,
-    RUN_EVALUATION_LOG_DIR,
-    UTF8,
-    SWEbenchInstance,
-)
-from swebench.harness.docker_build import close_logger  # type: ignore
-from swebench.harness.docker_build import (
-    BuildImageError,
-    build_container,
-    setup_logger,
-)
-from swebench.harness.docker_utils import cleanup_container  # type: ignore
-from swebench.harness.docker_utils import exec_run_with_timeout  # type: ignore
-from swebench.harness.docker_utils import remove_image  # type: ignore
-from swebench.harness.docker_utils import should_remove  # type: ignore
-from swebench.harness.docker_utils import (
-    copy_to_container,
-)
-from swebench.harness.grading import get_eval_report
-from swebench.harness.test_spec.test_spec import TestSpec, make_test_spec
-from swebench.harness.utils import EvaluationError
-
-import docker
-
-GIT_APPLY_CMDS = [
-    "git apply --verbose",
-    "git apply --verbose --reject",
-    "patch --batch --fuzz=5 -p1 -i",
-]
-
-
-def run_instance(
-    test_spec: TestSpec,
-    pred: Dict[str, Any],
-    rm_image: bool,
-    force_rebuild: bool,
-    client: docker.DockerClient,
-    run_id: str,
-    timeout: int | None = None,
-    rewrite_reports: bool = False,
-):
-    """
-    Run a single instance with the given prediction.
-
-    Args:
-        test_spec (TestSpec): TestSpec instance
-        pred (dict): Prediction w/ model_name_or_path, model_patch, instance_id
-        rm_image (bool): Whether to remove the image after running
-        force_rebuild (bool): Whether to force rebuild the image
-        client (docker.DockerClient): Docker client
-        run_id (str): Run ID
-        timeout (int): Timeout for running tests
-        rewrite_reports (bool): True if eval run is just to reformat existing report
-    """
-    # Set up logging directory
-    instance_id = test_spec.instance_id
-    model_name_or_path = pred.get(KEY_MODEL, "None").replace("/", "__")
-    log_dir = RUN_EVALUATION_LOG_DIR / run_id / model_name_or_path / instance_id
-
-    # Set up report file
-    report_path = log_dir / LOG_REPORT
-    if rewrite_reports:
-        test_output_path = log_dir / LOG_TEST_OUTPUT
-        if not test_output_path.exists():
-            raise ValueError(f"Test output file {test_output_path} does not exist")
-        report = get_eval_report(
-            test_spec=test_spec,
-            prediction=pred,
-            test_log_path=test_output_path,
-            include_tests_status=True,
-        )
-        # Write report to report.json
-        with open(report_path, "w") as f:
-            f.write(json.dumps(report, indent=4))
-        return instance_id, report
-    if report_path.exists():
-        return instance_id, json.loads(report_path.read_text())
-
-    if not test_spec.is_remote_image:
-        # Link the image build dir in the log dir
-        build_dir = INSTANCE_IMAGE_BUILD_DIR / test_spec.instance_image_key.replace(":", "__")
-        image_build_link = log_dir / "image_build_dir"
-        if not image_build_link.exists():
-            try:
-                # link the image build dir in the log dir
-                image_build_link.symlink_to(build_dir.absolute(), target_is_directory=True)
-            except:
-                # some error, idk why
-                pass
-
-    # Set up logger
-    log_dir.mkdir(parents=True, exist_ok=True)
-    log_file = log_dir / LOG_INSTANCE
-    logger = setup_logger(instance_id, log_file)
-
-    # Run the instance
-    container = None
-    try:
-        # Build + start instance container (instance image should already be built)
-        container = build_container(test_spec, client, run_id, logger, rm_image, force_rebuild)
-        container.start()
-        logger.info(f"Container for {instance_id} started: {container.id}")
-
-        # Copy model prediction as patch file to container
-        patch_file = Path(log_dir / "patch.diff")
-        patch_file.write_text(pred[KEY_PREDICTION] or "")
-        logger.info(f"Intermediate patch for {instance_id} written to {patch_file}, now applying to container...")
-        copy_to_container(container, patch_file, PurePosixPath(DOCKER_PATCH))  # type: ignore
-
-        # Attempt to apply patch to container (TODO: FIX THIS)
-        val: Optional[ExecResult] = None
-        for git_apply_cmd in GIT_APPLY_CMDS:
-            val = container.exec_run(  # type: ignore
-                f"{git_apply_cmd} {DOCKER_PATCH}",
-                workdir=DOCKER_WORKDIR,
-                user=DOCKER_USER,
-            )
-            if val.exit_code == 0:
-                logger.info(f"{APPLY_PATCH_PASS}:\n{val.output.decode(UTF8)}")
-                break
-            else:
-                logger.info(f"Failed to apply patch to container: {git_apply_cmd}")
-        if val is not None:
-            logger.info(f"{APPLY_PATCH_FAIL}:\n{val.output.decode(UTF8)}")
-            raise EvaluationError(
-                instance_id,
-                f"{APPLY_PATCH_FAIL}:\n{val.output.decode(UTF8)}",
-                logger,
-            )
-
-        # Get git diff before running eval script
-        git_diff_output_before = (
-            container.exec_run("git -c core.fileMode=false diff", workdir=DOCKER_WORKDIR).output.decode(UTF8).strip()  # type: ignore
-        )
-        logger.info(f"Git diff before:\n{git_diff_output_before}")
-
-        eval_file = Path(log_dir / "eval.sh")
-        eval_file.write_text(test_spec.eval_script)
-        logger.info(f"Eval script for {instance_id} written to {eval_file}; copying to container...")
-        copy_to_container(container, eval_file, PurePosixPath("/eval.sh"))  # type: ignore
-
-        # Run eval script, write output to logs
-        test_output, timed_out, total_runtime = exec_run_with_timeout(container, "/bin/bash /eval.sh", timeout)
-        test_output_path = log_dir / LOG_TEST_OUTPUT
-        logger.info(f"Test runtime: {total_runtime:_.2f} seconds")
-        with open(test_output_path, "w") as f:
-            f.write(test_output)
-            logger.info(f"Test output for {instance_id} written to {test_output_path}")
-            if timed_out:
-                f.write(f"\n\nTimeout error: {timeout} seconds exceeded.")
-                raise EvaluationError(
-                    instance_id,
-                    f"Test timed out after {timeout} seconds.",
-                    logger,
-                )
-
-        # Get git diff after running eval script (ignore permission changes)
-        git_diff_output_after = (
-            container.exec_run("git -c core.fileMode=false diff", workdir=str(DOCKER_WORKDIR)).output.decode(UTF8).strip()  # type: ignore
-        )
-
-        # Check if git diff changed after running eval script
-        logger.info(f"Git diff after:\n{git_diff_output_after}")
-        if git_diff_output_after != git_diff_output_before:
-            logger.info("Git diff changed after running eval script")
-
-        # Get report from test output
-        logger.info(f"Grading answer for {instance_id}...")
-        report = get_eval_report(
-            test_spec=test_spec,
-            prediction=pred,
-            test_log_path=test_output_path,
-            include_tests_status=True,
-        )
-        logger.info(f"report: {report}\n" f"Result for {instance_id}: resolved: {report[instance_id]['resolved']}")
-
-        # Write report to report.json
-        with open(report_path, "w") as f:
-            f.write(json.dumps(report, indent=4))
-        return instance_id, report
-    except EvaluationError as e:
-        error_msg = traceback.format_exc()
-        logger.info(error_msg)
-        print(e)
-    except BuildImageError as e:
-        error_msg = traceback.format_exc()
-        logger.info(error_msg)
-        print(e)
-    except Exception as e:
-        error_msg = f"Error in evaluating model for {instance_id}: {e}\n" f"{traceback.format_exc()}"
-        logger.error(error_msg)
-    finally:
-        # Remove instance container + image, close logger
-        cleanup_container(client, container, logger)
-        if rm_image:
-            remove_image(client, test_spec.instance_image_key, logger)
-        close_logger(logger)
-    return
-
-
-def evaluate(
-    prediction: Dict[str, Any],
-    instance: SWEbenchInstance,
-    cache_level: str,
-    clean: bool,
-    force_rebuild: bool,
-    run_id: str,
-    timeout: Optional[int],
-    namespace: Optional[str],
-    instance_image_tag: str,
-    rewrite_reports: bool,
-):
-    client = docker.from_env()
-    test_spec = make_test_spec(instance, namespace=namespace, instance_image_tag=instance_image_tag)
-
-    instance_image_ids = {
-        test_spec.instance_image_key,
-    }
-    existing_images = {tag for i in client.images.list(all=True) for tag in i.tags if tag in instance_image_ids}
-
-    return run_instance(
-        test_spec,
-        prediction,
-        should_remove(test_spec.instance_image_key, cache_level, clean, existing_images),
-        force_rebuild,
-        client,
-        run_id,
-        timeout,
-        rewrite_reports,
-    )

examples/claude_code/swebench_utils/logging.py

@@ -1,26 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Logging utility module for SWE-bench evaluation runs.
-
-This module provides a simple logging utility function that writes evaluation
-results and logs to timestamped files organized by run ID and instance ID.
-"""
-
-import datetime
-import os
-
-
-def log_for_evaluation(run_id: str, instance_id: str, text: str) -> None:
-    """Log a message for evaluation purposes of SWE-Bench.
-
-    The format follows the SWE-Bench evaluation framework.
-
-    Args:
-        run_id: The run ID of the evaluation.
-        instance_id: The instance ID of the evaluation.
-        text: The text to log.
-    """
-    os.makedirs(f"./logs/{run_id}", exist_ok=True)
-    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    with open(f"./logs/{run_id}/{instance_id}", mode="a") as f:
-        print(f"\n\n{current_time}\n{text}\n", file=f)

examples/claude_code/templates/handle_hook.template.sh

@@ -1,13 +0,0 @@
-#!/bin/bash
-
-# Read from stdin
-input=$(cat)
-
-# Define output file
-output_file="/tmp/hook.out"
-
-# Append input followed by two newlines
-echo -e "${input}\n\n" >> "$output_file"
-
-# Exit with status 0
-exit 0

examples/claude_code/templates/settings.template.json

@@ -1,94 +0,0 @@
-{
-  "hooks": {
-    "PreToolUse": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "PostToolUse": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "Notification": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "UserPromptSubmit": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "Stop": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "SubagentStop": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "PreCompact": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "SessionStart": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ],
-    "SessionEnd": [
-      {
-        "hooks": [
-          {
-            "type": "command",
-            "command": "/tmp/handle_hook.sh"
-          }
-        ]
-      }
-    ]
-  }
-}

examples/minimal/README.md

@@ -1,7 +1,5 @@
 # Minimal Component Showcase
 
-[![minimal CI 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)
-
 `examples/minimal` provides bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation.
 
 Each module have been documented with its own CLI usage in the module-level docstring. Use this directory as a reference when wiring the same pieces into a larger system.

examples/minimal/write_traces.py

@@ -22,7 +22,6 @@ from rich.console import Console
 
 from agentlightning import AgentOpsTracer, LightningStoreClient, OtelTracer, Span, emit_reward, setup_logging
 from agentlightning.store import InMemoryLightningStore
-from agentlightning.utils.otel import get_tracer_provider
 
 console = Console()
 
@@ -62,13 +61,13 @@ async def send_traces_via_otel(use_client: bool = False):
     assert "grpc-span-1" in span_names
     assert "grpc-span-2" in span_names
     assert "grpc-span-3" in span_names
-    assert "agentlightning.annotation" in span_names
+    assert "agentlightning.reward" in span_names
 
     last_span = traces[-1]
-    assert last_span.name == "agentlightning.annotation"
-    # NOTE: Try not to rely on this attribute like this example do. It may change in the future.
+    assert last_span.name == "agentlightning.reward"
+    # NOTE: Try not to rely on this attribute. It may change in the future.
     # Use utils from agentlightning.emitter to get the reward value.
-    assert last_span.attributes["agentlightning.reward.0.value"] == 1.0
+    assert last_span.attributes["reward"] == 1.0
 
     if use_client:
         # When using client, the resource should have rollout_id and attempt_id set
@@ -91,9 +90,6 @@ async def send_traces_via_agentops(use_client: bool = False):
     # Initialize the tracer lifespan
     # One lifespan can contain multiple traces
     with tracer.lifespan(store):
-        # Inspect current tracer provider
-        get_tracer_provider(inspect=True)
-
         # Initialize the capture of one single trace for one single rollout
         async with tracer.trace_context(
             "trace-1", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id

examples/tinker/README.md

@@ -1,9 +1,9 @@
 # Tinker + Agent-lightning Integration
 
-[![tinker CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml)
-
 This example shows how to use [Tinker's reinforcement-learning infrastructure](https://tinker-docs.thinkingmachines.ai/) as a fine-tuning backend for agents written against Agent-lightning. You author the agent exactly the way you would for deployment, while the bridge code reconstructs Tinker-compatible trajectories from Agent-lightning traces.
 
+**NOTE: The example is tested and compatible with Agent-lightning v0.2.x, but it's not yet maintained on CI due to the cost of running the Tinker training service.**
+
 ## How this differs from the original Tinker Cookbook RL recipe
 
 Real-world agent apps orchestrate logic in familiar frameworks (CrewAI, LangChain, AutoGen, OpenAI Agents, etc.) or by calling OpenAI-compatible REST APIs. A simple number-guessing agent might look like this:

pyproject.toml

@@ -98,8 +98,7 @@ torch-stable = [
   # This can work for both CPU and GPU.
   "torch>=2.8.0",
   "torchvision>=0.23.0",
-  # https://github.com/huggingface/transformers/issues/42369
-  "transformers>=4.55.0,!=4.57.2",
+  "transformers>=4.55.0",
   # vLLM 0.11.1 requires PyTorch 2.9.0, which is incompatible with flash-attn
   # https://github.com/Dao-AILab/flash-attention/issues/1967
   # Similar issues with vLLM 0.11.2
@@ -193,10 +192,7 @@ sql = [
 ]
 crewai = [
   # https://github.com/crewAIInc/crewAI/issues/3959
-  "crewai[tools]>=1.2.0,!=1.2.1,!=1.3.0,!=1.4.0,!=1.4.1,!=1.5.0",
-]
-swebench = [
-  "swebench",
+  "crewai[tools]==1.2.0",
 ]
 
 # Summarize into large installable groups.
@@ -207,7 +203,6 @@ agents = [
   {include-group = "sql"},
   {include-group = "anthropic"},
   {include-group = "crewai"},
-  {include-group = "swebench"},
 ]
 
 [tool.uv]

tests/algorithm/test_apo.py

@@ -12,7 +12,6 @@ import agentlightning.algorithm.apo.apo as apo_module
 from agentlightning.adapter import TraceAdapter
 from agentlightning.adapter.messages import TraceToMessages
 from agentlightning.algorithm.apo.apo import APO, RolloutResultForAPO, VersionedPromptTemplate, batch_iter_over_dataset
-from agentlightning.semconv import AGL_ANNOTATION
 from agentlightning.types import (
     Dataset,
     NamedResources,
@@ -23,6 +22,7 @@ from agentlightning.types import (
     Rollout,
     Span,
     SpanContext,
+    SpanNames,
     TraceStatus,
 )
 
@@ -120,7 +120,7 @@ def make_reward_span(rollout_id: str, attempt_id: str, reward: float, sequence_i
         trace_id=hex_id,
         span_id=span_hex,
         parent_id=None,
-        name=AGL_ANNOTATION,
+        name=SpanNames.REWARD.value,
         status=TraceStatus(status_code="OK"),
         attributes={"reward": reward},
         events=[],
@@ -424,7 +424,7 @@ async def test_get_rollout_results_adapts_spans() -> None:
     # Verify spans were serialized
     assert len(results[0]["spans"]) == 2
     assert results[0]["spans"][0]["rollout_id"] == "r-1"
-    assert results[0]["spans"][0]["name"] == AGL_ANNOTATION
+    assert results[0]["spans"][0]["name"] == SpanNames.REWARD.value
     assert results[0]["spans"][0]["attributes"]["reward"] == 1.0
     assert results[0]["spans"][1]["attributes"]["reward"] == 2.0
 

tests/benchmark/analysis.py

@@ -1,922 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Lightweight benchmark report for the Prometheus + Grafana stack shipped with Agent Lightning."""
-
-from __future__ import annotations
-
-import argparse
-import datetime as dt
-import json
-import math
-from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeGuard, cast
-from urllib import error, parse, request
-
-
-class PrometheusQueryError(RuntimeError):
-    """Raised when Prometheus returns an error payload."""
-
-
-class PrometheusClient:
-    """Tiny helper around the Prometheus HTTP API."""
-
-    def __init__(
-        self,
-        base_url: str,
-        timeout: float = 10.0,
-        default_time: Optional[dt.datetime] = None,
-    ):
-        self.base_url = base_url.rstrip("/")
-        self.timeout = timeout
-        self.default_time = default_time
-
-    def query_vector(self, expr: str, eval_time: Optional[dt.datetime] = None) -> List[Mapping[str, object]]:
-        params: Dict[str, str] = {"query": expr}
-        query_time = eval_time or self.default_time
-        if query_time is not None:
-            params["time"] = query_time.isoformat()
-        payload = self._get("/api/v1/query", params)
-        status = payload.get("status")
-        if not isinstance(status, str) or status != "success":
-            error_msg = payload.get("error", "unknown error")
-            raise PrometheusQueryError(str(error_msg))
-        data_obj = payload.get("data", {})
-        if isinstance(data_obj, dict):
-            data = cast(Dict[str, Any], data_obj)
-        else:
-            data = {}
-        result_type_obj = data.get("resultType")
-        result_type = result_type_obj if isinstance(result_type_obj, str) else None
-        raw_result_obj = data.get("result", [])
-        raw_result: List[object]
-        if isinstance(raw_result_obj, list):
-            raw_result = cast(List[object], raw_result_obj)
-        else:
-            raw_result = []
-        if result_type == "scalar":
-            if len(raw_result) >= 2:
-                ts = raw_result[0]
-                value = raw_result[1]
-                return [{"metric": {}, "value": [ts, value]}]
-            return []
-        vector_result: List[Mapping[str, object]] = [
-            cast(Mapping[str, object], item) for item in raw_result if isinstance(item, Mapping)
-        ]
-        if result_type == "matrix":
-            collapsed: List[Dict[str, object]] = []
-            for series in vector_result:
-                values_obj = series.get("values")
-                if isinstance(values_obj, list) and values_obj and isinstance(values_obj[-1], Sequence):
-                    last = cast(Sequence[object], values_obj[-1])
-                else:
-                    continue
-                metric_obj = series.get("metric")
-                if isinstance(metric_obj, Mapping):
-                    metric: Dict[str, object] = dict(cast(Mapping[str, object], metric_obj))
-                else:
-                    metric = {}
-                collapsed.append({"metric": metric, "value": list(last)})
-            return cast(List[Mapping[str, object]], collapsed)
-        if result_type == "vector":
-            return vector_result
-        return []
-
-    def query_scalar(self, expr: str, eval_time: Optional[dt.datetime] = None) -> Optional[float]:
-        samples = self.query_vector(expr, eval_time=eval_time)
-        if not samples:
-            return None
-        return _sample_value(samples[0])
-
-    def _get(self, path: str, data: Optional[Mapping[str, str]] = None) -> Dict[str, Any]:
-        encoded: Optional[bytes] = None
-        if data is not None:
-            encoded = parse.urlencode(data).encode()
-        req = request.Request(f"{self.base_url}{path}", data=encoded)
-        try:
-            with request.urlopen(req, timeout=self.timeout) as resp:
-                loaded = json.loads(resp.read().decode())
-                if isinstance(loaded, dict):
-                    return cast(Dict[str, Any], loaded)
-                return {}
-        except error.URLError as exc:  # pragma: no cover - network/infra issues
-            raise PrometheusQueryError(str(exc)) from exc
-
-
-def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Summarize benchmark metrics from Prometheus.")
-    parser.add_argument("--prom-url", default="http://localhost:9090", help="Base URL for the Prometheus API.")
-    parser.add_argument(
-        "--store-url",
-        default="http://localhost:4747/v1/agl",
-        help="Base URL for the Lightning Store API (without the /statistics suffix).",
-    )
-    parser.add_argument("--timeout", type=float, default=10.0, help="HTTP timeout in seconds.")
-    parser.add_argument("--start", type=str, help="ISO timestamp (e.g. 2024-05-01T12:00:00Z).")
-    parser.add_argument("--end", type=str, help="ISO timestamp (default: now).")
-    parser.add_argument(
-        "--duration",
-        type=str,
-        default="5m",
-        help="Fallback duration (e.g. 5m, 1h) used when --start is omitted.",
-    )
-    parser.add_argument("--top", type=int, default=8, help="Number of rows to show per table.")
-    return parser.parse_args(argv)
-
-
-def parse_timestamp(value: Optional[str], default: Optional[dt.datetime] = None) -> Optional[dt.datetime]:
-    if value is None:
-        return default
-    try:
-        if value.endswith("Z"):
-            value = value[:-1] + "+00:00"
-        return dt.datetime.fromisoformat(value).astimezone(dt.timezone.utc)
-    except ValueError as exc:  # pragma: no cover - invalid CLI input
-        raise SystemExit(f"Invalid timestamp '{value}': {exc}") from exc
-
-
-def parse_duration(text: str) -> dt.timedelta:
-    units = {"s": 1, "m": 60, "h": 3600}
-    if text.isdigit():
-        return dt.timedelta(seconds=int(text))
-    suffix = text[-1]
-    if suffix not in units:
-        raise SystemExit(f"Unsupported duration '{text}'. Use Ns/Nm/Nh.")
-    try:
-        value = int(text[:-1])
-    except ValueError as exc:  # pragma: no cover - invalid CLI input
-        raise SystemExit(f"Invalid duration '{text}': {exc}") from exc
-    return dt.timedelta(seconds=value * units[suffix])
-
-
-def format_window(seconds: float) -> str:
-    seconds = max(int(seconds), 1)
-    return f"{seconds}s"
-
-
-def compute_rate_window(duration_seconds: float) -> str:
-    return format_window(min(duration_seconds, 60.0))
-
-
-def compute_subquery_step(duration_seconds: float) -> str:
-    step_seconds = max(int(duration_seconds / 60), 1)
-    step_seconds = min(step_seconds, 15)
-    return f"{step_seconds}s"
-
-
-def _is_http_pair(value: Any) -> TypeGuard[Tuple[Any, Any]]:
-    if not isinstance(value, tuple):
-        return False
-    try:
-        value[0]
-        value[1]
-    except IndexError:
-        return False
-    return True
-
-
-def _sample_value(sample: Mapping[str, object]) -> Optional[float]:
-    value_obj = sample.get("value")
-    if not isinstance(value_obj, Sequence):
-        return None
-    value_seq = cast(Sequence[object], value_obj)
-    if len(value_seq) < 2:
-        return None
-    candidate = value_seq[1]
-    if isinstance(candidate, (int, float)):
-        return float(candidate)
-    if isinstance(candidate, str):
-        try:
-            return float(candidate)
-        except ValueError:
-            return None
-    return None
-
-
-def vector_to_map(
-    samples: Optional[Sequence[Mapping[str, object]]],
-    labels: Sequence[str],
-) -> Dict[Any, float]:
-    mapping: Dict[Any, float] = {}
-    if not samples:
-        return mapping
-    for sample in samples:
-        metric_obj = sample.get("metric", {})
-        if isinstance(metric_obj, Mapping):
-            metric: Dict[str, object] = dict(cast(Mapping[str, object], metric_obj))
-        else:
-            metric = {}
-        if len(labels) == 1:
-            key: Any = str(metric.get(labels[0], ""))
-        else:
-            key = tuple(str(metric.get(label, "")) for label in labels)
-        value = _sample_value(sample)
-        if value is not None:
-            mapping[key] = value
-    return mapping
-
-
-def safe_vector(client: PrometheusClient, expr: str) -> Optional[List[Mapping[str, object]]]:
-    try:
-        return client.query_vector(expr)
-    except PrometheusQueryError as exc:
-        print(f"[warn] Prometheus query failed: {exc} (expr={expr})")
-        return None
-
-
-def safe_scalar(client: PrometheusClient, expr: str) -> Optional[float]:
-    try:
-        return client.query_scalar(expr)
-    except PrometheusQueryError as exc:
-        print(f"[warn] Prometheus query failed: {exc} (expr={expr})")
-        return None
-
-
-def fetch_store_statistics(store_url: str, timeout: float) -> Optional[Dict[str, Any]]:
-    store_url = store_url.rstrip("/")
-    stats_url = f"{store_url}/statistics"
-    req = request.Request(stats_url)
-    try:
-        with request.urlopen(req, timeout=timeout) as resp:
-            loaded = json.loads(resp.read().decode())
-            if isinstance(loaded, Mapping):
-                return dict(cast(Mapping[str, Any], loaded))
-            return None
-    except error.URLError as exc:
-        print(f"[warn] Failed to fetch store statistics: {exc} (url={stats_url})")
-        return None
-    except json.JSONDecodeError as exc:
-        print(f"[warn] Failed to decode store statistics: {exc} (url={stats_url})")
-        return None
-
-
-# ---------------------------------------------------------------------------
-# Part 1 – high level throughput
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class CollectionThroughput:
-    name: str
-    count: Optional[float]
-    per_sec: Optional[float]
-
-
-STORE_TOTAL_FIELDS = {
-    "rollouts": "total_rollouts",
-    "spans": "total_spans",
-    "attempts": "total_attempts",
-    "resources": "total_resources",
-    "workers": "total_workers",
-}
-STORE_TOTAL_COLLECTIONS = tuple(STORE_TOTAL_FIELDS.keys())
-
-
-def _coerce_int(value: Any) -> Optional[int]:
-    if isinstance(value, bool):
-        return int(value)
-    if isinstance(value, int):
-        return value
-    if isinstance(value, float):
-        if math.isnan(value):
-            return None
-        return int(value)
-    if isinstance(value, str):
-        try:
-            return int(value)
-        except ValueError:
-            try:
-                return int(float(value))
-            except ValueError:
-                return None
-    return None
-
-
-def extract_store_totals(stats: Optional[Mapping[str, Any]]) -> Dict[str, Optional[int]]:
-    totals: Dict[str, Optional[int]] = {}
-    if not stats:
-        return totals
-    for display_name, field_name in STORE_TOTAL_FIELDS.items():
-        if field_name in stats:
-            totals[display_name] = _coerce_int(stats.get(field_name))
-        else:
-            totals[display_name] = None
-    return totals
-
-
-def gather_collection_throughput(
-    client: PrometheusClient, collections: Sequence[str], duration_seconds: float
-) -> List[CollectionThroughput]:
-    rows: List[CollectionThroughput] = []
-    window = format_window(duration_seconds)
-    for collection in collections:
-        # Successful insert operations reflect the number of new records.
-        expr = (
-            "sum("
-            f'increase(mongo_operation_total{{collection="{collection}", operation="insert", status="ok"}}[{window}])'
-            ")"
-        )
-        count = safe_scalar(client, expr)
-        if count is not None and count < 0:
-            count = 0.0
-        per_sec = (count / duration_seconds) if (count is not None and duration_seconds > 0) else None
-        rows.append(CollectionThroughput(collection, count, per_sec))
-    return rows
-
-
-# ---------------------------------------------------------------------------
-# Part 2 – CollectionBasedLightningStore method stats
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class StoreMethodStats:
-    method: str
-    ops_mean: float
-    ops_max: Optional[float]
-    ops_min: Optional[float]
-    p50: Optional[float]
-    p95: Optional[float]
-    p99: Optional[float]
-
-
-StatsSummary = Dict[str, Optional[float]]
-
-
-@dataclass
-class RolloutOutcomeStats:
-    status: str
-    rate: Optional[float]
-    p25: Optional[float]
-    p50: Optional[float]
-    p75: Optional[float]
-    max_latency: Optional[float]
-
-
-def gather_store_methods(
-    client: PrometheusClient,
-    window: str,
-    rate_window: str,
-    subquery_step: str,
-) -> Tuple[List[StoreMethodStats], StatsSummary]:
-    ops_expr = f"sum by (method)(rate(collection_store_total[{rate_window}]))"
-    ops_mean = vector_to_map(
-        safe_vector(client, f"avg_over_time(({ops_expr})[{window}:{subquery_step}])"),
-        ("method",),
-    )
-    ops_max = vector_to_map(
-        safe_vector(client, f"max_over_time(({ops_expr})[{window}:{subquery_step}])"),
-        ("method",),
-    )
-    ops_min = vector_to_map(
-        safe_vector(client, f"min_over_time(({ops_expr})[{window}:{subquery_step}])"),
-        ("method",),
-    )
-    p50 = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.50, sum by (le, method)(rate(collection_store_latency_seconds_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}])))",
-        ),
-        ("method",),
-    )
-    p99 = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.99, sum by (le, method)(rate(collection_store_latency_seconds_bucket[{window}])))",
-        ),
-        ("method",),
-    )
-    rows: List[StoreMethodStats] = []
-    for method, rate in sorted(ops_mean.items(), key=lambda item: str(item[0])):
-        p50_val = p50.get(method)
-        p95_val = p95.get(method)
-        p99_val = p99.get(method)
-        rows.append(
-            StoreMethodStats(
-                str(method or "-"),
-                rate,
-                ops_max.get(method),
-                ops_min.get(method),
-                p50_val,
-                p95_val,
-                p99_val,
-            )
-        )
-
-    overall: StatsSummary = {
-        "ops_mean": safe_scalar(
-            client,
-            f"avg_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
-        )
-        or 0.0,
-        "ops_max": safe_scalar(
-            client,
-            f"max_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
-        ),
-        "ops_min": safe_scalar(
-            client,
-            f"min_over_time((sum(rate(collection_store_total[{rate_window}])))[{window}:{subquery_step}])",
-        ),
-        "p50": safe_scalar(
-            client,
-            f"histogram_quantile(0.50, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
-        )
-        or 0.0,
-        "p95": safe_scalar(
-            client,
-            f"histogram_quantile(0.95, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
-        )
-        or 0.0,
-        "p99": safe_scalar(
-            client,
-            f"histogram_quantile(0.99, sum by (le)(rate(collection_store_latency_seconds_bucket[{window}])))",
-        )
-        or 0.0,
-    }
-    return rows, overall
-
-
-def gather_rollout_outcomes(
-    client: PrometheusClient,
-    window: str,
-    rate_window: str,
-) -> List[RolloutOutcomeStats]:
-    rate_map = vector_to_map(
-        safe_vector(client, f"sum by (status)(rate(collection_store_rollout_total[{rate_window}]))"),
-        ("status",),
-    )
-    p25_map = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.25, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
-        ),
-        ("status",),
-    )
-    p50_map = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.50, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
-        ),
-        ("status",),
-    )
-    p75_map = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.75, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
-        ),
-        ("status",),
-    )
-    max_map = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(1.00, "
-            f"sum by (le, status)(rate(collection_store_rollout_duration_seconds_bucket[{window}])))",
-        ),
-        ("status",),
-    )
-    statuses = sorted({*rate_map.keys(), *p25_map.keys(), *p50_map.keys(), *p75_map.keys(), *max_map.keys()}, key=str)
-    stats: List[RolloutOutcomeStats] = []
-    for status in statuses:
-        stats.append(
-            RolloutOutcomeStats(
-                status=str(status or "-"),
-                rate=rate_map.get(status),
-                p25=p25_map.get(status),
-                p50=p50_map.get(status),
-                p75=p75_map.get(status),
-                max_latency=max_map.get(status),
-            )
-        )
-    return stats
-
-
-# ---------------------------------------------------------------------------
-# Part 3 – HTTP traffic
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class HttpPathStats:
-    method: str
-    path: str
-    qps_mean: float
-    qps_max: Optional[float]
-    qps_min: Optional[float]
-    p50: Optional[float]
-    p95: Optional[float]
-    p99: Optional[float]
-
-
-def gather_http_paths(
-    client: PrometheusClient,
-    window: str,
-    rate_window: str,
-    subquery_step: str,
-) -> Tuple[List[HttpPathStats], StatsSummary]:
-    qps_expr = f"sum by (method, path)(rate(http_requests_total[{rate_window}]))"
-    qps_mean = vector_to_map(
-        safe_vector(client, f"avg_over_time(({qps_expr})[{window}:{subquery_step}])"),
-        ("method", "path"),
-    )
-    qps_max = vector_to_map(
-        safe_vector(client, f"max_over_time(({qps_expr})[{window}:{subquery_step}])"),
-        ("method", "path"),
-    )
-    qps_min = vector_to_map(
-        safe_vector(client, f"min_over_time(({qps_expr})[{window}:{subquery_step}])"),
-        ("method", "path"),
-    )
-    p50 = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.50, sum by (le, method, path)(rate(http_request_duration_seconds_bucket[{window}])))",
-        ),
-        ("method", "path"),
-    )
-    p95 = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.95, sum by (le, method, path)(rate(http_request_duration_seconds_bucket[{window}])))",
-        ),
-        ("method", "path"),
-    )
-    p99 = vector_to_map(
-        safe_vector(
-            client,
-            f"histogram_quantile(0.99, sum by (le, method, path)(rate(http_request_duration_seconds_bucket[{window}])))",
-        ),
-        ("method", "path"),
-    )
-
-    def normalize_http_key(raw_key: Any) -> Tuple[str, str]:
-        if _is_http_pair(raw_key):
-            method_raw, path_raw = raw_key
-            return (str(method_raw or "-"), str(path_raw or "-"))
-        return ("-", str(raw_key))
-
-    def normalize_dict(source: Mapping[Any, Optional[float]]) -> Dict[Tuple[str, str], Optional[float]]:
-        normalized: Dict[Tuple[str, str], Optional[float]] = {}
-        for raw_key, value in source.items():
-            normalized[normalize_http_key(raw_key)] = value
-        return normalized
-
-    qps_mean_norm = normalize_dict(cast(Mapping[Any, Optional[float]], qps_mean))
-    qps_max_norm = normalize_dict(cast(Mapping[Any, Optional[float]], qps_max))
-    qps_min_norm = normalize_dict(cast(Mapping[Any, Optional[float]], qps_min))
-    p50_norm = normalize_dict(cast(Mapping[Any, Optional[float]], p50))
-    p95_norm = normalize_dict(cast(Mapping[Any, Optional[float]], p95))
-    p99_norm = normalize_dict(cast(Mapping[Any, Optional[float]], p99))
-
-    path_stats: List[HttpPathStats] = []
-    for method_path in sorted(qps_mean_norm.keys()):
-        method_label, path_label = method_path
-        path_stats.append(
-            HttpPathStats(
-                method_label,
-                path_label,
-                qps_mean_norm.get(method_path, 0.0) or 0.0,
-                qps_max_norm.get(method_path),
-                qps_min_norm.get(method_path),
-                p50_norm.get(method_path),
-                p95_norm.get(method_path),
-                p99_norm.get(method_path),
-            )
-        )
-    overall: StatsSummary = {
-        "qps_mean": safe_scalar(
-            client,
-            f"avg_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
-        )
-        or 0.0,
-        "qps_max": safe_scalar(
-            client,
-            f"max_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
-        ),
-        "qps_min": safe_scalar(
-            client,
-            f"min_over_time((sum(rate(http_requests_total[{rate_window}])))" f"[{window}:{subquery_step}])",
-        ),
-        "p50": safe_scalar(
-            client, f"histogram_quantile(0.50, sum by (le)(rate(http_request_duration_seconds_bucket[{window}])))"
-        )
-        or 0.0,
-        "p95": safe_scalar(
-            client, f"histogram_quantile(0.95, sum by (le)(rate(http_request_duration_seconds_bucket[{window}])))"
-        )
-        or 0.0,
-        "p99": safe_scalar(
-            client, f"histogram_quantile(0.99, sum by (le)(rate(http_request_duration_seconds_bucket[{window}])))"
-        )
-        or 0.0,
-    }
-    return path_stats, overall
-
-
-# ---------------------------------------------------------------------------
-# Part 4 – diagnostics
-# ---------------------------------------------------------------------------
-
-
-def gather_diagnostics(client: PrometheusClient, window: str) -> Dict[str, Any]:
-    diagnostics: Dict[str, Any] = {}
-    diagnostics["mongo_ops"] = vector_to_map(
-        safe_vector(
-            client,
-            f"sum by (operation)(rate(mongo_operation_total{{operation!='ensure_collection'}}[{window}]))",
-        ),
-        ("operation",),
-    )
-    opcounters_samples = safe_vector(client, f"sum by (legacy_op_type)(rate(mongodb_ss_opcounters[{window}]))")
-    mongo_opcounters: Dict[str, float] = {}
-    if opcounters_samples:
-        for sample in opcounters_samples:
-            metric_obj = sample.get("metric", {})
-            if isinstance(metric_obj, Mapping):
-                metric: Dict[str, object] = dict(cast(Mapping[str, object], metric_obj))
-            else:
-                metric = {}
-            label_value = metric.get("legacy_op_type") or metric.get("type")
-            label = str(label_value) if label_value is not None else ""
-            value = _sample_value(sample)
-            if value is not None:
-                mongo_opcounters[str(label or "-")] = value
-    diagnostics["mongo_opcounters"] = mongo_opcounters
-    diagnostics["mongo_connections"] = safe_scalar(client, "avg(mongodb_ss_connections{conn_type='current'})")
-    diagnostics["cpu_usage"] = safe_scalar(client, f"1 - avg(rate(node_cpu_seconds_total{{mode='idle'}}[{window}]))")
-    diagnostics["memory_total"] = safe_scalar(client, "avg(node_memory_MemTotal_bytes)")
-    diagnostics["memory_available"] = safe_scalar(client, "avg(node_memory_MemAvailable_bytes)")
-    diagnostics["network_rx"] = safe_scalar(
-        client,
-        f"sum(rate(node_network_receive_bytes_total{{device!~'lo|docker.*'}}[{window}]))",
-    )
-    diagnostics["network_tx"] = safe_scalar(
-        client,
-        f"sum(rate(node_network_transmit_bytes_total{{device!~'lo|docker.*'}}[{window}]))",
-    )
-    diagnostics["disk_read_ops"] = safe_scalar(client, f"sum(rate(node_disk_reads_completed_total[{window}]))")
-    diagnostics["disk_write_ops"] = safe_scalar(client, f"sum(rate(node_disk_writes_completed_total[{window}]))")
-    diagnostics["disk_read_bytes"] = safe_scalar(client, f"sum(rate(node_disk_read_bytes_total[{window}]))")
-    diagnostics["disk_write_bytes"] = safe_scalar(client, f"sum(rate(node_disk_written_bytes_total[{window}]))")
-    return diagnostics
-
-
-# ---------------------------------------------------------------------------
-# Rendering helpers
-# ---------------------------------------------------------------------------
-
-
-def render_table(headers: Sequence[str], rows: Sequence[Sequence[str]]) -> List[str]:
-    if not rows:
-        return [f"(no data for {headers})"]
-    widths = [len(h) for h in headers]
-    rendered: List[List[str]] = []
-    for row in rows:
-        rendered_row = [str(cell) for cell in row]
-        for idx, cell in enumerate(rendered_row):
-            widths[idx] = max(widths[idx], len(cell))
-        rendered.append(rendered_row)
-
-    lines = [
-        " | ".join(headers[idx].ljust(widths[idx]) for idx in range(len(headers))),
-        "-+-".join("-" * widths[idx] for idx in range(len(headers))),
-    ]
-    for row in rendered:
-        lines.append(" | ".join(row[idx].ljust(widths[idx]) for idx in range(len(headers))))
-    return lines
-
-
-def fmt_rate(value: Optional[float]) -> str:
-    if value is None or math.isnan(value):
-        return "-"
-    return f"{value:.2f}/s"
-
-
-def fmt_latency(value: Optional[float]) -> str:
-    if value is None or math.isnan(value):
-        return "-"
-    if value < 0.5:
-        return f"{value * 1e3:.2f} ms"
-    return f"{value:.2f} s"
-
-
-def fmt_bytes(value: Optional[float]) -> str:
-    if value is None or math.isnan(value):
-        return "-"
-    units = ["B", "KB", "MB", "GB", "TB", "PB"]
-    idx = 0
-    current = value
-    while current >= 1024 and idx < len(units) - 1:
-        current /= 1024
-        idx += 1
-    return f"{current:.2f} {units[idx]}"
-
-
-def fmt_percentage(value: Optional[float]) -> str:
-    if value is None or math.isnan(value):
-        return "-"
-    return f"{value * 100:4.1f}%"
-
-
-def section(title: str, body: Iterable[str]) -> List[str]:
-    lines = [f"## {title}"]
-    lines.extend(body)
-    lines.append("")
-    return lines
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-
-def main(argv: Optional[Sequence[str]] = None) -> None:
-    args = parse_args(argv)
-    end = parse_timestamp(args.end, default=dt.datetime.now(dt.timezone.utc))
-    if end is None:
-        raise SystemExit("End timestamp could not be determined.")
-    start = parse_timestamp(args.start)
-    if start is None:
-        duration = parse_duration(args.duration)
-        start = end - duration
-    assert start is not None
-    duration_seconds = max((end - start).total_seconds(), 1.0)
-    window = format_window(duration_seconds)
-    rate_window = compute_rate_window(duration_seconds)
-    subquery_step = compute_subquery_step(duration_seconds)
-
-    client = PrometheusClient(args.prom_url, timeout=args.timeout, default_time=end)
-    store_stats = fetch_store_statistics(args.store_url, timeout=args.timeout)
-    store_totals = extract_store_totals(store_stats)
-    lines: List[str] = [
-        f"Agent Lightning benchmark report",
-        f"Range: {start.isoformat()} — {end.isoformat()} ({duration_seconds:.0f}s window)",
-        f"Prometheus: {args.prom_url}",
-        f"Store: {args.store_url}",
-        "",
-    ]
-
-    # Throughput
-    throughput_rows = gather_collection_throughput(
-        client, collections=STORE_TOTAL_COLLECTIONS, duration_seconds=duration_seconds
-    )
-    rows: List[List[str]] = []
-    for item in throughput_rows:
-        store_total = store_totals.get(item.name)
-        if store_total is not None:
-            count_value: Optional[int] = store_total
-        elif item.count is not None:
-            count_value = int(item.count)
-        else:
-            count_value = None
-        if count_value is None:
-            count_str = "-"
-        else:
-            count_str = f"{count_value:,}"
-        if count_value is not None and duration_seconds > 0:
-            per_sec_value = float(count_value) / duration_seconds
-        else:
-            per_sec_value = item.per_sec
-        rows.append([item.name, count_str, fmt_rate(per_sec_value)])
-    lines.extend(
-        section(
-            "Rollout / Attempt / Span / Resource / Worker Throughput",
-            render_table(["Collection", "Count", "Per Sec"], rows),
-        )
-    )
-
-    # Store internals
-    store_stats, store_overall = gather_store_methods(client, window, rate_window, subquery_step)
-    store_rows: List[List[str]] = [
-        [
-            stat.method,
-            fmt_rate(stat.ops_mean),
-            fmt_rate(stat.ops_max),
-            fmt_rate(stat.ops_min),
-            fmt_latency(stat.p50),
-            fmt_latency(stat.p95),
-            fmt_latency(stat.p99),
-        ]
-        for stat in store_stats
-    ]
-    store_lines = render_table(
-        ["Method", "Mean Ops/s", "Max Ops/s", "Min Ops/s", "P50", "P95", "P99"],
-        store_rows,
-    )
-    store_lines.append(
-        f"Overall: mean={fmt_rate(store_overall['ops_mean'])}, "
-        f"max={fmt_rate(store_overall['ops_max'])}, "
-        f"min={fmt_rate(store_overall['ops_min'])}, "
-        f"P50={fmt_latency(store_overall['p50'])}, "
-        f"P95={fmt_latency(store_overall['p95'])}, "
-        f"P99={fmt_latency(store_overall['p99'])}"
-    )
-    lines.extend(section("CollectionBasedLightningStore", store_lines))
-
-    rollout_outcomes = gather_rollout_outcomes(client, window, rate_window)
-    rollout_rows = [
-        [
-            stat.status,
-            fmt_rate(stat.rate),
-            fmt_latency(stat.p25),
-            fmt_latency(stat.p50),
-            fmt_latency(stat.p75),
-            fmt_latency(stat.max_latency),
-        ]
-        for stat in rollout_outcomes
-    ]
-    lines.extend(
-        section("Rollout Outcomes", render_table(["Status", "Rate", "P25", "P50", "P75", "Max"], rollout_rows))
-    )
-
-    # HTTP traffic
-    http_paths, http_overall = gather_http_paths(client, window, rate_window, subquery_step)
-    http_rows: List[List[str]] = [
-        [
-            stat.method,
-            stat.path,
-            fmt_rate(stat.qps_mean),
-            fmt_rate(stat.qps_max),
-            fmt_rate(stat.qps_min),
-            fmt_latency(stat.p50),
-            fmt_latency(stat.p95),
-            fmt_latency(stat.p99),
-        ]
-        for stat in http_paths
-    ]
-    http_lines = render_table(
-        ["Method", "Path", "Mean Req/s", "Max Req/s", "Min Req/s", "P50", "P95", "P99"], http_rows
-    )
-    http_lines.append(
-        f"Overall HTTP: mean={fmt_rate(http_overall['qps_mean'])}, "
-        f"max={fmt_rate(http_overall['qps_max'])}, "
-        f"min={fmt_rate(http_overall['qps_min'])}, "
-        f"P50={fmt_latency(http_overall['p50'])}, "
-        f"P95={fmt_latency(http_overall['p95'])}, "
-        f"P99={fmt_latency(http_overall['p99'])}"
-    )
-    lines.extend(section("HTTP Endpoints", http_lines))
-
-    # Diagnostics
-    diag = gather_diagnostics(client, window)
-    diagnostics_blocks: List[List[str]] = []
-
-    mongo_ops = cast(Dict[str, float], diag.get("mongo_ops", {}))
-    mongo_ops_rows = [
-        [operation or "-", fmt_rate(rate)]
-        for operation, rate in sorted(mongo_ops.items(), key=lambda item: str(item[0]))
-    ]
-    diagnostics_blocks.append(render_table(["Mongo Operation", "Ops/s"], mongo_ops_rows))
-
-    mongo_opcounters = cast(Dict[str, float], diag.get("mongo_opcounters", {}))
-    mongo_opcounters_rows = [
-        [op_type or "-", fmt_rate(rate)]
-        for op_type, rate in sorted(mongo_opcounters.items(), key=lambda item: str(item[0]))
-    ]
-    diagnostics_blocks.append(render_table(["MongoDB Opcounter", "Ops/s"], mongo_opcounters_rows))
-
-    mongo_misc_rows: List[List[str]] = []
-    if diag.get("mongo_connections") is not None:
-        mongo_misc_rows.append(["MongoDB connections (avg)", f"{diag['mongo_connections']:.2f}"])
-    if mongo_misc_rows:
-        diagnostics_blocks.append(render_table(["Mongo Metric", "Value"], mongo_misc_rows))
-
-    node_rows: List[List[str]] = []
-    if diag.get("cpu_usage") is not None:
-        node_rows.append(["CPU usage", fmt_percentage(diag["cpu_usage"])])
-    mem_total = diag.get("memory_total")
-    mem_available = diag.get("memory_available")
-    if mem_total and mem_available:
-        used = mem_total - mem_available
-        node_rows.append(
-            ["Memory usage", f"{fmt_bytes(used)} / {fmt_bytes(mem_total)} ({fmt_percentage(used / mem_total)})"]
-        )
-    node_rows.append(["Network rx", f"{fmt_bytes(diag.get('network_rx'))}/s"])
-    node_rows.append(["Network tx", f"{fmt_bytes(diag.get('network_tx'))}/s"])
-    node_rows.append(["Disk read ops", fmt_rate(diag.get("disk_read_ops"))])
-    node_rows.append(["Disk read bytes", f"{fmt_bytes(diag.get('disk_read_bytes'))}/s"])
-    node_rows.append(["Disk write ops", fmt_rate(diag.get("disk_write_ops"))])
-    node_rows.append(["Disk write bytes", f"{fmt_bytes(diag.get('disk_write_bytes'))}/s"])
-    diagnostics_blocks.append(render_table(["Node Metric", "Value"], node_rows))
-
-    diagnostics_lines: List[str] = []
-    for idx, block in enumerate(diagnostics_blocks):
-        diagnostics_lines.extend(block)
-        if idx != len(diagnostics_blocks) - 1:
-            diagnostics_lines.append("")
-
-    lines.extend(section("Diagnostics", diagnostics_lines))
-
-    print("\n".join(lines))
-
-
-if __name__ == "__main__":  # pragma: no cover - manual execution
-    main()

tests/benchmark/benchmark_store.py

@@ -2,38 +2,15 @@
 
 """Benchmarking store performance by writing and querying spans from the store."""
 
-import argparse
 import asyncio
-import os
 import random
-import sys
-import threading
-from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple, cast
-
-from rich.console import Console
+from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple
 
 import agentlightning as agl
-from agentlightning.utils.otel import get_tracer
+from agentlightning.emitter.utils import get_tracer
 
 from .utils import flatten_dict, random_dict
 
-console = Console()
-
-MAX_RUNTIME_SECONDS = 45 * 60
-
-
-def _abort_due_to_timeout() -> None:
-    sys.stderr.write(f"[benchmark] Exiting after exceeding the {MAX_RUNTIME_SECONDS // 60} minute timeout.\n")
-    sys.stderr.flush()
-    os._exit(1)
-
-
-def _start_timeout_guard(timeout_seconds: float) -> threading.Timer:
-    timer = threading.Timer(timeout_seconds, _abort_due_to_timeout)
-    timer.daemon = True
-    timer.start()
-    return timer
-
 
 def generate_attributes() -> Dict[str, Any]:
     return flatten_dict(
@@ -46,51 +23,58 @@ def generate_attributes() -> Dict[str, Any]:
     )
 
 
-def make_agent(max_rounds: int, sleep_seconds: float) -> agl.LitAgent[str]:
-    @agl.rollout
-    async def agent(task: str, llm: agl.LLM):
-        tracer = get_tracer()
-        rounds = random.randint(1, max_rounds)
-        selected_round = random.randint(0, rounds - 1)
+@agl.rollout
+async def agent(task: str, llm: agl.LLM):
+    tracer = get_tracer()
+    rounds = random.randint(1, 10)
+    selected_round = random.randint(0, rounds - 1)
 
-        for i in range(rounds):
-            with tracer.start_as_current_span(f"agent{i}") as span:
-                # Nested Span
-                with tracer.start_as_current_span(f"round{i}_1") as span:
-                    await asyncio.sleep(random.uniform(0.0, sleep_seconds))
-                    span.set_attributes(generate_attributes())
-                    if i == selected_round:
-                        span.set_attribute("task", task)
+    for i in range(rounds):
+        with tracer.start_as_current_span(f"agent{i}") as span:
+            # Nested Span
+            with tracer.start_as_current_span(f"round{i}_1") as span:
+                await asyncio.sleep(random.uniform(0.0, 1.0))
+                span.set_attributes(generate_attributes())
+                if i == selected_round:
+                    span.set_attribute("task", task)
 
-                # Nested Span
-                with tracer.start_as_current_span(f"round{i}_2") as span:
-                    await asyncio.sleep(random.uniform(0.0, sleep_seconds))
-                    span.set_attributes(generate_attributes())
+            # Nested Span
+            with tracer.start_as_current_span(f"round{i}_2") as span:
+                await asyncio.sleep(random.uniform(0.0, 1.0))
+                span.set_attributes(generate_attributes())
 
-            if random.uniform(0, 1) < 0.5:
-                agl.emit_reward(random.uniform(0.0, 1.0))
+        if random.uniform(0, 1) < 0.5:
+            agl.emit_reward(random.uniform(0.0, 1.0))
 
-        # Final Span
-        with tracer.start_as_current_span("final") as span:
-            await asyncio.sleep(random.uniform(0.0, sleep_seconds))
-            span.set_attributes(generate_attributes())
+    # Final Span
+    with tracer.start_as_current_span("final") as span:
+        await asyncio.sleep(random.uniform(0.0, 1.0))
+        span.set_attributes(generate_attributes())
 
-        agl.emit_reward(random.uniform(1.0, 2.0))
-
-    return agent
+    agl.emit_reward(random.uniform(1.0, 2.0))
 
 
 def check_spans(spans: Sequence[agl.Span], task: str) -> None:
     """Check if the spans contain the task."""
-    found_task = any(span.attributes.get("task") == task for span in spans)
-
-    final_reward = agl.find_final_reward(spans)
-    if final_reward is None:
-        raise ValueError("Final reward is not found")
-    if not (final_reward >= 1 and final_reward <= 2):
-        raise ValueError(f"Final reward {final_reward} is not in the range of 1 to 2")
+    found_task = False
+    last_reward_in_12 = None
+    for span in spans:
+        if span.attributes.get("task") == task:
+            found_task = True
+        if span.name == agl.SpanNames.REWARD.value:
+            if span.attributes.get("reward") is None:
+                raise ValueError("Reward is not set for a reward span")
+            rew = float(span.attributes.get("reward"))  # type: ignore
+            if rew >= 1 and rew <= 2:
+                last_reward_in_12 = True
+            else:
+                last_reward_in_12 = False
     if not found_task:
         raise ValueError(f"Task {task} is not found in the spans")
+    if last_reward_in_12 is None:
+        raise ValueError("Last reward is not found")
+    elif not last_reward_in_12:
+        raise ValueError("Last reward is not in the range of 1 to 2")
 
 
 class AlgorithmBatch(agl.Algorithm):
@@ -138,7 +122,6 @@ class AlgorithmBatch(agl.Algorithm):
         submitted = 0
 
         while submitted < total_tasks:
-            print(f"Submitting batch {submitted} of {total_tasks}")
             batch_count = min(batch_size, total_tasks - submitted)
             batch_rollouts: List[Tuple[str, str]] = []
             await store.add_resources(
@@ -183,7 +166,6 @@ class AlgorithmBatch(agl.Algorithm):
         active_rollouts: Dict[str, str] = {}
 
         while completed < total_tasks:
-            console.print(f"Completed {completed} of {total_tasks} rollouts")
             if submitted < total_tasks and len(active_rollouts) < remaining_tasks:
                 batch_count = min(batch_size, total_tasks - submitted)
                 await store.add_resources(
@@ -235,7 +217,6 @@ class AlgorithmBatch(agl.Algorithm):
         async def handle_single(task_index: int) -> None:
             task_name = f"task-{task_index}"
             async with semaphore:
-                console.print(f"Submitting task {task_index} of {total_tasks}")
                 await store.add_resources(
                     {
                         "llm": agl.LLM(
@@ -260,70 +241,20 @@ class AlgorithmBatch(agl.Algorithm):
         await asyncio.gather(*all_tasks)
 
 
-def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Benchmark LightningStore implementations with synthetic rollouts.")
-    parser.add_argument("--store-url", default="http://localhost:4747", help="Lightning Store endpoint base URL.")
-    parser.add_argument(
-        "--mode",
-        choices=("batch", "batch_partial", "single"),
-        default="batch",
-        help="Algorithm mode to exercise different submission patterns.",
-    )
-    parser.add_argument("--total-tasks", type=int, default=128 * 128, help="Total number of rollouts to submit.")
-    parser.add_argument("--batch-size", type=int, default=128, help="Batch size for batch-style modes.")
-    parser.add_argument(
-        "--remaining-tasks",
-        type=int,
-        default=512,
-        help="Target number of in-flight rollouts before submitting more (batch_partial mode).",
-    )
-    parser.add_argument("--concurrency", type=int, default=32, help="Maximum concurrent rollouts for single mode.")
-    parser.add_argument("--n-runners", type=int, default=32, help="Number of runner processes to launch.")
-    parser.add_argument("--max-rounds", type=int, default=10, help="Maximum number of rounds for each rollout.")
-    parser.add_argument("--sleep-seconds", type=float, default=1.0, help="Sleep seconds for each rollout.")
-    args = parser.parse_args(argv)
-
-    if args.total_tasks <= 0:
-        parser.error("--total-tasks must be positive")
-    if args.n_runners <= 0:
-        parser.error("--n-runners must be positive")
-    if args.mode in {"batch", "batch_partial"} and (args.batch_size is None or args.batch_size <= 0):
-        parser.error("--batch-size must be positive for batch modes")
-    if args.mode == "batch_partial" and (args.remaining_tasks is None or args.remaining_tasks <= 0):
-        parser.error("--remaining-tasks must be positive for batch_partial mode")
-    if args.mode == "single" and (args.concurrency is None or args.concurrency <= 0):
-        parser.error("--concurrency must be positive for single mode")
-    if args.max_rounds <= 0:
-        parser.error("--max-rounds must be positive")
-    if args.sleep_seconds <= 0:
-        parser.error("--sleep-seconds must be positive")
-
-    return args
-
-
-def main(argv: Optional[Sequence[str]] = None) -> None:
-    args = parse_args(argv)
-    store = agl.LightningStoreClient(args.store_url)
-    timeout_guard = _start_timeout_guard(MAX_RUNTIME_SECONDS)
+def main() -> None:
+    store = agl.LightningStoreClient("http://localhost:4747")
     try:
         trainer = agl.Trainer(
             store=store,
-            algorithm=AlgorithmBatch(
-                mode=cast(Literal["batch", "batch_partial", "single"], args.mode),
-                total_tasks=args.total_tasks,
-                batch_size=args.batch_size,
-                remaining_tasks=args.remaining_tasks,
-                concurrency=args.concurrency,
-            ),
-            n_runners=args.n_runners,
+            algorithm=AlgorithmBatch(mode="batch", total_tasks=1024, batch_size=128),
+            n_runners=32,
             strategy={
                 "type": "cs",
                 "managed_store": False,
             },
         )
-        trainer.fit(make_agent(max_rounds=args.max_rounds, sleep_seconds=args.sleep_seconds))
+        trainer.fit(agent)
     finally:
-        timeout_guard.cancel()
         asyncio.run(store.close())
 
 

tests/emitter/test_annotation.py

@@ -1,72 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-from typing import Any, Dict
-
-import opentelemetry.trace as trace_api
-import pytest
-from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.trace import TraceFlags
-
-from agentlightning.emitter import annotation as annotation_module
-from agentlightning.emitter.annotation import emit_annotation
-from agentlightning.semconv import AGL_ANNOTATION
-
-
-class DummyReadableSpan(ReadableSpan):
-    def __init__(self) -> None:
-        super().__init__(
-            name="dummy",
-            context=trace_api.SpanContext(
-                trace_id=0x1,
-                span_id=0x2,
-                is_remote=False,
-                trace_flags=TraceFlags(TraceFlags.SAMPLED),
-                trace_state=trace_api.TraceState(),
-            ),
-            resource=Resource.create({}),
-        )
-
-    def __enter__(self) -> "DummyReadableSpan":
-        return self
-
-    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
-        return False
-
-
-class DummyTracer:
-    def __init__(self, span: DummyReadableSpan) -> None:
-        self._span = span
-        self.last_name: str | None = None
-        self.last_attributes: Dict[str, Any] | None = None
-
-    def start_span(self, name: str, attributes: Dict[str, Any] | None = None) -> DummyReadableSpan:
-        self.last_name = name
-        self.last_attributes = attributes or {}
-        return self._span
-
-
-def test_emit_annotation_flattens_and_respects_propagation(monkeypatch: pytest.MonkeyPatch) -> None:
-    span = DummyReadableSpan()
-    tracer = DummyTracer(span)
-    captured: Dict[str, Any] = {}
-
-    def fake_get_tracer(*_: Any, **kwargs: Any) -> DummyTracer:
-        captured["propagate"] = kwargs.get("use_active_span_processor")
-        return tracer
-
-    monkeypatch.setattr(annotation_module, "get_tracer", fake_get_tracer)
-
-    result = emit_annotation({"meta": {"tag": "foo"}, "score": 1.5}, propagate=False)
-
-    assert result is span
-    assert captured["propagate"] is False
-    assert tracer.last_name == AGL_ANNOTATION
-    assert tracer.last_attributes == {"meta.tag": "foo", "score": 1.5}
-
-
-def test_emit_annotation_rejects_non_primitive_values() -> None:
-    with pytest.raises(TypeError):
-        emit_annotation({"bad": {"set": {1}}})

tests/emitter/test_emitter.py

@@ -0,0 +1,105 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+from typing import Any, Dict, Optional
+
+import pytest
+from opentelemetry.semconv.attributes import exception_attributes
+
+from agentlightning.emitter import emit_exception, emit_message, emit_object
+from agentlightning.emitter import exception as exception_module
+from agentlightning.emitter import message as message_module
+from agentlightning.emitter import object as object_module
+from agentlightning.types.tracer import SpanAttributeNames, SpanNames
+
+
+class DummySpan:
+    def __init__(self) -> None:
+        self.recorded_exception: Optional[Exception] = None
+        self.status: Optional[Any] = None
+
+    def __enter__(self) -> "DummySpan":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
+        return False
+
+    def record_exception(self, exception: Exception) -> None:
+        self.recorded_exception = exception
+
+    def set_status(self, status: Any) -> None:
+        self.status = status
+
+
+class DummyTracer:
+    def __init__(self, span: DummySpan) -> None:
+        self._span = span
+        self.last_name: Optional[str] = None
+        self.last_attributes: Optional[Dict[str, Any]] = None
+
+    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
+        self.last_name = name
+        self.last_attributes = attributes or {}
+        return self._span
+
+
+def test_emit_message_valid(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = DummySpan()
+    tracer = DummyTracer(span)
+    monkeypatch.setattr(message_module, "get_tracer", lambda: tracer)
+
+    emit_message("hello world")
+
+    assert tracer.last_name == SpanNames.MESSAGE.value
+    assert tracer.last_attributes == {SpanAttributeNames.MESSAGE.value: "hello world"}
+
+
+def test_emit_message_requires_string(caplog: pytest.LogCaptureFixture) -> None:
+    # emit_message logs an error but doesn't raise an exception for invalid input
+    emit_message(123)  # type: ignore[arg-type]
+    assert "Message must be a string" in caplog.text
+
+
+def test_emit_object_serializes_payload(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = DummySpan()
+    tracer = DummyTracer(span)
+    monkeypatch.setattr(object_module, "get_tracer", lambda: tracer)
+
+    payload = {"foo": "bar", "baz": [1, 2, 3]}
+    emit_object(payload)
+
+    assert tracer.last_name == SpanNames.OBJECT.value
+    assert tracer.last_attributes is not None
+    assert json.loads(tracer.last_attributes[SpanAttributeNames.OBJECT.value]) == payload
+
+
+def test_emit_object_requires_json_serializable(caplog: pytest.LogCaptureFixture) -> None:
+    # emit_object logs an error but doesn't raise an exception for non-serializable input
+    emit_object(object())  # type: ignore[arg-type]
+    assert "Object must be JSON serializable" in caplog.text
+
+
+def test_emit_exception_records_exception(monkeypatch: pytest.MonkeyPatch) -> None:
+    span = DummySpan()
+    tracer = DummyTracer(span)
+    monkeypatch.setattr(exception_module, "get_tracer", lambda: tracer)
+
+    exc: Optional[Exception] = None
+    try:
+        raise ValueError("boom")
+    except ValueError as err:
+        emit_exception(err)
+        exc = err
+
+    assert tracer.last_name == SpanNames.EXCEPTION.value
+    assert tracer.last_attributes is not None
+    assert tracer.last_attributes[exception_attributes.EXCEPTION_TYPE] == "ValueError"
+    assert tracer.last_attributes[exception_attributes.EXCEPTION_MESSAGE] == "boom"
+    assert tracer.last_attributes[exception_attributes.EXCEPTION_ESCAPED] is True
+    assert span.recorded_exception is exc
+
+
+def test_emit_exception_requires_exception_instance(caplog: pytest.LogCaptureFixture) -> None:
+    # emit_exception logs an error but doesn't raise an exception for invalid input
+    emit_exception("boom")  # type: ignore[arg-type]
+    assert "Expected an BaseException instance" in caplog.text

tests/emitter/test_exception.py

@@ -1,72 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-from typing import Any, Dict, Optional
-
-import pytest
-from opentelemetry.semconv.attributes import exception_attributes
-
-from agentlightning.emitter import emit_exception
-from agentlightning.emitter import exception as exception_module
-from agentlightning.semconv import AGL_EXCEPTION
-
-
-class DummySpan:
-    def __init__(self) -> None:
-        self.recorded_exception: Optional[Exception] = None
-
-    def __enter__(self) -> "DummySpan":
-        return self
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
-        return False
-
-    def record_exception(self, exception: Exception) -> None:
-        self.recorded_exception = exception
-
-
-class DummyTracer:
-    def __init__(self, span: DummySpan) -> None:
-        self._span = span
-        self.last_name: Optional[str] = None
-        self.last_attributes: Optional[Dict[str, Any]] = None
-
-    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
-        self.last_name = name
-        self.last_attributes = attributes or {}
-        return self._span
-
-
-def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
-    tracer = DummyTracer(span)
-
-    def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
-        return tracer
-
-    monkeypatch.setattr(exception_module, "get_tracer", fake_get_tracer)
-    return tracer
-
-
-def test_emit_exception_records_exception(monkeypatch: pytest.MonkeyPatch) -> None:
-    span = DummySpan()
-    tracer = _stub_tracer(monkeypatch, span)
-
-    exc: Optional[Exception] = None
-    try:
-        raise ValueError("boom")
-    except ValueError as err:
-        emit_exception(err)
-        exc = err
-
-    assert tracer.last_name == AGL_EXCEPTION
-    assert tracer.last_attributes is not None
-    assert tracer.last_attributes[exception_attributes.EXCEPTION_TYPE] == "ValueError"
-    assert tracer.last_attributes[exception_attributes.EXCEPTION_MESSAGE] == "boom"
-    assert tracer.last_attributes[exception_attributes.EXCEPTION_ESCAPED] is True
-    assert span.recorded_exception is exc
-
-
-def test_emit_exception_requires_exception_instance() -> None:
-    with pytest.raises(TypeError):
-        emit_exception("boom")  # type: ignore[arg-type]

tests/emitter/test_message.py

@@ -1,83 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, cast
-
-import pytest
-
-from agentlightning.emitter import emit_message
-from agentlightning.emitter import message as message_module
-from agentlightning.emitter.message import get_message_value
-from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
-from agentlightning.types.tracer import SpanLike
-
-
-@dataclass
-class FakeSpan:
-    attributes: Optional[Dict[str, Any]]
-
-
-class DummySpan:
-    def __enter__(self) -> "DummySpan":
-        return self
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
-        return False
-
-
-class DummyTracer:
-    def __init__(self, span: DummySpan) -> None:
-        self._span = span
-        self.last_name: Optional[str] = None
-        self.last_attributes: Optional[Dict[str, Any]] = None
-
-    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
-        self.last_name = name
-        self.last_attributes = attributes or {}
-        return self._span
-
-
-def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
-    tracer = DummyTracer(span)
-
-    def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
-        return tracer
-
-    monkeypatch.setattr(message_module, "get_tracer", fake_get_tracer)
-    return tracer
-
-
-def test_get_message_value_returns_string() -> None:
-    span = FakeSpan(attributes={LightningSpanAttributes.MESSAGE_BODY.value: "hello"})
-
-    assert get_message_value(cast(SpanLike, span)) == "hello"
-
-
-def test_get_message_value_returns_none_when_missing() -> None:
-    span = FakeSpan(attributes={})
-
-    assert get_message_value(cast(SpanLike, span)) is None
-
-
-def test_get_message_value_rejects_non_string() -> None:
-    span = FakeSpan(attributes={LightningSpanAttributes.MESSAGE_BODY.value: ["not", "string"]})
-
-    with pytest.raises(TypeError):
-        get_message_value(cast(SpanLike, span))
-
-
-def test_emit_message_valid(monkeypatch: pytest.MonkeyPatch) -> None:
-    span = DummySpan()
-    tracer = _stub_tracer(monkeypatch, span)
-
-    emit_message("hello world")
-
-    assert tracer.last_name == AGL_MESSAGE
-    assert tracer.last_attributes == {LightningSpanAttributes.MESSAGE_BODY.value: "hello world"}
-
-
-def test_emit_message_requires_string() -> None:
-    with pytest.raises(TypeError):
-        emit_message(123)  # type: ignore[arg-type]

tests/emitter/test_object.py

@@ -1,177 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import base64
-import json
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, cast
-
-import pytest
-
-from agentlightning.emitter import emit_object
-from agentlightning.emitter import object as object_module
-from agentlightning.emitter.object import encode_object, get_object_value
-from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
-from agentlightning.types.tracer import SpanLike
-
-
-@dataclass
-class FakeSpan:
-    attributes: Optional[Dict[str, Any]]
-
-
-class DummySpan:
-    def __enter__(self) -> "DummySpan":
-        return self
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
-        return False
-
-
-class DummyTracer:
-    def __init__(self, span: DummySpan) -> None:
-        self._span = span
-        self.last_name: Optional[str] = None
-        self.last_attributes: Optional[Dict[str, Any]] = None
-
-    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
-        self.last_name = name
-        self.last_attributes = attributes or {}
-        return self._span
-
-
-def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
-    tracer = DummyTracer(span)
-
-    def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
-        return tracer
-
-    monkeypatch.setattr(object_module, "get_tracer", fake_get_tracer)
-    return tracer
-
-
-def test_encode_object_for_primitives() -> None:
-    encoded = encode_object("hello")
-
-    assert encoded == {
-        LightningSpanAttributes.OBJECT_TYPE.value: "str",
-        LightningSpanAttributes.OBJECT_LITERAL.value: "hello",
-    }
-
-
-def test_encode_object_for_bytes() -> None:
-    payload = b"binary"
-    encoded = encode_object(payload)
-    expected_literal = base64.b64encode(payload).decode("utf-8")
-
-    assert encoded == {
-        LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
-        LightningSpanAttributes.OBJECT_LITERAL.value: expected_literal,
-    }
-
-
-def test_encode_object_for_dict_serializes_json() -> None:
-    payload = {"key": 1, "nested": [1, 2]}
-    encoded = encode_object(payload)
-
-    assert encoded[LightningSpanAttributes.OBJECT_TYPE.value] == "dict"
-    assert json.loads(encoded[LightningSpanAttributes.OBJECT_JSON.value]) == payload
-
-
-def test_encode_object_raises_for_unserializable() -> None:
-    with pytest.raises(RuntimeError):
-        encode_object(object())
-
-
-def test_get_object_value_from_json_attribute() -> None:
-    payload = {"answer": 42}
-    span = FakeSpan(
-        attributes={
-            LightningSpanAttributes.OBJECT_JSON.value: json.dumps(payload),
-            LightningSpanAttributes.OBJECT_TYPE.value: "dict",
-        }
-    )
-
-    assert get_object_value(cast(SpanLike, span)) == payload
-
-
-def test_get_object_value_from_literal_types() -> None:
-    span = FakeSpan(
-        attributes={
-            LightningSpanAttributes.OBJECT_LITERAL.value: "123",
-            LightningSpanAttributes.OBJECT_TYPE.value: "int",
-        }
-    )
-    assert get_object_value(cast(SpanLike, span)) == 123
-
-    span = FakeSpan(
-        attributes={
-            LightningSpanAttributes.OBJECT_LITERAL.value: "3.14",
-            LightningSpanAttributes.OBJECT_TYPE.value: "float",
-        }
-    )
-    assert get_object_value(cast(SpanLike, span)) == pytest.approx(3.14)  # type: ignore
-
-    span = FakeSpan(
-        attributes={
-            LightningSpanAttributes.OBJECT_LITERAL.value: "true",
-            LightningSpanAttributes.OBJECT_TYPE.value: "bool",
-        }
-    )
-    assert get_object_value(cast(SpanLike, span)) is True
-
-    literal = base64.b64encode(b"hi").decode("utf-8")
-    span = FakeSpan(
-        attributes={
-            LightningSpanAttributes.OBJECT_LITERAL.value: literal,
-            LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
-        }
-    )
-    assert get_object_value(cast(SpanLike, span)) == b"hi"
-
-
-def test_get_object_value_returns_none_when_missing() -> None:
-    span = FakeSpan(attributes={})
-
-    assert get_object_value(cast(SpanLike, span)) is None
-
-
-def test_get_object_value_raises_for_invalid_json() -> None:
-    span = FakeSpan(
-        attributes={
-            LightningSpanAttributes.OBJECT_JSON.value: "not json",
-        }
-    )
-
-    with pytest.raises(RuntimeError):
-        get_object_value(cast(SpanLike, span))
-
-
-def test_emit_object_serializes_payload(monkeypatch: pytest.MonkeyPatch) -> None:
-    span = DummySpan()
-    tracer = _stub_tracer(monkeypatch, span)
-
-    payload = {"foo": "bar", "baz": [1, 2, 3]}
-    emit_object(payload)
-
-    assert tracer.last_name == AGL_OBJECT
-    assert tracer.last_attributes is not None
-    assert json.loads(tracer.last_attributes[LightningSpanAttributes.OBJECT_JSON.value]) == payload
-
-
-def test_emit_object_requires_json_serializable() -> None:
-    with pytest.raises(RuntimeError):
-        emit_object(object())  # type: ignore[arg-type]
-
-
-def test_get_object_value_raises_for_unknown_literal_type() -> None:
-    span = FakeSpan(
-        attributes={
-            LightningSpanAttributes.OBJECT_LITERAL.value: "value",
-            LightningSpanAttributes.OBJECT_TYPE.value: "complex",
-        }
-    )
-
-    with pytest.raises(RuntimeError):
-        get_object_value(cast(SpanLike, span))

tests/emitter/test_reward.py

@@ -1,18 +1,16 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import importlib
 import json
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, cast
 
-import pytest
-
-reward_module = importlib.import_module("agentlightning.emitter.reward")
-from agentlightning.emitter.reward import emit_reward, get_rewards_from_span
-from agentlightning.reward import find_final_reward, find_reward_spans, get_reward_value, is_reward_span
-from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
-from agentlightning.types import SpanLike
-from agentlightning.utils.otel import make_link_attributes, make_tag_attributes
+from agentlightning.reward import (
+    find_final_reward,
+    find_reward_spans,
+    get_reward_value,
+    is_reward_span,
+)
+from agentlightning.types import SpanLike, SpanNames
 
 
 @dataclass
@@ -21,77 +19,10 @@ class FakeSpan:
     attributes: Optional[Dict[str, Any]] = None
 
 
-@dataclass
-class AttributeSpan:
-    attributes: Optional[Dict[str, Any]]
-
-
 def make_span(name: str, attributes: Optional[Dict[str, Any]] = None) -> SpanLike:
     return cast(SpanLike, FakeSpan(name=name, attributes=attributes))
 
 
-def _capture_emit_annotation(monkeypatch: pytest.MonkeyPatch) -> tuple[Dict[str, Any], object]:
-    captured: Dict[str, Any] = {}
-    sentinel = object()
-
-    def fake_emit_annotation(payload: Dict[str, Any], *, propagate: bool) -> object:
-        captured["payload"] = payload
-        captured["propagate"] = propagate
-        return sentinel
-
-    monkeypatch.setattr(reward_module, "emit_annotation", fake_emit_annotation)
-    return captured, sentinel
-
-
-def test_emit_reward_example_scalar(monkeypatch: pytest.MonkeyPatch) -> None:
-    captured, sentinel = _capture_emit_annotation(monkeypatch)
-
-    result = emit_reward(1.0)
-
-    assert result is sentinel
-    assert captured["propagate"] is True
-    dimensions = captured["payload"][LightningSpanAttributes.REWARD.value]
-    assert dimensions == [{"name": "primary", "value": 1.0}]
-
-
-def test_emit_reward_example_multi_dimensional(monkeypatch: pytest.MonkeyPatch) -> None:
-    captured, _ = _capture_emit_annotation(monkeypatch)
-
-    emit_reward({"task_completion": 1.0, "efficiency": 0.8}, primary_key="task_completion")
-
-    dimensions = captured["payload"][LightningSpanAttributes.REWARD.value]
-    assert dimensions == [
-        {"name": "task_completion", "value": 1.0},
-        {"name": "efficiency", "value": 0.8},
-    ]
-
-
-def test_emit_reward_example_with_links(monkeypatch: pytest.MonkeyPatch) -> None:
-    captured, _ = _capture_emit_annotation(monkeypatch)
-
-    link_attrs = make_link_attributes({"gen_ai.response.id": "response-123", "span_id": "abcd-efgh"})
-    emit_reward(0.5, attributes=link_attrs)
-
-    payload = captured["payload"]
-    assert payload[f"{LightningSpanAttributes.LINK.value}.0.key_match"] == "gen_ai.response.id"
-    assert payload[f"{LightningSpanAttributes.LINK.value}.0.value_match"] == "response-123"
-    reward_dimensions = payload[LightningSpanAttributes.REWARD.value]
-    assert reward_dimensions == [{"name": "primary", "value": 0.5}]
-
-
-def test_emit_reward_example_with_tags(monkeypatch: pytest.MonkeyPatch) -> None:
-    captured, _ = _capture_emit_annotation(monkeypatch)
-
-    tag_attrs = make_tag_attributes(["fast", "reliable"])
-    emit_reward(0.7, attributes=tag_attrs)
-
-    payload = captured["payload"]
-    assert payload[f"{LightningSpanAttributes.TAG.value}.0"] == "fast"
-    assert payload[f"{LightningSpanAttributes.TAG.value}.1"] == "reliable"
-    reward_dimensions = payload[LightningSpanAttributes.REWARD.value]
-    assert reward_dimensions == [{"name": "primary", "value": 0.7}]
-
-
 def test_get_reward_value_from_agentops_dict() -> None:
     span = make_span(
         name="any",
@@ -112,7 +43,7 @@ def test_get_reward_value_from_agentops_json_string() -> None:
 
 def test_get_reward_value_from_reward_span_attributes() -> None:
     span = make_span(
-        name=AGL_ANNOTATION,
+        name=SpanNames.REWARD.value,
         attributes={"reward": 0.75},
     )
 
@@ -142,7 +73,7 @@ def test_is_reward_span_false_when_no_reward() -> None:
 
 def test_find_reward_spans_filters_correctly() -> None:
     reward_span = make_span(
-        name=AGL_ANNOTATION,
+        name=SpanNames.REWARD.value,
         attributes={"reward": 2.0},
     )
     non_reward_span = make_span(name="other", attributes={})
@@ -155,7 +86,7 @@ def test_find_reward_spans_filters_correctly() -> None:
 def test_find_final_reward_returns_last_reward_value() -> None:
     spans = [
         make_span(name="first", attributes={}),
-        make_span(name=AGL_ANNOTATION, attributes={"reward": 1.0}),
+        make_span(name=SpanNames.REWARD.value, attributes={"reward": 1.0}),
         make_span(name="agentops", attributes={"agentops.task.output": {"type": "reward", "value": 5.5}}),
     ]
 
@@ -169,75 +100,3 @@ def test_find_final_reward_returns_none_when_no_reward() -> None:
     ]
 
     assert find_final_reward(spans) is None
-
-
-def test_emit_reward_scalar_converts_to_primary_dimension(monkeypatch: pytest.MonkeyPatch) -> None:
-    captured: Dict[str, Any] = {}
-    sentinel_span = object()
-
-    def fake_emit_annotation(payload: Dict[str, Any], *, propagate: bool) -> object:
-        captured["payload"] = payload
-        captured["propagate"] = propagate
-        return sentinel_span
-
-    monkeypatch.setattr(reward_module, "emit_annotation", fake_emit_annotation)
-
-    result = emit_reward(2, attributes={"extra": "value"}, propagate=False)
-
-    assert result is sentinel_span
-    assert captured["propagate"] is False
-    rewards = captured["payload"][LightningSpanAttributes.REWARD.value]
-    assert rewards == [{"name": "primary", "value": 2.0}]
-    assert captured["payload"]["extra"] == "value"
-
-
-def test_emit_reward_dict_requires_primary_key(monkeypatch: pytest.MonkeyPatch) -> None:
-    captured: Dict[str, Any] = {}
-
-    def fake_emit_annotation(payload: Dict[str, Any], *, propagate: bool) -> Dict[str, Any]:
-        captured["payload"] = payload
-        return payload
-
-    monkeypatch.setattr(reward_module, "emit_annotation", fake_emit_annotation)
-
-    emit_reward({"score": 0.8, "other": 0.2}, primary_key="score")
-
-    rewards = captured["payload"][LightningSpanAttributes.REWARD.value]
-    assert [dim["name"] for dim in rewards] == ["score", "other"]
-
-    with pytest.raises(ValueError):
-        emit_reward({"score": 0.8}, primary_key=None)
-
-    with pytest.raises(ValueError):
-        emit_reward({"score": 0.8}, primary_key="missing")
-
-    with pytest.raises(ValueError):
-        emit_reward({"score": "bad"}, primary_key="score")
-
-
-def test_emit_reward_rejects_non_numeric() -> None:
-    with pytest.raises(TypeError):
-        emit_reward("bad")  # type: ignore[arg-type]
-
-
-def test_get_rewards_from_span_roundtrip() -> None:
-    attributes = {
-        f"{LightningSpanAttributes.REWARD.value}.0.name": "primary",
-        f"{LightningSpanAttributes.REWARD.value}.0.value": 1.0,
-        f"{LightningSpanAttributes.REWARD.value}.1.name": "aux",
-        f"{LightningSpanAttributes.REWARD.value}.1.value": 0.25,
-    }
-    span = AttributeSpan(attributes=attributes)
-
-    rewards = get_rewards_from_span(cast(SpanLike, span))
-
-    assert rewards == [
-        RewardPydanticModel(name="primary", value=1.0),
-        RewardPydanticModel(name="aux", value=0.25),
-    ]
-
-
-def test_get_rewards_from_span_returns_empty_when_missing() -> None:
-    span = AttributeSpan(attributes={})
-
-    assert get_rewards_from_span(cast(SpanLike, span)) == []

tests/llm_proxy/test_llm_proxy_gpu.py

@@ -127,7 +127,6 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer, otlp_enabled:
         print(f">>> Span: {span.name}")
         print(f">>> Start time: {span.start_time}")
         print(f">>> End time: {span.end_time}")
-        print(f">>> Attributes: {span.attributes.keys()}")
         assert span.start_time is not None, f"Span {span.name} has no start time"
         assert span.end_time is not None, f"Span {span.name} has no end time"
 

tests/runner/test_agent_integration.py

@@ -67,7 +67,8 @@ async def test_runner_integration_basic_rollout() -> None:
     assert rollouts and rollouts[0].status == "succeeded"
     attempts = await store.query_attempts(rollouts[0].rollout_id)
     spans = await store.query_spans(rollouts[0].rollout_id, attempts[-1].attempt_id)
-    assert any(span.attributes.get("agentlightning.reward.0.value") == 1.0 for span in spans)
+    print(store.__dict__)
+    assert any(span.attributes.get("reward") == 1.0 for span in spans)
 
 
 @pytest.mark.asyncio
@@ -228,10 +229,8 @@ async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenA
 
         last_spans = [span for span in spans if span.sequence_id == max(span.sequence_id for span in spans)]
         assert len(last_spans) == 1
-        assert last_spans[0].name == "agentlightning.annotation"
-        assert (
-            last_spans[0].attributes.get("agentlightning.reward.0.value") == 0.5
-        ), f"Expected reward to be 0.5, found {last_spans[0].attributes}"
+        assert last_spans[0].name == "agentlightning.reward"
+        assert last_spans[0].attributes.get("reward") == 0.5
     finally:
         teardown_runner(runner)
         await proxy.stop()

tests/runner/test_agent_runner.py

@@ -17,11 +17,10 @@ from agentlightning.litagent import LitAgent
 from agentlightning.reward import emit_reward, find_final_reward
 from agentlightning.runner import LitAgentRunner
 from agentlightning.runner.base import Runner
-from agentlightning.semconv import AGL_ANNOTATION
 from agentlightning.store.base import UNSET, LightningStore, Unset
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.tracer.base import Tracer
-from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, Worker
+from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, SpanNames, Worker
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -260,7 +259,7 @@ async def test_step_emits_reward_for_float_result() -> None:
 
     rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
     spans = await store.query_spans(rollout_id, attempt_id)
-    rewards = [span.attributes.get("agentlightning.reward.0.value") for span in spans if span.name == AGL_ANNOTATION]
+    rewards = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
     assert rewards == [0.75]
 
 
@@ -287,7 +286,7 @@ async def test_step_handles_non_llm_resource() -> None:
 
     rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
     spans = await store.query_spans(rollout_id, attempt_id)
-    rewards = [span.attributes.get("agentlightning.reward.0.value") for span in spans if span.name == AGL_ANNOTATION]
+    rewards = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
     assert rewards == [0.1]
 
 
@@ -513,9 +512,7 @@ async def test_agent_emits_multiple_rewards() -> None:
 
     rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
     spans = await store.query_spans(rollout_id, attempt_id)
-    reward_values = [
-        span.attributes.get("agentlightning.reward.0.value") for span in spans if span.name == AGL_ANNOTATION
-    ]
+    reward_values = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
     assert reward_values == [0.2, 0.6]
 
 
@@ -706,72 +703,6 @@ async def test_step_with_custom_resources_returns_rollout() -> None:
     assert result.resources_id is not None
 
 
-@pytest.mark.asyncio
-async def test_step_registers_worker_id_on_start_rollout(monkeypatch: pytest.MonkeyPatch) -> None:
-    """runner.step should pass the formatted worker ID down to the store."""
-
-    class WorkerAwareAgent(LitAgent[Dict[str, Any]]):
-        def validation_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> float:
-            return 1.0
-
-    agent = WorkerAwareAgent()
-    runner, store, _ = await setup_runner(agent)
-
-    expected_worker_label = runner.get_worker_id()
-    captured: Dict[str, Optional[str]] = {}
-    original_start_rollout = store.start_rollout
-
-    async def wrapped_start_rollout(*args: Any, **kwargs: Any):
-        captured["worker_id"] = kwargs.get("worker_id")
-        return await original_start_rollout(*args, **kwargs)
-
-    monkeypatch.setattr(store, "start_rollout", wrapped_start_rollout)
-
-    try:
-        await runner.step({"task": "worker-aware"})
-    finally:
-        teardown_runner(runner)
-
-    assert captured["worker_id"] == expected_worker_label
-
-
-@pytest.mark.asyncio
-async def test_iter_passes_worker_id_to_dequeue(monkeypatch: pytest.MonkeyPatch) -> None:
-    """iter() should poll the store with the formatted worker identifier."""
-
-    class IdleAgent(LitAgent[Dict[str, Any]]):
-        def validation_rollout(
-            self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
-        ) -> float:  # pragma: no cover - not invoked
-            return 0.0
-
-    agent = IdleAgent()
-    runner, store, _ = await setup_runner(agent, poll_interval=0.01)
-
-    expected_worker_label = runner.get_worker_id()
-    captured: Dict[str, Optional[str]] = {}
-    event = ThreadingEvent()
-
-    async def fake_dequeue(*, worker_id: Optional[str] = None):
-        captured["worker_id"] = worker_id
-        event.set()
-        return None
-
-    async def fast_sleep(self: LitAgentRunner[Any], event: Optional[ExecutionEvent] = None) -> None:
-        if event is not None:
-            event.set()
-
-    monkeypatch.setattr(store, "dequeue_rollout", fake_dequeue)
-    monkeypatch.setattr(LitAgentRunner, "_sleep_until_next_poll", fast_sleep)
-
-    try:
-        await runner.iter(event=event)
-    finally:
-        teardown_runner(runner)
-
-    assert captured["worker_id"] == expected_worker_label
-
-
 @pytest.mark.asyncio
 async def test_emit_heartbeat_updates_worker_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
     snapshot = {"cpu_pct": 42.0, "mem_pct": 10.5}

tests/store/dummy_store.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Literal, Optional
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -10,7 +10,6 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
-    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
@@ -43,9 +42,8 @@ class DummyLightningStore(LightningStore):
         resources_id: Optional[str] = None,
         config: Optional[RolloutConfig] = None,
         metadata: Optional[Dict[str, Any]] = None,
-        worker_id: Optional[str] = None,
     ) -> AttemptedRollout:
-        self.calls.append(("start_rollout", (input, mode, resources_id, config, metadata, worker_id), {}))
+        self.calls.append(("start_rollout", (input, mode, resources_id, config, metadata), {}))
         return self.return_values["start_rollout"]
 
     async def enqueue_rollout(
@@ -59,25 +57,12 @@ class DummyLightningStore(LightningStore):
         self.calls.append(("enqueue_rollout", (input, mode, resources_id, config, metadata), {}))
         return self.return_values["enqueue_rollout"]
 
-    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
-        self.calls.append(("enqueue_many_rollouts", (rollouts,), {}))
-        return self.return_values["enqueue_many_rollouts"]
-
     async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         self.calls.append(("dequeue_rollout", (worker_id,), {}))
         return self.return_values["dequeue_rollout"]
 
-    async def dequeue_many_rollouts(
-        self,
-        *,
-        limit: int = 1,
-        worker_id: Optional[str] = None,
-    ) -> Sequence[AttemptedRollout]:
-        self.calls.append(("dequeue_many_rollouts", (), {"limit": limit, "worker_id": worker_id}))
-        return self.return_values["dequeue_many_rollouts"]
-
-    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
-        self.calls.append(("start_attempt", (rollout_id, worker_id), {}))
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        self.calls.append(("start_attempt", (rollout_id,), {}))
         return self.return_values["start_attempt"]
 
     async def query_rollouts(self, *args: Any, **kwargs: Any) -> List[Rollout]:
@@ -116,21 +101,17 @@ class DummyLightningStore(LightningStore):
         self.calls.append(("query_resources", args, kwargs))
         return self.return_values["query_resources"]
 
-    async def add_span(self, span: Span) -> Optional[Span]:
+    async def add_span(self, span: Span) -> Span:
         self.calls.append(("add_span", (span,), {}))
         return self.return_values["add_span"]
 
-    async def add_many_spans(self, spans: Sequence[Span]) -> List[Span]:
-        self.calls.append(("add_many_spans", (spans,), {}))
-        return self.return_values["add_many_spans"]
-
     async def add_otel_span(
         self,
         rollout_id: str,
         attempt_id: str,
         readable_span: ReadableSpan,
         sequence_id: Optional[int] = None,
-    ) -> Optional[Span]:
+    ) -> Span:
         self.calls.append(("add_otel_span", (rollout_id, attempt_id, readable_span, sequence_id), {}))
         return self.return_values["add_otel_span"]
 
@@ -142,10 +123,6 @@ class DummyLightningStore(LightningStore):
         self.calls.append(("get_next_span_sequence_id", (rollout_id, attempt_id), {}))
         return self.return_values["get_next_span_sequence_id"]
 
-    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> List[int]:
-        self.calls.append(("get_many_span_sequence_ids", (rollout_attempt_ids,), {}))
-        return self.return_values["get_many_span_sequence_ids"]
-
     async def query_spans(self, *args: Any, **kwargs: Any) -> List[Span]:
         self.calls.append(("query_spans", args, kwargs))
         return self.return_values["query_spans"]
@@ -229,13 +206,10 @@ def minimal_dummy_store() -> DummyLightningStore:
             "update_resources": None,
             "get_resources_by_id": None,
             "get_latest_resources": None,
-            "query_resources": [],
             "add_span": None,
-            "add_many_spans": [],
             "add_otel_span": None,
             "wait_for_rollouts": [],
             "get_next_span_sequence_id": 0,
-            "get_many_span_sequence_ids": [],
             "query_spans": [],
             "update_rollout": None,
             "update_attempt": None,

tests/store/test_client_server.py

@@ -18,16 +18,7 @@ from yarl import URL
 from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
-from agentlightning.types import (
-    LLM,
-    EnqueueRolloutRequest,
-    OtelResource,
-    PaginatedResult,
-    PromptTemplate,
-    RolloutConfig,
-    Span,
-    TraceStatus,
-)
+from agentlightning.types import LLM, OtelResource, PaginatedResult, PromptTemplate, RolloutConfig, Span, TraceStatus
 from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncherArgs
 
 
@@ -150,206 +141,6 @@ async def test_server_accepts_custom_launcher_args(store_fixture: LightningStore
         await server.stop()
 
 
-@pytest.mark.asyncio
-async def test_server_client_statistics_match(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
-    """Server and client should report identical statistics."""
-    server, client = server_client
-    await client.start_rollout(input={"source": "statistics"})
-
-    server_stats = await server.statistics()
-    client_stats = await client.statistics()
-
-    assert {k: v for k, v in server_stats.items() if k != "uptime"} == {
-        k: v for k, v in client_stats.items() if k != "uptime"
-    }
-    assert server_stats["uptime"] < client_stats["uptime"]  # type: ignore
-    expected_name = server.store.__class__.__name__ if server.store is not None else server_stats["name"]  # type: ignore
-    assert server_stats["name"] == expected_name  # type: ignore
-    assert server_stats["total_rollouts"] >= 1  # type: ignore
-
-
-@pytest.mark.asyncio
-async def test_client_start_rollout_propagates_worker_id(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient],
-) -> None:
-    server, client = server_client
-    attempt = await client.start_rollout(input={"source": "remote-worker"}, worker_id="client-worker-start")
-
-    assert attempt.attempt.worker_id == "client-worker-start"
-    worker = await server.get_worker_by_id("client-worker-start")
-    assert worker is not None
-    assert worker.status == "busy"
-    assert worker.current_rollout_id == attempt.rollout_id
-    assert worker.current_attempt_id == attempt.attempt.attempt_id
-
-
-@pytest.mark.asyncio
-async def test_client_start_attempt_propagates_worker_id(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient],
-) -> None:
-    server, client = server_client
-    initial = await client.start_rollout(input={"source": "retry-worker"})
-    retry = await client.start_attempt(initial.rollout_id, worker_id="client-worker-retry")
-
-    assert retry.attempt.sequence_id == 2
-    assert retry.attempt.worker_id == "client-worker-retry"
-    worker = await server.get_worker_by_id("client-worker-retry")
-    assert worker is not None
-    assert worker.status == "busy"
-    assert worker.current_rollout_id == retry.rollout_id
-    assert worker.current_attempt_id == retry.attempt.attempt_id
-
-
-@pytest.mark.asyncio
-async def test_client_enqueue_many_rollouts_uses_batch_payload(monkeypatch: MonkeyPatch) -> None:
-    client = LightningStoreClient("http://localhost:9000")
-    captured: Dict[str, Any] = {}
-
-    async def fake_request_json(_, method: str, path: str, *, json: Any = None, params: Any = None) -> Any:
-        captured.update({"method": method, "path": path, "json": json})
-        count = len(json["rollouts"]) if json and "rollouts" in json else 0  # type: ignore[index]
-        return [{"rollout_id": f"bulk-{idx}", "input": {"idx": idx}, "start_time": float(idx)} for idx in range(count)]
-
-    monkeypatch.setattr(LightningStoreClient, "_request_json", fake_request_json, raising=False)  # type: ignore
-
-    requests = [
-        EnqueueRolloutRequest(input={"idx": 0}, mode="train", metadata={"batch": "left"}),
-        EnqueueRolloutRequest(input={"idx": 1}, resources_id="resources-1"),
-    ]
-    rollouts = await client.enqueue_many_rollouts(requests)
-
-    assert captured["method"] == "post"
-    assert captured["path"] == "/queues/rollouts/enqueue"
-    assert len(captured["json"]["rollouts"]) == 2  # type: ignore[index]
-    assert captured["json"]["rollouts"][0]["mode"] == "train"  # type: ignore[index]
-    assert captured["json"]["rollouts"][1]["resources_id"] == "resources-1"  # type: ignore[index]
-    assert len(rollouts) == 2
-    await client.close()
-
-
-@pytest.mark.asyncio
-async def test_client_dequeue_methods_share_batch_logic(monkeypatch: MonkeyPatch) -> None:
-    client = LightningStoreClient("http://localhost:9001")
-
-    def attempt_payload(idx: int) -> Dict[str, Any]:
-        attempt_id = f"attempt-{idx}"
-        rollout_id = f"rollout-{idx}"
-        return {
-            "rollout_id": rollout_id,
-            "input": {"idx": idx},
-            "start_time": float(idx),
-            "status": "preparing",
-            "attempt": {
-                "rollout_id": rollout_id,
-                "attempt_id": attempt_id,
-                "sequence_id": 1,
-                "start_time": float(idx),
-                "status": "preparing",
-                "worker_id": "batch-worker",
-            },
-        }
-
-    payload_queue = [
-        [attempt_payload(0), attempt_payload(1)],
-        [attempt_payload(0)],
-    ]
-
-    class FakeResponse:
-        def __init__(self, body: Any):
-            self._body = body
-            self.status = 200
-
-        async def __aenter__(self) -> "FakeResponse":
-            return self
-
-        async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
-            return None
-
-        def raise_for_status(self) -> None:
-            return None
-
-        async def json(self) -> Any:
-            return self._body
-
-    class RecordingSession:
-        def __init__(self) -> None:
-            self.calls: list[Dict[str, Any]] = []
-
-        def post(self, url: str, json: Dict[str, Any]) -> FakeResponse:
-            self.calls.append({"url": url, "json": json})
-            body = payload_queue.pop(0)
-            return FakeResponse(body)
-
-    session = RecordingSession()
-
-    async def fake_get_session() -> RecordingSession:
-        return session
-
-    monkeypatch.setattr(client, "_get_session", fake_get_session)
-
-    batch = await client.dequeue_many_rollouts(limit=2, worker_id="batch-worker")
-    assert len(batch) == 2
-    single = await client.dequeue_rollout(worker_id="batch-worker")
-    assert single is not None
-    assert [call["json"]["limit"] for call in session.calls] == [2, 1]
-    await client.close()
-
-
-@pytest.mark.asyncio
-async def test_client_dequeue_many_rollouts_skips_network_for_non_positive_limit(monkeypatch: MonkeyPatch) -> None:
-    client = LightningStoreClient("http://localhost:9002")
-
-    async def fail_get_session() -> None:
-        pytest.fail("Client should not request a session when limit <= 0")
-
-    monkeypatch.setattr(client, "_get_session", fail_get_session)
-
-    assert await client.dequeue_many_rollouts(limit=0, worker_id="idle") == []
-    assert await client.dequeue_many_rollouts(limit=-5) == []
-    await client.close()
-
-
-@pytest.mark.asyncio
-async def test_client_concurrent_enqueue_many_rollouts(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient],
-) -> None:
-    _, client = server_client
-
-    async def enqueue_batch(batch_idx: int) -> list[str]:
-        requests = [EnqueueRolloutRequest(input={"batch": batch_idx, "idx": item}) for item in range(3)]
-        rollouts = await client.enqueue_many_rollouts(requests)
-        return [rollout.rollout_id for rollout in rollouts]
-
-    batches = await asyncio.gather(*(enqueue_batch(batch_idx) for batch_idx in range(5)))
-    all_ids = {rollout_id for batch in batches for rollout_id in batch}
-    assert len(all_ids) == 15
-
-    queried = await client.query_rollouts(limit=-1)
-    assert isinstance(queried, PaginatedResult)
-    assert queried.total >= 15
-
-
-@pytest.mark.asyncio
-async def test_client_concurrent_dequeue_many_rollouts(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient],
-) -> None:
-    server, client = server_client
-    requests = [EnqueueRolloutRequest(input={"idx": idx}) for idx in range(6)]
-    # Seed queue from the server to avoid races with background processing
-    await asyncio.gather(*(server.enqueue_rollout(**req.model_dump()) for req in requests))
-
-    async def consume(limit: int, worker: str):
-        return await client.dequeue_many_rollouts(limit=limit, worker_id=worker)
-
-    batches = await asyncio.gather(
-        consume(3, "worker-a"),
-        consume(3, "worker-b"),
-    )
-    claimed_ids = {attempt.rollout_id for batch in batches for attempt in batch}
-    assert len(claimed_ids) == 6
-    assert await client.dequeue_many_rollouts(limit=1) == []
-
-
 @pytest.mark.asyncio
 async def test_add_resources_via_server(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
     """Test that add_resources works correctly via server."""
@@ -504,7 +295,7 @@ async def test_client_server_end_to_end(
     dequeued = await server.dequeue_rollout(worker_id=server_worker_id)
     server_worker_after_dequeue = await server.get_worker_by_id(server_worker_id)
     assert server_worker_after_dequeue is not None
-    assert server_worker_after_dequeue.status == "busy"  # should be busy after dequeue
+    assert server_worker_after_dequeue.status == "idle"
     assert server_worker_after_dequeue.last_dequeue_time is not None
     dequeue_time = server_worker_after_dequeue.last_dequeue_time
     started_attempt = await server.start_attempt(queued_rollout.rollout_id)
@@ -582,7 +373,7 @@ async def test_client_server_end_to_end(
     assert dequeued_client is not None
     client_worker_after_dequeue = await client.get_worker_by_id(client_worker_id)
     assert client_worker_after_dequeue is not None
-    assert client_worker_after_dequeue.status == "busy"  # should be busy after dequeue
+    assert client_worker_after_dequeue.status == "idle"
     assert client_worker_after_dequeue.last_dequeue_time is not None
     client_dequeue_time = client_worker_after_dequeue.last_dequeue_time
     started_client_attempt = await client.start_attempt(dequeued_client.rollout_id)
@@ -603,7 +394,7 @@ async def test_client_server_end_to_end(
 
     client_span = _make_span(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id, 101, "client-span")
     stored_span = await client.add_span(client_span)
-    assert stored_span is not None and stored_span.name == "client-span"
+    assert stored_span.name == "client-span"
     assert await client.get_next_span_sequence_id(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id) == 102
 
     with patch("agentlightning.store.client_server.Span.from_opentelemetry", autospec=True) as mocked:
@@ -912,87 +703,6 @@ async def test_client_query_spans_filters_and_pagination(
     assert [span.span_id for span in paged] == [spans[1].span_id]
 
 
-@pytest.mark.asyncio
-async def test_server_get_many_span_sequence_ids_and_add_many_spans_mixed_batches(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient],
-) -> None:
-    server, _ = server_client
-
-    first = await server.start_rollout(input={"origin": "batch-server"})
-    second = await server.start_rollout(input={"origin": "batch-server-2"})
-    await server.update_rollout(first.rollout_id, status="requeuing")
-    retried = await server.start_attempt(first.rollout_id)
-
-    sequence_pairs = [
-        (first.rollout_id, first.attempt.attempt_id),
-        (second.rollout_id, second.attempt.attempt_id),
-        (first.rollout_id, retried.attempt.attempt_id),
-        (first.rollout_id, first.attempt.attempt_id),
-    ]
-    sequence_ids = await server.get_many_span_sequence_ids(sequence_pairs)
-    assert sequence_ids == [1, 1, 2, 3]
-
-    next_single = await server.get_next_span_sequence_id(first.rollout_id, first.attempt.attempt_id)
-    assert next_single == 4
-
-    batch_spans = [
-        _make_span(first.rollout_id, first.attempt.attempt_id, 10, "server-batch-1"),
-        _make_span(second.rollout_id, second.attempt.attempt_id, 11, "server-batch-2"),
-        _make_span(retried.rollout_id, retried.attempt.attempt_id, 12, "server-batch-retry"),
-    ]
-    stored_spans = await server.add_many_spans(batch_spans)
-    assert {span.name for span in stored_spans} == {
-        "server-batch-1",
-        "server-batch-2",
-        "server-batch-retry",
-    }
-
-    spans_first = await server.query_spans(first.rollout_id)
-    assert any(span.name == "server-batch-1" for span in spans_first)
-    assert any(span.name == "server-batch-retry" for span in spans_first)
-    spans_second = await server.query_spans(second.rollout_id)
-    assert any(span.name == "server-batch-2" for span in spans_second)
-
-
-@pytest.mark.asyncio
-async def test_client_handles_optional_span_results_and_batch_insert(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient],
-    mock_readable_span: ReadableSpan,
-) -> None:
-    _, client = server_client
-
-    first = await client.start_rollout(input={"origin": "client-span"})
-    second = await client.start_rollout(input={"origin": "client-span-2"})
-    first_attempt_id = first.attempt.attempt_id
-    second_attempt_id = second.attempt.attempt_id
-
-    base_span = _make_span(first.rollout_id, first_attempt_id, 1, "client-span-1")
-    stored = await client.add_span(base_span)
-    assert stored is not None
-    assert await client.add_span(base_span) is None
-
-    batch_spans = [
-        _make_span(first.rollout_id, first_attempt_id, 2, "client-span-2"),
-        _make_span(second.rollout_id, second_attempt_id, 1, "client-span-other"),
-        base_span,
-    ]
-    inserted = await client.add_many_spans(batch_spans)
-    assert [span.name for span in inserted] == ["client-span-2", "client-span-other"]
-
-    sequence_ids = await client.get_many_span_sequence_ids(
-        [
-            (first.rollout_id, first_attempt_id),
-            (second.rollout_id, second_attempt_id),
-            (first.rollout_id, "latest"),
-        ]
-    )
-    assert sequence_ids == [3, 2, 4]
-
-    with patch("agentlightning.store.client_server.Span.from_opentelemetry", autospec=True) as mocked_span_factory:
-        mocked_span_factory.return_value = base_span
-        assert await client.add_otel_span(first.rollout_id, first_attempt_id, mock_readable_span) is None
-
-
 @pytest.mark.asyncio
 async def test_concurrent_add_otel_span_sequence_ids_unique(
     server_client: Tuple[LightningStoreServer, LightningStoreClient], mock_readable_span: ReadableSpan
@@ -1011,7 +721,7 @@ async def test_concurrent_add_otel_span_sequence_ids_unique(
         spans = await asyncio.gather(
             *[client.add_otel_span(rollout_id, attempt_id, mock_readable_span) for _ in range(20)]
         )
-    sequence_ids = [span.sequence_id for span in spans]  # type: ignore
+    sequence_ids = [span.sequence_id for span in spans]
     assert len(set(sequence_ids)) == 20
     assert set(sequence_ids) == set(range(1, 21))
 

tests/store/test_collection_memory.py

@@ -3,22 +3,8 @@
 from __future__ import annotations
 
 import asyncio
-import threading
 import time
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Awaitable,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Literal,
-    Mapping,
-    Sequence,
-    Tuple,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Literal, Mapping, Sequence, Tuple, Union
 from uuid import uuid4
 
 import pytest
@@ -28,16 +14,12 @@ import agentlightning.store.collection.memory as memory_module
 from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, ListBasedCollection
 from agentlightning.store.collection.base import Collection
 from agentlightning.store.collection.memory import _item_matches_filters  # pyright: ignore[reportPrivateUsage]
-from agentlightning.store.collection.memory import _LoopAwareAsyncLock  # pyright: ignore[reportPrivateUsage]
-from agentlightning.store.collection.memory import _ThreadSafeAsyncLock  # pyright: ignore[reportPrivateUsage]
 from agentlightning.types import Rollout
 from tests.store.conftest import QueueItem, SampleItem
 
 if TYPE_CHECKING:
     from pymongo.asynchronous.database import AsyncDatabase
 
-    from agentlightning.store.collection.mongo import MongoLightningCollections
-
 
 def _build_collection(items: Iterable[SampleItem] = ()) -> ListBasedCollection[SampleItem]:
     return ListBasedCollection(list(items), SampleItem, ("partition", "index"))
@@ -110,20 +92,6 @@ async def test_list_collection_insert_duplicate_raises(sample_collection: Collec
         await sample_collection.insert([duplicate])
 
 
-@pytest.mark.asyncio()
-async def test_list_collection_insert_rejects_duplicate_payload(sample_collection: Collection[SampleItem]) -> None:
-    """Ensure duplicate items within the same insert batch are rejected."""
-    starting_size = await sample_collection.size()
-    dup_a = SampleItem(partition="omega", index=1, name="dup-a", status="new")
-    dup_b = SampleItem(partition="omega", index=1, name="dup-b", status="new")
-
-    with pytest.raises(ValueError, match="duplicate primary key"):
-        await sample_collection.insert([dup_a, dup_b])
-
-    assert await sample_collection.size() == starting_size
-    assert await sample_collection.get({"partition": {"exact": "omega"}}) is None
-
-
 @pytest.mark.asyncio()
 async def test_list_collection_insert_wrong_type(sample_collection: Collection[SampleItem]) -> None:
     class Another(BaseModel):
@@ -184,126 +152,6 @@ async def test_list_collection_upsert_updates_when_existing(sample_collection: C
     assert fetched == replacement
 
 
-@pytest.mark.asyncio()
-async def test_list_collection_upsert_get_or_insert_semantics(sample_collection: Collection[SampleItem]) -> None:
-    filters: Mapping[str, Any] = {"partition": {"exact": "beta"}, "index": {"exact": 2}}
-    original = await sample_collection.get(filters)
-    assert original is not None
-
-    replacement = SampleItem(
-        partition="beta",
-        index=2,
-        name="replacement",
-        status="queued",
-        tags=["patched"],
-        score=999,
-        rank=999,
-        updated_time=99.0,
-        payload={"priority": 99},
-        metadata="replacement",
-    )
-
-    await sample_collection.upsert([replacement], update_fields=[])
-
-    fetched = await sample_collection.get(filters)
-    assert fetched == original and fetched is not None
-    assert fetched.name == original.name
-    assert fetched.status == original.status
-
-
-@pytest.mark.asyncio()
-async def test_list_collection_upsert_updates_selected_fields(sample_collection: Collection[SampleItem]) -> None:
-    filters: Mapping[str, Any] = {"partition": {"exact": "beta"}, "index": {"exact": 1}}
-    original = await sample_collection.get(filters)
-    assert original is not None
-
-    incoming = SampleItem(
-        partition="beta",
-        index=1,
-        name="beta-incoming",
-        status="in-progress",
-        tags=["different"],
-        score=-1.0,
-        rank=42,
-        updated_time=123.45,
-        payload={"priority": -1},
-        metadata="incoming",
-    )
-
-    await sample_collection.upsert([incoming], update_fields=["status", "updated_time"])
-
-    fetched = await sample_collection.get(filters)
-    assert fetched is not None
-    assert fetched.status == incoming.status
-    assert fetched.updated_time == incoming.updated_time
-    # Ensure unspecified fields (e.g. name/tags) remain the same as the original document.
-    assert fetched.name == original.name
-    assert fetched.tags == original.tags
-
-
-@pytest.mark.asyncio()
-async def test_list_collection_update_returns_mutated_items(sample_collection: Collection[SampleItem]) -> None:
-    replacements = [
-        SampleItem(partition="alpha", index=1, name="alpha-new", status="patched"),
-        SampleItem(partition="delta", index=1, name="delta-new", status="patched"),
-    ]
-
-    returned = await sample_collection.update(replacements)
-    assert list(returned) == replacements
-
-    for expected in replacements:
-        fetched = await sample_collection.get(
-            {"partition": {"exact": expected.partition}, "index": {"exact": expected.index}}
-        )
-        assert fetched == expected
-
-    original_beta = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 1}})
-    assert original_beta is not None
-
-    partial_payload = SampleItem(
-        partition="beta", index=1, name="ignored", status="partial", metadata="updated-metadata"
-    )
-    partial_returned = await sample_collection.update([partial_payload], update_fields=["status", "metadata"])
-    assert len(partial_returned) == 1
-
-    fetched_partial = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 1}})
-    assert fetched_partial == partial_returned[0]
-    assert fetched_partial is not None
-    assert fetched_partial.status == partial_payload.status
-    assert fetched_partial.metadata == partial_payload.metadata
-    assert fetched_partial.name == original_beta.name
-
-
-@pytest.mark.asyncio()
-async def test_list_collection_upsert_returns_mutated_items(sample_collection: Collection[SampleItem]) -> None:
-    new_item = SampleItem(partition="omega", index=99, name="omega-new", status="queued")
-    inserted = await sample_collection.upsert([new_item])
-    assert list(inserted) == [new_item]
-
-    fetched_new = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 99}})
-    assert fetched_new == new_item
-
-    original_existing = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
-    assert original_existing is not None
-
-    incoming = SampleItem(
-        partition="beta",
-        index=2,
-        name="beta-incoming-new-name",
-        status="processing",
-        tags=["beta", "patched"],
-    )
-    updated = await sample_collection.upsert([incoming], update_fields=["status", "tags"])
-    assert len(updated) == 1
-
-    fetched_existing = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
-    assert fetched_existing == updated[0]
-    assert fetched_existing is not None
-    assert fetched_existing.status == incoming.status
-    assert fetched_existing.tags == incoming.tags
-    assert fetched_existing.name == original_existing.name
-
-
 @pytest.mark.asyncio()
 async def test_list_collection_delete_multiple_items(sample_collection: Collection[SampleItem]) -> None:
     await sample_collection.delete(
@@ -865,139 +713,6 @@ async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyVa
     assert await dict_key_value.size() == 1
 
 
-def test_thread_safe_async_lock_blocks_threads() -> None:
-    lock = _ThreadSafeAsyncLock()
-    allow_second = threading.Event()
-    second_has_lock = threading.Event()
-    release_first = threading.Event()
-
-    def first() -> None:
-        with lock:
-            allow_second.set()
-            release_first.wait()
-
-    def second() -> None:
-        allow_second.wait()
-        with lock:
-            second_has_lock.set()
-
-    t1 = threading.Thread(target=first)
-    t2 = threading.Thread(target=second)
-    t1.start()
-    t2.start()
-
-    assert allow_second.wait(timeout=1)
-    assert not second_has_lock.wait(0.05)
-
-    release_first.set()
-    t1.join(timeout=1)
-    t2.join(timeout=1)
-
-    assert second_has_lock.is_set()
-
-
-@pytest.mark.asyncio()
-async def test_thread_safe_async_lock_serializes_async_tasks() -> None:
-    lock = _ThreadSafeAsyncLock()
-    first_entered = asyncio.Event()
-    release_first = asyncio.Event()
-    second_acquired = asyncio.Event()
-
-    async def first() -> None:
-        async with lock:
-            first_entered.set()
-            await release_first.wait()
-
-    async def second() -> None:
-        await first_entered.wait()
-        async with lock:
-            second_acquired.set()
-
-    task1 = asyncio.create_task(first())
-    task2 = asyncio.create_task(second())
-
-    await first_entered.wait()
-    await asyncio.sleep(0)
-    assert not second_acquired.is_set()
-
-    release_first.set()
-    await asyncio.gather(task1, task2)
-
-    assert second_acquired.is_set()
-
-
-@pytest.mark.asyncio()
-async def test_loop_aware_async_lock_serializes_tasks() -> None:
-    lock = _LoopAwareAsyncLock()
-    first_entered = asyncio.Event()
-    release_first = asyncio.Event()
-    second_acquired = asyncio.Event()
-
-    async def first() -> None:
-        async with lock:
-            first_entered.set()
-            await release_first.wait()
-
-    async def second() -> None:
-        await first_entered.wait()
-        async with lock:
-            second_acquired.set()
-
-    task1 = asyncio.create_task(first())
-    task2 = asyncio.create_task(second())
-
-    await first_entered.wait()
-    await asyncio.sleep(0)
-    assert not second_acquired.is_set()
-
-    release_first.set()
-    await asyncio.gather(task1, task2)
-
-    assert second_acquired.is_set()
-
-
-@pytest.mark.asyncio()
-async def test_loop_aware_async_lock_reuses_loop_specific_lock() -> None:
-    lock = _LoopAwareAsyncLock()
-    first_lock: asyncio.Lock | None = None
-
-    async with lock as acquired:
-        first_lock = acquired
-        assert first_lock.locked()
-
-    assert first_lock is not None and not first_lock.locked()
-
-    async with lock as acquired_again:
-        assert acquired_again is first_lock
-
-
-@pytest.mark.asyncio()
-async def test_loop_aware_async_lock_distinguishes_event_loops() -> None:
-    lock = _LoopAwareAsyncLock()
-    main_loop_lock: asyncio.Lock | None = None
-
-    async with lock as acquired:
-        main_loop_lock = acquired
-
-    locks_from_threads: List[asyncio.Lock] = []
-
-    def _worker() -> None:
-        async def runner() -> None:
-            async with lock as acquired:
-                locks_from_threads.append(acquired)
-
-        asyncio.run(runner())
-
-    worker = threading.Thread(target=_worker)
-    worker.start()
-    worker.join(timeout=2)
-
-    assert worker.is_alive() is False
-    assert main_loop_lock is not None
-    assert locks_from_threads
-    assert locks_from_threads[0] is not main_loop_lock
-
-
 @pytest.mark.asyncio()
 async def test_dict_key_value_does_not_mutate_input_mapping(dict_key_value_data: Dict[str, int]) -> None:
     key_value = DictBasedKeyValue(dict_key_value_data)
@@ -1006,55 +721,6 @@ async def test_dict_key_value_does_not_mutate_input_mapping(dict_key_value_data:
     assert dict_key_value_data == {"alpha": 1, "beta": 2}
 
 
-@pytest.mark.asyncio()
-async def test_inmemory_atomic_read_only_skips_lock() -> None:
-    collections = memory_module.InMemoryLightningCollections(lock_type="asyncio")
-
-    class FailingLock:
-        async def __aenter__(self) -> None:
-            raise AssertionError("read-only atomic block should not acquire the lock")
-
-        async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
-            return None
-
-    collections._lock = {"default": FailingLock()}  # type: ignore[attr-defined]
-
-    async with collections.atomic(mode="r", snapshot=False):
-        # Should complete without touching the failing lock.
-        assert collections.rollouts is not None
-
-
-@pytest.mark.asyncio()
-async def test_inmemory_atomic_snapshot_or_write_acquires_lock() -> None:
-    collections = memory_module.InMemoryLightningCollections(lock_type="asyncio")
-
-    class RecordingLock:
-        def __init__(self) -> None:
-            self.enter_count = 0
-            self.exit_count = 0
-
-        async def __aenter__(self) -> None:
-            self.enter_count += 1
-
-        async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
-            self.exit_count += 1
-
-    lock = RecordingLock()
-    collections._lock = {"default": lock}  # type: ignore[attr-defined]
-
-    async with collections.atomic(mode="rw", snapshot=False):
-        assert collections.attempts is not None
-
-    assert (lock.enter_count, lock.exit_count) == (1, 1)
-
-    lock.enter_count = lock.exit_count = 0
-
-    async with collections.atomic(mode="r", snapshot=True):
-        assert collections.spans is not None
-
-    assert (lock.enter_count, lock.exit_count) == (1, 1)
-
-
 @pytest.mark.mongo
 @pytest.mark.asyncio()
 async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[Any]) -> None:
@@ -1104,31 +770,6 @@ async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[
         assert not await span_kv.has("span-123")
 
 
-@pytest.mark.mongo
-@pytest.mark.asyncio()
-async def test_mongo_based_collection_rejects_duplicate_payload(temporary_mongo_database: AsyncDatabase[Any]) -> None:
-    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
-
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
-        collection = MongoBasedCollection[Any](
-            client_pool,
-            temporary_mongo_database.name,
-            f"duplicate-check-{uuid4().hex}",
-            "partition-dup",
-            ["rollout_id"],
-            Rollout,
-        )
-        await collection.ensure_collection()
-        start_time = time.time()
-        first = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
-        duplicate = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
-
-        with pytest.raises(ValueError, match="duplicate primary key"):
-            await collection.insert([first, duplicate])
-
-        assert await collection.size() == 0
-
-
 @pytest.mark.mongo
 @pytest.mark.asyncio()
 async def test_mongo_ensure_collection_creates_partition_scoped_index(
@@ -1211,110 +852,3 @@ async def test_mongo_ensure_collection_repeats_without_altering_indexes(
                 unique_indexes.append((index["name"], list(index["key"].items())))  # type: ignore
 
         assert unique_indexes == [("uniq_partition_index", [("partition_id", 1), ("index", 1)])]
-
-
-async def _with_mongo_collections(
-    db: AsyncDatabase[Any],
-    callback: Callable[[MongoLightningCollections], Awaitable[Any]],
-) -> Any:
-    from agentlightning.store.collection.mongo import MongoClientPool, MongoLightningCollections
-
-    async with MongoClientPool(db.client) as client_pool:
-        collections = MongoLightningCollections(
-            client_pool=client_pool,
-            database_name=db.name,
-            partition_id=f"partition-{uuid4().hex}",
-        )
-        return await callback(collections)
-
-
-async def _initialize_counter(collections: MongoLightningCollections, key: str) -> None:
-    async def _init(coll: MongoLightningCollections) -> None:
-        await coll.span_sequence_ids.set(key, 0)
-
-    await collections.execute(_init, commit=False)
-
-
-async def _read_counter(collections: MongoLightningCollections, key: str) -> int:
-    async def _read(coll: MongoLightningCollections) -> int:
-        value = await coll.span_sequence_ids.get(key)
-        assert value is not None
-        return value
-
-    return await collections.execute(_read, commit=False)
-
-
-async def _contention_run(
-    collections: MongoLightningCollections,
-    *,
-    key: str,
-    commit: bool,
-    concurrency: int,
-) -> int:
-    read_lock = asyncio.Lock()
-    ready = asyncio.Event()
-    readers_seen = 0
-
-    async def _barrier() -> None:
-        nonlocal readers_seen
-        async with read_lock:
-            readers_seen += 1
-            if readers_seen == concurrency:
-                ready.set()
-        await ready.wait()
-
-    async def worker(_: int) -> None:
-        first_attempt = True
-
-        async def callback(coll: MongoLightningCollections) -> None:
-            nonlocal first_attempt
-            value = await coll.span_sequence_ids.get(key)
-            assert value is not None
-            if first_attempt:
-                await _barrier()
-                first_attempt = False
-            await asyncio.sleep(0)
-            await coll.span_sequence_ids.set(key, value + 1)
-
-        await collections.execute(callback, commit=commit, snapshot=True, mode="rw")
-
-    await asyncio.gather(*(worker(i) for i in range(concurrency)))
-    return await _read_counter(collections, key)
-
-
-@pytest.mark.mongo
-@pytest.mark.asyncio()
-async def test_mongo_execute_without_commit_allows_lost_updates(
-    temporary_mongo_database: AsyncDatabase[Any],
-) -> None:
-    async def scenario(collections: MongoLightningCollections) -> None:
-        counter_key = f"counter-{uuid4().hex}"
-        await _initialize_counter(collections, counter_key)
-        final_value = await _contention_run(
-            collections,
-            key=counter_key,
-            commit=False,
-            concurrency=6,
-        )
-        assert final_value == 1
-
-    await _with_mongo_collections(temporary_mongo_database, scenario)
-
-
-@pytest.mark.mongo
-@pytest.mark.asyncio()
-async def test_mongo_execute_with_commit_retries_until_success(
-    temporary_mongo_database: AsyncDatabase[Any],
-) -> None:
-    async def scenario(collections: MongoLightningCollections) -> None:
-        counter_key = f"counter-{uuid4().hex}"
-        await _initialize_counter(collections, counter_key)
-        final_value = await _contention_run(
-            collections,
-            key=counter_key,
-            commit=True,
-            concurrency=6,
-        )
-        assert final_value == 6
-
-    await _with_mongo_collections(temporary_mongo_database, scenario)

tests/store/test_implementation.py

@@ -29,9 +29,7 @@ from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.memory import InMemoryLightningStore, estimate_model_size
 from agentlightning.types import (
     LLM,
-    Attempt,
     AttemptedRollout,
-    EnqueueRolloutRequest,
     Event,
     Link,
     OtelResource,
@@ -63,34 +61,6 @@ def test_paginated_result_behaves_like_sequence() -> None:
     assert repr(result2) == "<PaginatedResult (1: of 5) ['a', ...]>"
 
 
-@pytest.mark.asyncio
-async def test_statistics_updates_counts(store_fixture: LightningStore, mock_readable_span: Mock) -> None:
-    """Statistics should reflect entity counts across the store."""
-    initial = await store_fixture.statistics()
-
-    rollout = await store_fixture.start_rollout(input={"origin": "stats"})
-    await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span)
-    await store_fixture.add_resources(
-        {
-            "stat_llm": LLM(
-                resource_type="llm",
-                endpoint="http://localhost:8000/v1",
-                model="stats-model",
-            )
-        }
-    )
-    await store_fixture.update_worker("stats-worker", heartbeat_stats={"cpu": 0.5})
-
-    updated = await store_fixture.statistics()
-    assert updated["name"] == store_fixture.__class__.__name__  # type: ignore
-    assert updated["total_rollouts"] == initial.get("total_rollouts", 0) + 1  # type: ignore
-    assert updated["total_attempts"] == initial.get("total_attempts", 0) + 1  # type: ignore
-    assert updated["total_spans"] == initial.get("total_spans", 0) + 1  # type: ignore
-    assert updated["total_resources"] == initial.get("total_resources", 0) + 1  # type: ignore
-    assert updated["total_workers"] == initial.get("total_workers", 0) + 1  # type: ignore
-    assert updated["uptime"] > initial.get("uptime", 0.0)  # type: ignore
-
-
 # Core CRUD Operations Tests
 
 
@@ -676,49 +646,6 @@ async def test_requeue_mechanism(store_fixture: LightningStore) -> None:
     assert latest_attempt.sequence_id == 2
 
 
-@pytest.mark.asyncio
-async def test_enqueue_many_rollouts_preserves_order(store_fixture: LightningStore) -> None:
-    """enqueue_many_rollouts should enqueue tasks in the provided order with matching metadata."""
-
-    requests = [
-        EnqueueRolloutRequest(input={"idx": 0}, metadata={"batch": "a"}),
-        EnqueueRolloutRequest(input={"idx": 1}, mode="train"),
-        EnqueueRolloutRequest(input={"idx": 2}, config=RolloutConfig(timeout_seconds=3.5)),
-    ]
-
-    rollouts = await store_fixture.enqueue_many_rollouts(requests)
-
-    assert [rollout.input["idx"] for rollout in rollouts] == [0, 1, 2]
-    assert all(rollout.status == "queuing" for rollout in rollouts)
-    assert rollouts[0].metadata == {"batch": "a"}
-    assert rollouts[1].mode == "train"
-    assert rollouts[2].config.timeout_seconds == 3.5
-
-    for expected_idx in range(3):
-        dequeued = await store_fixture.dequeue_rollout()
-        assert dequeued is not None
-        assert dequeued.input["idx"] == expected_idx
-
-
-@pytest.mark.asyncio
-async def test_dequeue_many_rollouts_with_limit(store_fixture: LightningStore) -> None:
-    """dequeue_many_rollouts should honor limit and propagate worker IDs to attempts."""
-
-    requests = [EnqueueRolloutRequest(input={"idx": idx}) for idx in range(4)]
-    await store_fixture.enqueue_many_rollouts(requests)
-
-    first_batch = await store_fixture.dequeue_many_rollouts(limit=2, worker_id="bulk-worker")
-    assert len(first_batch) == 2
-    assert [attempt.input["idx"] for attempt in first_batch] == [0, 1]
-    assert all(attempt.attempt.worker_id == "bulk-worker" for attempt in first_batch)
-
-    second_batch = await store_fixture.dequeue_many_rollouts(limit=5, worker_id="bulk-worker")
-    assert len(second_batch) == 2
-    assert [attempt.input["idx"] for attempt in second_batch] == [2, 3]
-
-    assert await store_fixture.dequeue_many_rollouts(limit=1) == []
-
-
 @pytest.mark.asyncio
 async def test_update_and_query_workers(store_fixture: LightningStore) -> None:
     """Workers can be created, heartbeats recorded, and telemetry auto-updated."""
@@ -762,49 +689,6 @@ async def test_update_and_query_workers(store_fixture: LightningStore) -> None:
         await store_fixture.update_worker("worker-1", heartbeat_stats=None)  # type: ignore[arg-type]
 
 
-@pytest.mark.asyncio
-async def test_start_rollout_assigns_worker(store_fixture: LightningStore) -> None:
-    """start_rollout should immediately associate attempts with the provided worker."""
-    attempted = await store_fixture.start_rollout(input={"task": "direct"}, worker_id="worker-direct")
-
-    assert attempted.attempt.worker_id == "worker-direct"
-    worker = await store_fixture.get_worker_by_id("worker-direct")
-    assert worker is not None
-    assert worker.status == "busy"
-    assert worker.current_rollout_id == attempted.rollout_id
-    assert worker.current_attempt_id == attempted.attempt.attempt_id
-
-
-@pytest.mark.asyncio
-async def test_dequeue_rollout_assigns_worker(store_fixture: LightningStore) -> None:
-    """dequeue_rollout should stamp attempts and worker telemetry with worker_id."""
-    await store_fixture.enqueue_rollout(input={"task": "queued"})
-    dequeued = await store_fixture.dequeue_rollout(worker_id="worker-dequeue")
-
-    assert dequeued is not None
-    assert dequeued.attempt.worker_id == "worker-dequeue"
-    worker = await store_fixture.get_worker_by_id("worker-dequeue")
-    assert worker is not None
-    assert worker.status == "busy"
-    assert worker.current_rollout_id == dequeued.rollout_id
-    assert worker.current_attempt_id == dequeued.attempt.attempt_id
-
-
-@pytest.mark.asyncio
-async def test_start_attempt_assigns_worker(store_fixture: LightningStore) -> None:
-    """Manual retries should also update worker state when worker_id is provided."""
-    initial = await store_fixture.start_rollout(input={"task": "retry-seed"})
-    retry = await store_fixture.start_attempt(initial.rollout_id, worker_id="worker-retry")
-
-    assert retry.attempt.sequence_id == 2
-    assert retry.attempt.worker_id == "worker-retry"
-    worker = await store_fixture.get_worker_by_id("worker-retry")
-    assert worker is not None
-    assert worker.status == "busy"
-    assert worker.current_rollout_id == retry.rollout_id
-    assert worker.current_attempt_id == retry.attempt.attempt_id
-
-
 @pytest.mark.asyncio
 async def test_query_workers_supports_filters(store_fixture: LightningStore) -> None:
     """Worker queries should support filtering, sorting, and pagination."""
@@ -1214,116 +1098,20 @@ async def test_span_sequence_generation(store_fixture: LightningStore, mock_read
     assert seq_id == 1
 
     span1 = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-    assert span1 is not None and span1.sequence_id == 2
+    assert span1.sequence_id == 2
 
     # Next span gets sequence_id 3
     seq_id = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, attempt_id)
     assert seq_id == 3
 
     span2 = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-    assert span2 is not None and span2.sequence_id == 4
+    assert span2.sequence_id == 4
 
     # Different attempt reuses the same rollout_id
     seq_id = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, "attempt-does-not-exist")
     assert seq_id == 5
 
 
-@pytest.mark.asyncio
-async def test_get_many_span_sequence_ids_handles_mixed_batches(store_fixture: LightningStore) -> None:
-    """Bulk sequence allocation should work across mixed rollouts and attempts."""
-    first = await store_fixture.start_rollout(input={"origin": "bulk-seq"})
-    second = await store_fixture.start_rollout(input={"origin": "bulk-seq-2"})
-    await store_fixture.update_rollout(first.rollout_id, status="requeuing")
-    retried = await store_fixture.start_attempt(first.rollout_id)
-
-    sequence_pairs = [
-        (first.rollout_id, first.attempt.attempt_id),
-        (second.rollout_id, second.attempt.attempt_id),
-        (first.rollout_id, retried.attempt.attempt_id),
-        (first.rollout_id, first.attempt.attempt_id),
-    ]
-    sequence_ids = await store_fixture.get_many_span_sequence_ids(sequence_pairs)
-    assert sequence_ids == [1, 1, 2, 3]
-
-    next_first = await store_fixture.get_next_span_sequence_id(first.rollout_id, first.attempt.attempt_id)
-    next_second = await store_fixture.get_next_span_sequence_id(second.rollout_id, second.attempt.attempt_id)
-    assert next_first == 4
-    assert next_second == 2
-
-
-@pytest.mark.asyncio
-async def test_add_many_spans_handles_mixed_rollouts_and_attempts(store_fixture: LightningStore) -> None:
-    """Batch span insertion should handle mixed rollouts and skip duplicates."""
-    first = await store_fixture.start_rollout(input={"origin": "batch-spans"})
-    second = await store_fixture.start_rollout(input={"origin": "batch-spans-2"})
-    await store_fixture.update_rollout(first.rollout_id, status="requeuing")
-    retried = await store_fixture.start_attempt(first.rollout_id)
-
-    def _build_span(seed: int, rollout_id: str, attempt_id: str) -> Span:
-        trace_hex = f"{seed:032x}"
-        span_hex = f"{seed:016x}"
-        return Span(
-            rollout_id=rollout_id,
-            attempt_id=attempt_id,
-            sequence_id=seed,
-            trace_id=trace_hex,
-            span_id=span_hex,
-            parent_id=None,
-            name=f"batch-span-{seed}",
-            status=TraceStatus(status_code="OK"),
-            attributes={},
-            events=[],
-            links=[],
-            start_time=None,
-            end_time=None,
-            context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}),
-            parent=None,
-            resource=OtelResource(attributes={}, schema_url=""),
-        )
-
-    span_first = _build_span(1, first.rollout_id, first.attempt.attempt_id)
-    span_retry = _build_span(2, retried.rollout_id, retried.attempt.attempt_id)
-    span_second = _build_span(3, second.rollout_id, second.attempt.attempt_id)
-    duplicate_first = _build_span(1, first.rollout_id, first.attempt.attempt_id)
-
-    stored = await store_fixture.add_many_spans([span_first, span_retry, span_second, duplicate_first])
-    assert {span.span_id for span in stored} == {span_first.span_id, span_retry.span_id, span_second.span_id}
-
-    spans_first = await store_fixture.query_spans(first.rollout_id)
-    assert {span.span_id for span in spans_first} >= {span_first.span_id, span_retry.span_id}
-    spans_second = await store_fixture.query_spans(second.rollout_id)
-    assert {span.span_id for span in spans_second} >= {span_second.span_id}
-
-
-@pytest.mark.asyncio
-async def test_add_span_returns_none_on_duplicate(store_fixture: LightningStore) -> None:
-    """Duplicate span IDs should return None instead of raising."""
-    attempted = await store_fixture.start_rollout(input={"origin": "duplicate-span"})
-    attempt_id = attempted.attempt.attempt_id
-
-    span = Span(
-        rollout_id=attempted.rollout_id,
-        attempt_id=attempt_id,
-        sequence_id=1,
-        trace_id="0" * 32,
-        span_id="0" * 16,
-        parent_id=None,
-        name="dup-span",
-        status=TraceStatus(status_code="OK"),
-        attributes={},
-        events=[],
-        links=[],
-        start_time=None,
-        end_time=None,
-        context=SpanContext(trace_id="0" * 32, span_id="0" * 16, is_remote=False, trace_state={}),
-        parent=None,
-        resource=OtelResource(attributes={}, schema_url=""),
-    )
-
-    assert await store_fixture.add_span(span) is span
-    assert await store_fixture.add_span(span) is None
-
-
 @pytest.mark.asyncio
 async def test_span_updates_attempt_status(store_fixture: LightningStore, mock_readable_span: Mock) -> None:
     """Adding a span should persist attempt heartbeat and transition status to running."""
@@ -1343,35 +1131,6 @@ async def test_span_updates_attempt_status(store_fixture: LightningStore, mock_r
     assert updated_attempt.last_heartbeat_time is not None
 
 
-@pytest.mark.asyncio
-async def test_spans_promote_preparing_attempt_with_heartbeat(
-    store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch
-) -> None:
-    """Spans should set heartbeat time and promote preparing attempts/rollouts to running."""
-    rollout = await store_fixture.enqueue_rollout(input={"test": "preparing-heartbeat"})
-    dequeued = await store_fixture.dequeue_rollout()
-    assert dequeued is not None
-
-    attempts_before = await store_fixture.query_attempts(rollout.rollout_id)
-    assert attempts_before
-    attempt_id = attempts_before[0].attempt_id
-    assert attempts_before[0].status == "preparing"
-    assert attempts_before[0].last_heartbeat_time is None
-
-    heartbeat_time = 1234.5
-    monkeypatch.setattr("agentlightning.store.collection_based.time.time", lambda: heartbeat_time)
-
-    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-
-    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
-    assert attempt_after.status == "running"
-    assert attempt_after.last_heartbeat_time == heartbeat_time
-
-    rollout_after = await store_fixture.get_rollout_by_id(rollout.rollout_id)
-    assert rollout_after is not None
-    assert rollout_after.status == "running"
-
-
 @pytest.mark.asyncio
 async def test_unresponsive_attempt_recovers_after_span(
     store_fixture: LightningStore, mock_readable_span: Mock
@@ -1420,77 +1179,6 @@ async def test_running_attempt_updates_heartbeat(
     assert attempt_after_second.last_heartbeat_time == first_heartbeat + 100.0
 
 
-@pytest.mark.asyncio
-async def test_span_post_add_preserves_concurrent_updates(
-    store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch
-) -> None:
-    """Concurrent attempt updates should not be clobbered when spans record heartbeats."""
-    rollout = await store_fixture.enqueue_rollout(input={"test": "span-concurrency"})
-    dequeued = await store_fixture.dequeue_rollout()
-    assert dequeued is not None
-    attempt_id = dequeued.attempt.attempt_id
-
-    original_post = store_fixture._post_add_spans  # type: ignore
-
-    async def patched_post(spans: List[Span], rollout_id: str, mutated_attempt_id: str) -> None:
-        await store_fixture.update_attempt(rollout_id, mutated_attempt_id, metadata={"concurrent": True})
-        await original_post(spans, rollout_id, mutated_attempt_id)
-
-    monkeypatch.setattr(store_fixture, "_post_add_spans", patched_post)
-
-    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-
-    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
-    assert attempt_after.metadata == {"concurrent": True}
-
-
-@pytest.mark.asyncio
-async def test_healthcheck_marks_unresponsive_and_updates_worker(
-    store_fixture: LightningStore, monkeypatch: pytest.MonkeyPatch
-) -> None:
-    """Healthcheck should mark attempts unresponsive and sync worker state via the helper."""
-
-    class TimeStub:
-        def __init__(self, value: float):
-            self.value = value
-
-        def __call__(self) -> float:
-            return self.value
-
-    time_stub = TimeStub(100.0)
-    monkeypatch.setattr("agentlightning.store.collection_based.time.time", time_stub)
-
-    rollout = await store_fixture.enqueue_rollout(
-        input={"test": "healthcheck-worker"},
-        config=RolloutConfig(unresponsive_seconds=1.0),
-    )
-    dequeued = await store_fixture.dequeue_rollout(worker_id="worker-sync")
-    assert dequeued is not None
-    attempt_id = dequeued.attempt.attempt_id
-    await store_fixture.update_attempt(rollout.rollout_id, attempt_id, worker_id="worker-sync")
-
-    original_sync = store_fixture._sync_workers_with_attempts  # type: ignore
-    sync_calls: List[str] = []
-
-    async def tracking_sync(attempts: Sequence[Attempt]) -> None:
-        for attempt in attempts:
-            sync_calls.append(attempt.attempt_id)
-        await original_sync(attempts)
-
-    monkeypatch.setattr(store_fixture, "_sync_workers_with_attempts", tracking_sync)
-
-    time_stub.value = 105.0
-    await store_fixture.get_rollout_by_id(rollout.rollout_id)
-
-    worker = await store_fixture.get_worker_by_id("worker-sync")
-    assert worker is not None
-    assert worker.status == "unknown"
-
-    attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0]
-    assert attempt_after.status == "unresponsive"
-    assert sync_calls == [attempt_after.attempt_id]
-
-
 @pytest.mark.asyncio
 async def test_duplicate_span_id_error(
     store_fixture: LightningStore, mock_readable_span: Mock, caplog: pytest.LogCaptureFixture
@@ -1508,8 +1196,7 @@ async def test_duplicate_span_id_error(
 
     await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
 
-    duplicate = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
-    assert duplicate is None
+    await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
     assert "Duplicated span added" in caplog.text
 
 
@@ -1524,7 +1211,7 @@ async def test_span_with_explicit_sequence_id(store_fixture: LightningStore, moc
 
     # Add span with explicit sequence_id
     span = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span, sequence_id=100)
-    assert span is not None and span.sequence_id == 100
+    assert span.sequence_id == 100
 
     next_seq = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, attempt_id)
     assert next_seq == 101
@@ -2466,7 +2153,7 @@ async def test_concurrent_span_additions(store_fixture: LightningStore, mock_rea
     assert rollout is not None
 
     async def add_span(index: int) -> Span:
-        return await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span)  # type: ignore
+        return await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span)
 
     # Add 30 spans concurrently
     tasks = [add_span(i) for i in range(30)]
@@ -2517,7 +2204,7 @@ async def test_concurrent_resource_updates(store_fixture: LightningStore) -> Non
 @pytest.mark.asyncio
 async def test_update_nonexistent_rollout(store_fixture: LightningStore) -> None:
     """Test updating non-existent rollout raises error."""
-    with pytest.raises(ValueError, match=r"Item.*does not exist"):
+    with pytest.raises(ValueError, match="Rollout nonexistent not found"):
         await store_fixture.update_rollout(rollout_id="nonexistent", status="failed")
 
 
@@ -3067,7 +2754,7 @@ async def test_full_lifecycle_success(store_fixture: LightningStore, mock_readab
 
     # 3. Add span (transitions to running)
     span = await store_fixture.add_otel_span(rollout.rollout_id, attempt.attempt_id, mock_readable_span)
-    assert span is not None and span.sequence_id == 1
+    assert span.sequence_id == 1
 
     # Check status transitions
     rollouts = await store_fixture.query_rollouts(status=["running"])

tests/store/test_restful.py

@@ -12,7 +12,7 @@ Test categories:
 
 import asyncio
 import contextlib
-from typing import Any, AsyncGenerator, Dict, List, Tuple
+from typing import AsyncGenerator, List, Tuple
 
 import aiohttp
 import pytest
@@ -69,77 +69,6 @@ async def _run_server_with_cors(cors_origins: List[str] | str | None = None):
         await server.stop()
 
 
-async def _request_rollouts_page(
-    session: aiohttp.ClientSession, api_endpoint: str, method: str, payload: Dict[str, Any]
-) -> Dict[str, Any]:
-    """Send either GET or POST /search request for rollouts and return parsed payload."""
-
-    if method == "get":
-        async with session.get(f"{api_endpoint}/rollouts", params=payload) as resp:
-            assert resp.status == 200
-            return await resp.json()
-    async with session.post(f"{api_endpoint}/rollouts/search", json=payload) as resp:
-        assert resp.status == 200
-        return await resp.json()
-
-
-async def _request_attempts_page(
-    session: aiohttp.ClientSession,
-    api_endpoint: str,
-    rollout_id: str,
-    method: str,
-    payload: Dict[str, Any],
-) -> Dict[str, Any]:
-    """Send either GET or POST /search request for attempts and return parsed payload."""
-
-    base = f"{api_endpoint}/rollouts/{rollout_id}/attempts"
-    if method == "get":
-        async with session.get(base, params=payload) as resp:
-            assert resp.status == 200
-            return await resp.json()
-    async with session.post(f"{base}/search", json=payload) as resp:
-        assert resp.status == 200
-        return await resp.json()
-
-
-async def _request_spans_page(
-    session: aiohttp.ClientSession,
-    api_endpoint: str,
-    rollout_id: str,
-    method: str,
-    payload: Dict[str, Any],
-) -> Dict[str, Any]:
-    """Send either GET or POST /search request for spans and return parsed payload."""
-
-    base_params = {"rollout_id": rollout_id}
-    base_params.update(payload)
-    if method == "get":
-        async with session.get(f"{api_endpoint}/spans", params=base_params) as resp:
-            assert resp.status == 200
-            return await resp.json()
-    async with session.post(f"{api_endpoint}/spans/search", json=base_params) as resp:
-        assert resp.status == 200
-        return await resp.json()
-
-
-async def _request_workers_page(
-    session: aiohttp.ClientSession,
-    api_endpoint: str,
-    method: str,
-    payload: Dict[str, Any],
-) -> Dict[str, Any]:
-    """Send either GET or POST /search request for workers and return parsed payload."""
-
-    base = f"{api_endpoint}/workers"
-    if method == "get":
-        async with session.get(base, params=payload) as resp:
-            assert resp.status == 200
-            return await resp.json()
-    async with session.post(f"{base}/search", json=payload) as resp:
-        assert resp.status == 200
-        return await resp.json()
-
-
 @pytest_asyncio.fixture
 async def server_client(
     store_fixture: LightningStore,
@@ -208,152 +137,6 @@ async def test_cors_allows_wildcard_origin() -> None:
                 assert allow_credentials == "true"
 
 
-@pytest.mark.asyncio
-async def test_enqueue_endpoint_batches_payloads(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    _server, _client, session, api_endpoint = server_client
-
-    single_payload = {"rollouts": [{"input": {"task": "single"}}]}
-    async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json=single_payload) as resp:
-        assert resp.status == 201
-        body = await resp.json()
-        assert isinstance(body, list)
-        assert len(body) == 1  # type: ignore
-        assert body[0]["input"] == {"task": "single"}
-
-    batch_payload = {
-        "rollouts": [
-            {"input": {"task": "batch-1"}, "metadata": {"batch": 1}},
-            {"input": {"task": "batch-2"}},
-        ]
-    }
-    async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json=batch_payload) as resp:
-        assert resp.status == 201
-        body = await resp.json()
-        assert len(body) == 2
-        assert [item["input"]["task"] for item in body] == ["batch-1", "batch-2"]
-        assert body[0]["metadata"] == {"batch": 1}
-
-
-@pytest.mark.asyncio
-async def test_dequeue_endpoint_returns_batches(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    server, _client, session, api_endpoint = server_client
-
-    for idx in range(3):
-        await server.enqueue_rollout(input={"idx": idx})
-
-    async with session.post(
-        f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 2, "worker_id": "rest-worker"}
-    ) as resp:
-        assert resp.status == 200
-        body = await resp.json()
-        assert len(body) == 2
-        assert all(item["attempt"]["worker_id"] == "rest-worker" for item in body)
-
-    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue") as resp:
-        assert resp.status == 200
-        body = await resp.json()
-        assert len(body) == 1
-
-    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue") as resp:
-        assert resp.status == 200
-        body = await resp.json()
-        assert body == []
-
-
-@pytest.mark.asyncio
-async def test_enqueue_endpoint_requires_rollouts_field(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    _server, _client, session, api_endpoint = server_client
-
-    async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json={}) as resp:
-        assert resp.status == 422
-
-
-@pytest.mark.asyncio
-async def test_dequeue_endpoint_zero_limit_returns_empty(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    server, _client, session, api_endpoint = server_client
-    await server.enqueue_rollout(input={"idx": 0})
-
-    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 0}) as resp:
-        assert resp.status == 200
-        assert await resp.json() == []
-
-    async with session.post(f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 1}) as resp:
-        assert resp.status == 200
-        body = await resp.json()
-        assert len(body) == 1
-
-
-@pytest.mark.asyncio
-async def test_statistics_endpoint_returns_counts(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    """The statistics endpoint should expose store-level metrics."""
-    server, client, session, api_endpoint = server_client
-    await client.start_rollout(input={"source": "stats-endpoint"})
-
-    async with session.get(f"{api_endpoint}/statistics") as resp:
-        assert resp.status == 200
-        payload = await resp.json()
-
-    expected_name = server.store.__class__.__name__ if server.store is not None else payload["name"]
-    assert payload["name"] == expected_name
-    assert payload["total_rollouts"] >= 1
-
-
-@pytest.mark.asyncio
-async def test_rest_start_rollout_propagates_worker_id(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    server, _client, session, api_endpoint = server_client
-    payload = {"input": {"source": "rest-worker"}, "worker_id": "rest-start-worker"}
-
-    async with session.post(f"{api_endpoint}/rollouts", json=payload) as resp:
-        assert resp.status == 201
-        data = await resp.json()
-
-    assert data["attempt"]["worker_id"] == "rest-start-worker"
-    worker = await server.get_worker_by_id("rest-start-worker")
-    assert worker is not None
-    assert worker.status == "busy"
-    assert worker.current_rollout_id == data["rollout_id"]
-    assert worker.current_attempt_id == data["attempt"]["attempt_id"]
-
-
-@pytest.mark.asyncio
-async def test_rest_start_attempt_propagates_worker_id(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    server, _client, session, api_endpoint = server_client
-
-    async with session.post(f"{api_endpoint}/rollouts", json={"input": {"source": "rest-retry"}}) as resp:
-        assert resp.status == 201
-        base_rollout = await resp.json()
-
-    attempt_worker = "rest-attempt-worker"
-    async with session.post(
-        f"{api_endpoint}/rollouts/{base_rollout['rollout_id']}/attempts",
-        json={"worker_id": attempt_worker},
-    ) as resp:
-        assert resp.status == 201
-        retry_payload = await resp.json()
-
-    assert retry_payload["attempt"]["sequence_id"] == 2
-    assert retry_payload["attempt"]["worker_id"] == attempt_worker
-    worker = await server.get_worker_by_id(attempt_worker)
-    assert worker is not None
-    assert worker.status == "busy"
-    assert worker.current_rollout_id == retry_payload["rollout_id"]
-    assert worker.current_attempt_id == retry_payload["attempt"]["attempt_id"]
-
-
 # Rollouts Pagination, Sorting, and Filtering Tests
 
 
@@ -407,24 +190,6 @@ async def test_rollouts_pagination_disabled(
         assert len(data["items"]) == 15
 
 
-@pytest.mark.asyncio
-@pytest.mark.parametrize("method", ["get", "post"])
-async def test_rollouts_search_supports_get_and_post(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
-) -> None:
-    """Ensure both GET and POST /search endpoints behave the same for rollouts."""
-
-    server, _client, session, api_endpoint = server_client
-    for i in range(3):
-        await server.enqueue_rollout(input={"index": i})
-
-    data = await _request_rollouts_page(session, api_endpoint, method, {"limit": 2, "offset": 0})
-    assert data["total"] == 3
-    assert data["limit"] == 2
-    assert data["offset"] == 0
-    assert len(data["items"]) == 2
-
-
 @pytest.mark.asyncio
 async def test_rollouts_sorting_by_start_time(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -699,22 +464,6 @@ async def test_attempts_pagination_basic(
         assert len(data["items"]) == 5
 
 
-@pytest.mark.asyncio
-@pytest.mark.parametrize("method", ["get", "post"])
-async def test_attempts_search_supports_get_and_post(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
-) -> None:
-    server, _client, session, api_endpoint = server_client
-
-    rollout = await server.enqueue_rollout(input={"test": "attempt-search"})
-    await server.start_attempt(rollout.rollout_id)
-
-    data = await _request_attempts_page(session, api_endpoint, rollout.rollout_id, method, {"limit": 1, "offset": 0})
-    assert data["total"] == 1
-    assert data["limit"] == 1
-    assert len(data["items"]) == 1
-
-
 @pytest.mark.asyncio
 async def test_attempts_sorting(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -980,30 +729,6 @@ async def test_spans_pagination_basic(
         assert len(data["items"]) == 10
 
 
-@pytest.mark.asyncio
-@pytest.mark.parametrize("method", ["get", "post"])
-async def test_spans_search_supports_get_and_post(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
-) -> None:
-    server, _client, session, api_endpoint = server_client
-
-    attempted = await server.start_rollout(input={"test": "span-search"})
-    attempt_id = attempted.attempt.attempt_id
-    for seq in range(3):
-        await server.add_span(_make_span(attempted.rollout_id, attempt_id, seq + 1, f"span-{seq}"))
-
-    data = await _request_spans_page(
-        session,
-        api_endpoint,
-        attempted.rollout_id,
-        method,
-        {"attempt_id": attempt_id, "limit": 2, "offset": 0},
-    )
-    assert data["total"] == 3
-    assert data["limit"] == 2
-    assert len(data["items"]) == 2
-
-
 @pytest.mark.asyncio
 async def test_spans_sorting_by_start_time(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
@@ -1374,50 +1099,6 @@ async def test_request_json_spans_returns_pagination_metadata(
     assert len(data["items"]) == 1
 
 
-# Update semantics tests
-
-
-@pytest.mark.asyncio
-async def test_update_rollout_distinguishes_unset_fields(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    server, _client, session, api_endpoint = server_client
-    rollout = await server.enqueue_rollout(input={"id": "unset-check"})
-
-    async with session.post(
-        f"{api_endpoint}/rollouts/{rollout.rollout_id}",
-        json={"metadata": {"foo": "bar"}},
-    ) as resp:
-        assert resp.status == 200
-
-    async with session.post(
-        f"{api_endpoint}/rollouts/{rollout.rollout_id}",
-        json={"status": None},
-    ) as resp:
-        assert resp.status == 400
-
-
-@pytest.mark.asyncio
-async def test_update_attempt_distinguishes_unset_fields(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
-) -> None:
-    server, _client, session, api_endpoint = server_client
-    attempted = await server.start_rollout(input={"id": "attempt-unset"})
-    attempt_id = attempted.attempt.attempt_id
-
-    async with session.post(
-        f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}",
-        json={"status": "running"},
-    ) as resp:
-        assert resp.status == 200
-
-    async with session.post(
-        f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}",
-        json={"worker_id": None},
-    ) as resp:
-        assert resp.status == 400
-
-
 # Client Compatibility Tests
 
 
@@ -1464,27 +1145,6 @@ async def test_client_query_with_filters(
     assert rollouts[0].rollout_id == r2.rollout_id
 
 
-@pytest.mark.asyncio
-@pytest.mark.parametrize("method", ["get", "post"])
-async def test_workers_search_supports_get_and_post(
-    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str
-) -> None:
-    server, _client, session, api_endpoint = server_client
-
-    await server.update_worker("worker-1", heartbeat_stats={"cpu": 0.5})
-    await server.update_worker("worker-2", heartbeat_stats={"cpu": 0.7})
-
-    data = await _request_workers_page(
-        session,
-        api_endpoint,
-        method,
-        {"limit": 1, "offset": 0, "sort_by": "worker_id", "sort_order": "asc"},
-    )
-    assert data["total"] == 2
-    assert data["limit"] == 1
-    assert len(data["items"]) == 1
-
-
 @pytest.mark.asyncio
 async def test_workers_endpoint_supports_updates(
     server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],

tests/store/test_threading.py

@@ -17,7 +17,6 @@ from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
-    EnqueueRolloutRequest,
     NamedResources,
     OtelResource,
     ResourcesUpdate,
@@ -178,11 +177,9 @@ async def test_threaded_store_delegates_all_methods() -> None:
         "get_resources_by_id": resources_update,
         "get_latest_resources": resources_update,
         "add_span": span,
-        "add_many_spans": [span],
         "add_otel_span": span,
         "wait_for_rollouts": [base_rollout],
         "get_next_span_sequence_id": 42,
-        "get_many_span_sequence_ids": [101],
         "query_spans": [span],
         "update_rollout": updated_rollout,
         "update_attempt": updated_attempt,
@@ -214,11 +211,9 @@ async def test_threaded_store_delegates_all_methods() -> None:
     assert await threaded_store.get_resources_by_id("resources-1") == resources_update
     assert await threaded_store.get_latest_resources() == resources_update
     assert await threaded_store.add_span(span) == span
-    assert await threaded_store.add_many_spans([span]) == [span]
     assert await threaded_store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id=5) == span
     assert await threaded_store.wait_for_rollouts(rollout_ids=[rollout_id], timeout=1.0) == [base_rollout]
     assert await threaded_store.get_next_span_sequence_id(rollout_id, attempt_id) == 42
-    assert await threaded_store.get_many_span_sequence_ids([(rollout_id, attempt_id)]) == [101]
     assert await threaded_store.query_spans(rollout_id, attempt_id="latest") == [span]
     assert (
         await threaded_store.update_rollout(
@@ -259,11 +254,9 @@ async def test_threaded_store_delegates_all_methods() -> None:
         "get_resources_by_id",
         "get_latest_resources",
         "add_span",
-        "add_many_spans",
         "add_otel_span",
         "wait_for_rollouts",
         "get_next_span_sequence_id",
-        "get_many_span_sequence_ids",
         "query_spans",
         "update_rollout",
         "update_attempt",
@@ -274,42 +267,6 @@ async def test_threaded_store_delegates_all_methods() -> None:
     assert [name for name, *_ in dummy_store.calls] == expected_order
 
 
-@pytest.mark.asyncio
-async def test_threaded_store_enqueue_many_rollouts_delegates() -> None:
-    requests = [
-        EnqueueRolloutRequest(input={"idx": 0}, mode="train", metadata={"batch": "left"}),
-        EnqueueRolloutRequest(input={"idx": 1}, mode=None, resources_id="resources-1"),
-    ]
-    rollouts = [
-        Rollout(rollout_id="bulk-0", input={"idx": 0}, start_time=0.0),
-        Rollout(rollout_id="bulk-1", input={"idx": 1}, start_time=1.0),
-    ]
-    dummy_store = DummyLightningStore({"enqueue_many_rollouts": rollouts})
-    threaded_store = LightningStoreThreaded(dummy_store)
-
-    result = await threaded_store.enqueue_many_rollouts(requests)
-    assert result == rollouts
-    assert dummy_store.calls[-1][0] == "enqueue_many_rollouts"
-    assert dummy_store.calls[-1][1][0] == requests
-
-
-@pytest.mark.asyncio
-async def test_threaded_store_dequeue_many_rollouts_delegates() -> None:
-    attempt_a = Attempt(rollout_id="bulk-0", attempt_id="attempt-0", sequence_id=1, start_time=0.0)
-    attempt_b = Attempt(rollout_id="bulk-1", attempt_id="attempt-1", sequence_id=1, start_time=0.0)
-    attempts = [
-        AttemptedRollout(rollout_id="bulk-0", input={"idx": 0}, start_time=0.0, attempt=attempt_a),
-        AttemptedRollout(rollout_id="bulk-1", input={"idx": 1}, start_time=0.0, attempt=attempt_b),
-    ]
-    dummy_store = DummyLightningStore({"dequeue_many_rollouts": attempts})
-    threaded_store = LightningStoreThreaded(dummy_store)
-
-    result = await threaded_store.dequeue_many_rollouts(limit=2, worker_id="thread-worker")
-    assert result == attempts
-    assert dummy_store.calls[-1][0] == "dequeue_many_rollouts"
-    assert dummy_store.calls[-1][2] == {"limit": 2, "worker_id": "thread-worker"}
-
-
 def test_threaded_store_serializes_update_attempt_calls() -> None:
     store = SlowAttemptStore()
     threaded_store = LightningStoreThreaded(store)
@@ -402,28 +359,3 @@ def test_threaded_store_serializes_add_resources_calls() -> None:
     assert len(updates) == num_adds
     resource_ids = {update.resources_id for update in updates}
     assert len(resource_ids) == num_adds  # All IDs should be unique
-
-
-@pytest.mark.asyncio
-async def test_threaded_store_statistics_match_underlying(inmemory_store: InMemoryLightningStore) -> None:
-    """Threaded store should proxy statistics from wrapped store."""
-    rollout = await inmemory_store.start_rollout(input={"source": "threaded-stats"})
-    await inmemory_store.add_span(make_span(rollout.rollout_id, rollout.attempt.attempt_id))
-    await inmemory_store.add_resources(
-        {
-            "threaded_llm": LLM(
-                resource_type="llm",
-                endpoint="http://localhost:9000/v1",
-                model="threaded-model",
-            )
-        }
-    )
-    await inmemory_store.update_worker("threaded-worker", heartbeat_stats={"cpu": 0.1})
-
-    threaded_store = LightningStoreThreaded(inmemory_store)
-    threaded_stats = await threaded_store.statistics()
-    inmemory_stats = await inmemory_store.statistics()
-    assert {k: v for k, v in threaded_stats.items() if k != "uptime"} == {
-        k: v for k, v in inmemory_stats.items() if k != "uptime"
-    }
-    assert threaded_stats["uptime"] < inmemory_stats["uptime"]  # type: ignore