Implement AgentRunnerV2 (#125)

.github/workflows/tests-full.yml

@@ -46,3 +46,5 @@ jobs:
           pytest -v tests
         env:
           PYTEST_ADDOPTS: "--color=yes"
+          OPENAI_BASE_URL: ${{ secrets.OPENAI_API_BASE }}  # We will use BASE_URL instead of API_BASE in pytest
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}

agentlightning/adapter/triplet.py

@@ -11,7 +11,7 @@ from opentelemetry import trace as trace_api
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel
 
-from agentlightning.types import Triplet
+from agentlightning.types import SpanNames, Triplet
 
 from .base import TraceAdapter
 
@@ -289,6 +289,10 @@ class TraceTree:
                         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:

agentlightning/litagent.py

@@ -8,10 +8,10 @@ import logging
 import weakref
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, Generic, Optional, TypeVar, Union
 
-from .types import LLM, NamedResources, Rollout, RolloutRawResult, Task
+from .types import LLM, NamedResources, Rollout, RolloutRawResultV2, RolloutV2, Task
 
 if TYPE_CHECKING:
-    from .runner import AgentRunner
+    from .runner import BaseRunner
     from .tracer import BaseTracer
     from .trainer import Trainer
 
@@ -57,7 +57,7 @@ class LitAgent(Generic[T]):
         """
         self.trained_agents = trained_agents
         self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
-        self._runner_ref: weakref.ReferenceType[AgentRunner] | None = None
+        self._runner_ref: weakref.ReferenceType[BaseRunner[T]] | None = None
 
     @property
     def is_async(self) -> bool:
@@ -114,22 +114,22 @@ class LitAgent(Generic[T]):
         """
         return self.trainer.tracer
 
-    def set_runner(self, runner: AgentRunner) -> None:
+    def set_runner(self, runner: BaseRunner[T]) -> None:
         """
         Set the runner for this agent.
 
         Args:
-            runner: The AgentRunner instance that will handle the execution of rollouts.
+            runner: The runner instance that will handle the execution of rollouts.
         """
         self._runner_ref = weakref.ref(runner)
 
     @property
-    def runner(self) -> AgentRunner:
+    def runner(self) -> BaseRunner[T]:
         """
         Get the runner for this agent.
 
         Returns:
-            The AgentRunner instance associated with this agent.
+            The runner instance associated with this agent.
         """
         if self._runner_ref is None:
             raise ValueError("Runner has not been set for this agent.")
@@ -138,12 +138,14 @@ class LitAgent(Generic[T]):
             raise ValueError("Runner reference is no longer valid (object has been garbage collected).")
         return runner
 
-    def on_rollout_start(self, task: Task, runner: AgentRunner, tracer: BaseTracer) -> None:
+    def on_rollout_start(self, task: Task, runner: BaseRunner[T], tracer: BaseTracer) -> None:
         """Hook called immediately before a rollout begins.
 
+        Deprecated in favor of `on_rollout_start` in the `Hook` interface.
+
         Args:
             task: The :class:`Task` object that will be processed.
-            runner: The :class:`AgentRunner` managing the rollout.
+            runner: The :class:`BaseRunner` managing the rollout.
             tracer: The tracer instance associated with the runner.
 
         Subclasses can override this method to implement custom logic such as
@@ -151,20 +153,22 @@ class LitAgent(Generic[T]):
         no-op.
         """
 
-    def on_rollout_end(self, task: Task, rollout: Rollout, runner: AgentRunner, tracer: BaseTracer) -> None:
+    def on_rollout_end(self, task: Task, rollout: RolloutV2, runner: BaseRunner[T], tracer: BaseTracer) -> None:
         """Hook called after a rollout completes.
 
+        Deprecated in favor of `on_rollout_end` in the `Hook` interface.
+
         Args:
             task: The :class:`Task` object that was processed.
             rollout: The resulting :class:`Rollout` object.
-            runner: The :class:`AgentRunner` managing the rollout.
+            runner: The :class:`BaseRunner` managing the rollout.
             tracer: The tracer instance associated with the runner.
 
         Subclasses can override this method for cleanup or additional
         logging. By default, this is a no-op.
         """
 
-    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    def rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
         """Main entry point for executing a rollout.
 
         This method determines whether to call the synchronous or
@@ -193,7 +197,7 @@ class LitAgent(Generic[T]):
         """
         raise NotImplementedError("Agents must implement the `rollout` method.")
 
-    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
         """Asynchronous version of the main rollout method.
 
         This method determines whether to call the synchronous or
@@ -219,7 +223,7 @@ class LitAgent(Generic[T]):
         """
         raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")
 
-    def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    def training_rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
         """Defines the agent's behavior for a single training task.
 
         This method should contain the logic for how the agent processes an
@@ -235,7 +239,7 @@ class LitAgent(Generic[T]):
         """
         return self.rollout(task, resources, rollout)
 
-    def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    def validation_rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
         """Defines the agent's behavior for a single validation task.
 
         By default, this method redirects to `training_rollout`. Override it
@@ -253,7 +257,9 @@ class LitAgent(Generic[T]):
         """
         return self.rollout(task, resources, rollout)
 
-    async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    async def training_rollout_async(
+        self, task: T, resources: NamedResources, rollout: RolloutV2
+    ) -> RolloutRawResultV2:
         """Asynchronous version of `training_rollout`.
 
         This method should be implemented by agents that perform asynchronous
@@ -270,7 +276,9 @@ class LitAgent(Generic[T]):
         """
         return await self.rollout_async(task, resources, rollout)
 
-    async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    async def validation_rollout_async(
+        self, task: T, resources: NamedResources, rollout: RolloutV2
+    ) -> RolloutRawResultV2:
         """Asynchronous version of `validation_rollout`.
 
         By default, this method redirects to `training_rollout_async`.
@@ -289,10 +297,10 @@ class LitAgent(Generic[T]):
 
 
 LlmRolloutFunc = Union[
-    Callable[[T, LLM, Rollout], RolloutRawResult],
-    Callable[[T, LLM], RolloutRawResult],
-    Callable[[T, LLM, Rollout], Coroutine[Any, Any, RolloutRawResult]],
-    Callable[[T, LLM], Coroutine[Any, Any, RolloutRawResult]],
+    Callable[[T, LLM, Rollout], RolloutRawResultV2],
+    Callable[[T, LLM], RolloutRawResultV2],
+    Callable[[T, LLM, Rollout], Coroutine[Any, Any, RolloutRawResultV2]],
+    Callable[[T, LLM], Coroutine[Any, Any, RolloutRawResultV2]],
 ]
 
 
@@ -331,7 +339,7 @@ class LitAgentLLM(LitAgent[T]):
     def is_async(self) -> bool:
         return self._is_async
 
-    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    def rollout(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
         """Execute a synchronous rollout using the wrapped function.
 
         Args:
@@ -353,7 +361,7 @@ class LitAgentLLM(LitAgent[T]):
         else:
             return self.llm_rollout_func(task, llm=llm)  # type: ignore
 
-    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: RolloutV2) -> RolloutRawResultV2:
         """Execute an asynchronous rollout using the wrapped function.
 
         Args:

agentlightning/llm_proxy.py

@@ -14,6 +14,7 @@ import time
 from typing import Any, Awaitable, Callable, Dict, Iterable, List, Optional, Sequence, TypedDict, Union, cast
 
 import litellm
+import opentelemetry.trace as trace_api
 import uvicorn
 import yaml
 from fastapi import Request, Response
@@ -23,7 +24,7 @@ from litellm.proxy.proxy_server import app, save_worker_config  # pyright: ignor
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 
-from agentlightning.types import LLM
+from agentlightning.types import LLM, ProxyLLM
 
 from .store.base import LightningStore
 
@@ -208,13 +209,35 @@ class LightningSpanExporter(SpanExporter):
     def __init__(self, store: Optional[LightningStore] = None):
         self._store = store
         self._buffer: List[ReadableSpan] = []
-        self._lock = threading.RLock()
+        self._lock: Optional[threading.RLock] = None
 
         # Single dedicated event loop running in a daemon thread.
         # This decouples OTEL SDK threads from our async store I/O.
-        self._loop = asyncio.new_event_loop()
-        self._loop_thread = threading.Thread(target=self._run_loop, name="LightningSpanExporterLoop", daemon=True)
-        self._loop_thread.start()
+        # Deferred creation until first use.
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[threading.Thread] = None
+
+    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
+        """Lazily initialize the event loop and thread on first use.
+
+        Returns:
+            asyncio.AbstractEventLoop: The initialized event loop.
+        """
+        if self._loop is None:
+            self._loop = asyncio.new_event_loop()
+            self._loop_thread = threading.Thread(target=self._run_loop, name="LightningSpanExporterLoop", daemon=True)
+            self._loop_thread.start()
+        return self._loop
+
+    def _ensure_lock(self) -> threading.RLock:
+        """Lazily initialize the lock on first use.
+
+        Returns:
+            threading.RLock: The initialized lock.
+        """
+        if self._lock is None:
+            self._lock = threading.RLock()
+        return self._lock
 
     def _get_store(self) -> LightningStore:
         """Return the LightningStore to use.
@@ -231,6 +254,7 @@ class LightningSpanExporter(SpanExporter):
 
     def _run_loop(self) -> None:
         """Run the private asyncio loop forever on the exporter thread."""
+        assert self._loop is not None, "Loop should be initialized before thread starts"
         asyncio.set_event_loop(self._loop)
         self._loop.run_forever()
 
@@ -240,13 +264,18 @@ class LightningSpanExporter(SpanExporter):
         Safe to call at process exit.
 
         """
+        if self._loop is None:
+            return
+
         try:
 
             def _stop():
+                assert self._loop is not None
                 self._loop.stop()
 
             self._loop.call_soon_threadsafe(_stop)
-            self._loop_thread.join(timeout=2.0)
+            if self._loop_thread is not None:
+                self._loop_thread.join(timeout=2.0)
             self._loop.close()
         except Exception:
             logger.exception("Error during exporter shutdown")
@@ -264,18 +293,19 @@ class LightningSpanExporter(SpanExporter):
             SpanExportResult: SUCCESS on flush success, else FAILURE.
         """
         # Buffer append under lock to protect against concurrent exporters.
-        with self._lock:
+        with self._ensure_lock():
             for span in spans:
                 self._buffer.append(span)
 
-        # Run the async flush on our private loop, synchronously from caller’s POV.
+        # Run the async flush on our private loop, synchronously from caller's POV.
         async def _locked_flush():
             # Take the lock inside the coroutine to serialize with other flushes.
-            with self._lock:
+            with self._ensure_lock():
                 return await self._maybe_flush()
 
         try:
-            fut = asyncio.run_coroutine_threadsafe(_locked_flush(), self._loop)
+            loop = self._ensure_loop()
+            fut = asyncio.run_coroutine_threadsafe(_locked_flush(), loop)
             fut.result()  # Bubble up any exceptions from the coroutine.
         except Exception as e:
             logger.exception("Export flush failed: %s", e)
@@ -433,6 +463,14 @@ class LightningOpenTelemetry(OpenTelemetry):
 
     def __init__(self, store: LightningStore | None = None):
         config = OpenTelemetryConfig(exporter=LightningSpanExporter(store))
+
+        # Check for tracer initialization
+        if (
+            hasattr(trace_api, "_TRACER_PROVIDER")
+            and trace_api._TRACER_PROVIDER is not None  # pyright: ignore[reportPrivateUsage]
+        ):
+            logger.error("Tracer is already initialized. OpenTelemetry may not work as expected.")
+
         super().__init__(config=config)  # pyright: ignore[reportUnknownMemberType]
 
 
@@ -602,8 +640,8 @@ class LLMProxy:
 
     def as_resource(
         self,
-        rollout_id: str,
-        attempt_id: str,
+        rollout_id: str | None = None,
+        attempt_id: str | None = None,
         model: str | None = None,
         sampling_parameters: Dict[str, Any] | None = None,
     ) -> LLM:
@@ -613,8 +651,8 @@ class LLMProxy:
             ``http://{host}:{port}/rollout/{rollout_id}/attempt/{attempt_id}``
 
         Args:
-            rollout_id: Rollout identifier used for span attribution.
-            attempt_id: Attempt identifier used for span attribution.
+            rollout_id: Rollout identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
+            attempt_id: Attempt identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
             model: Logical model name to use. If omitted and exactly one model
                 is configured, that model is used.
             sampling_parameters: Optional default sampling parameters.
@@ -633,11 +671,20 @@ class LLMProxy:
                     f"Multiple or zero models found in model_list: {self.model_list}. Please specify the model."
                 )
 
-        return LLM(
-            endpoint=f"http://{self.host}:{self.port}/rollout/{rollout_id}/attempt/{attempt_id}",
-            model=model,
-            sampling_parameters=dict(sampling_parameters or {}),
-        )
+        if rollout_id is None and attempt_id is None:
+            return ProxyLLM(
+                endpoint=f"http://{self.host}:{self.port}",
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        elif rollout_id is not None and attempt_id is not None:
+            return LLM(
+                endpoint=f"http://{self.host}:{self.port}/rollout/{rollout_id}/attempt/{attempt_id}",
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        else:
+            raise ValueError("Either rollout_id and attempt_id must be provided, or neither.")
 
 
 def _get_default_ipv4_address() -> str:

agentlightning/reward.py

@@ -2,10 +2,29 @@
 
 import asyncio
 import inspect
+import logging
 import warnings
-from typing import Any, Callable, Literal, Optional, TypedDict, TypeVar
+from typing import (
+    Any,
+    Callable,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    TypedDict,
+    TypeVar,
+    Union,
+    cast,
+)
 
+import opentelemetry.trace as trace_api
 from agentops.sdk.decorators import operation
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import get_tracer_provider
+
+from agentlightning.types import Span, SpanNames
+
+logger = logging.getLogger(__name__)
 
 
 class RewardSpanData(TypedDict):
@@ -69,3 +88,68 @@ def reward(fn: FnType) -> FnType:
             return result
 
         return wrapper  # type: ignore
+
+
+def emit_reward(reward: float) -> ReadableSpan:
+    """
+    Record a new reward as a new span.
+    """
+    logger.debug(f"Emitting reward: {reward}")
+    if isinstance(reward, (int, bool)):
+        reward = float(reward)
+    if not isinstance(reward, float):
+        raise ValueError(f"Reward must be a number, got: {type(reward)}")
+
+    # Check for tracer initialization
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+
+    tracer_provider = get_tracer_provider()
+
+    tracer = tracer_provider.get_tracer("agentlightning")
+    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
+
+
+SpanLike = Union[ReadableSpan, Span]
+
+
+def find_reward_spans(spans: Sequence[SpanLike]) -> List[SpanLike]:
+    """
+    Find all reward spans in the given list of spans.
+
+    Args:
+        spans: A list of spans (either ReadableSpan or Span).
+
+    Returns:
+        A list of spans whose name matches the reward span name.
+    """
+    return [span for span in spans if span.name == SpanNames.REWARD.value]
+
+
+def get_last_reward(spans: Sequence[SpanLike]) -> Optional[float]:
+    """
+    Get the last reward value from a list of spans.
+
+    Args:
+        spans: A list of spans (either ReadableSpan or Span).
+
+    Returns:
+        The reward value from the last reward span, or None if not found.
+    """
+    reward_spans = find_reward_spans(spans)
+    if len(reward_spans) == 0:
+        return None
+    attributes = reward_spans[-1].attributes
+    if attributes:
+        reward = attributes.get("reward", None)
+        if not isinstance(reward, float):
+            logger.error(f"Reward is not a number, got: {type(reward)}. This may cause undefined behaviors.")
+            return cast(float, reward)
+        return reward
+    return None

agentlightning/runner/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .agent import AgentRunnerV2
+from .base import BaseRunner
+from .legacy import AgentRunner
+
+__all__ = [
+    "BaseRunner",
+    "AgentRunner",
+    "AgentRunnerV2",
+]

agentlightning/runner/agent.py

@@ -0,0 +1,505 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent runner implementation for executing agent rollouts.
+
+This module provides the concrete implementation of the runner interface,
+handling the execution of agent rollouts with support for tracing, hooks,
+and distributed worker coordination.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+import uuid
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, TypeVar, cast
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.litagent import LitAgent
+from agentlightning.reward import emit_reward, get_last_reward
+from agentlightning.store.base import LightningStore
+from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import (
+    AttemptedRollout,
+    Hook,
+    NamedResources,
+    RolloutMode,
+    RolloutRawResultV2,
+    RolloutV2,
+    Span,
+)
+
+if TYPE_CHECKING:
+    from agentlightning.execution.events import Event
+
+from .base import BaseRunner
+
+T_task = TypeVar("T_task")
+
+logger = logging.getLogger(__name__)
+
+
+class AgentRunnerV2(BaseRunner[T_task]):
+    """Runner implementation for executing agent tasks with distributed support.
+
+    This runner manages the complete lifecycle of agent rollout execution,
+    including task polling, resource management, tracing, and hooks. It supports
+    both continuous iteration over tasks from the store and single-step execution.
+
+    Attributes:
+        worker_id: The unique identifier for this worker process.
+    """
+
+    def __init__(self, tracer: BaseTracer, max_tasks: Optional[int] = None, poll_interval: float = 5.0) -> None:
+        """Initialize the agent runner.
+
+        Args:
+            tracer: The tracer instance for recording execution traces and spans.
+            max_tasks: Maximum number of tasks to process in iter() mode. If None,
+                the runner will continue indefinitely until interrupted.
+            poll_interval: Time in seconds to wait between polling attempts when
+                no tasks are available in the store.
+        """
+        super().__init__()
+        self._tracer = tracer
+        self._max_tasks = max_tasks
+        self._poll_interval = poll_interval
+
+        # Set later
+        self._agent: Optional[LitAgent[T_task]] = None
+        self._hooks: Sequence[Hook] = []
+        self._store: Optional[LightningStore] = None
+        self.worker_id: Optional[int] = None
+
+    def init(self, agent: LitAgent[T_task], *, hooks: Optional[Sequence[Hook]] = None, **kwargs: Any) -> None:
+        """Initialize the runner with the agent.
+
+        This sets up the agent-runner relationship, registers hooks, and
+        initializes the tracer.
+
+        Args:
+            agent: The LitAgent instance to be managed by this runner.
+            hooks: Optional sequence of Hook objects to be called at various
+                lifecycle stages (on_trace_start, on_trace_end, on_rollout_start,
+                on_rollout_end).
+            **kwargs: Additional initialization arguments (currently unused).
+        """
+        self._agent = agent
+        self._agent.set_runner(self)
+        self._hooks = [*hooks] if hooks is not None else []
+
+        self._tracer.init()
+
+    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        """Initialize the runner for each worker with worker_id and store.
+
+        This method is called once per worker in a distributed setup to provide
+        the worker with its ID and store connection.
+
+        Args:
+            worker_id: Unique identifier for this worker process.
+            store: The LightningStore instance for task coordination and data persistence.
+            **kwargs: Additional worker-specific initialization arguments (currently unused).
+        """
+        self._store = store
+        self.worker_id = worker_id
+
+        self._tracer.init_worker(worker_id)
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Teardown the runner and clean up all resources.
+
+        This method resets all internal state including the agent, store,
+        hooks, and worker ID, and calls the tracer's teardown method.
+
+        Args:
+            *args: Additional teardown arguments (currently unused).
+            **kwargs: Additional teardown keyword arguments (currently unused).
+        """
+        self._agent = None
+        self._store = None
+        self.worker_id = None
+        self._hooks = []
+
+        self._tracer.teardown()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Teardown the runner for a specific worker.
+
+        This method cleans up worker-specific resources and resets the worker ID.
+
+        Args:
+            worker_id: The unique identifier of the worker being torn down.
+            *args: Additional teardown arguments (currently unused).
+            **kwargs: Additional teardown keyword arguments (currently unused).
+        """
+        self.worker_id = None
+
+        self._tracer.teardown_worker(worker_id)
+
+    def get_agent(self) -> LitAgent[T_task]:
+        """Get the agent instance.
+
+        Returns:
+            The LitAgent instance managed by this runner.
+
+        Raises:
+            ValueError: If the agent has not been initialized via init().
+        """
+        if self._agent is None:
+            raise ValueError("Agent not initialized. Call init() first.")
+        return self._agent
+
+    def get_store(self) -> LightningStore:
+        """Get the store instance.
+
+        Returns:
+            The LightningStore instance for this worker.
+
+        Raises:
+            ValueError: If the store has not been initialized via init_worker().
+        """
+        if self._store is None:
+            raise ValueError("Store not initialized. Call init_worker() first.")
+        return self._store
+
+    def get_worker_id(self) -> str:
+        """Get the formatted worker ID string.
+
+        Returns:
+            A formatted string like "Worker-0" if initialized, or "Worker-Unknown"
+            if the worker ID has not been set.
+        """
+        return f"Worker-{self.worker_id}" if self.worker_id is not None else "Worker-Unknown"
+
+    def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
+        """Generate a standardized log prefix for the current worker.
+
+        This creates a consistent prefix format for log messages to identify
+        which worker and rollout the message is associated with.
+
+        Args:
+            rollout_id: Optional rollout ID to include in the prefix.
+
+        Returns:
+            A formatted log prefix string like "[Worker 0 | Rollout xyz]",
+            "[Worker 0]", "[Rollout xyz]", or "[Default Worker]".
+        """
+        if self.worker_id is not None:
+            if rollout_id:
+                return f"[Worker {self.worker_id} | Rollout {rollout_id}]"
+            else:
+                return f"[Worker {self.worker_id}]"
+        if rollout_id:
+            return f"[Rollout {rollout_id}]"
+        return "[Default Worker]"
+
+    async def _trigger_hooks(
+        self,
+        hook_type: Literal["on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end"],
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Trigger all registered hooks of a specific type.
+
+        This method calls the specified hook method on all registered hooks,
+        catching and logging any exceptions that occur during hook execution
+        to prevent them from disrupting the main execution flow.
+
+        Args:
+            hook_type: The type of hook to trigger. Valid values are:
+                "on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end".
+            *args: Positional arguments to pass to the hook methods.
+            **kwargs: Keyword arguments to pass to the hook methods.
+        """
+        for hook in self._hooks:
+            try:
+                await getattr(hook, hook_type)(*args, **kwargs)
+            except Exception:
+                logger.exception(f"{self._log_prefix()} Exception during {hook_type} hook {hook}.")
+
+    async def _post_process_rollout_result(
+        self, rollout: AttemptedRollout, raw_result: RolloutRawResultV2
+    ) -> List[ReadableSpan] | List[Span]:
+        """Standardizes the agent's return value and report what's needed to report to the store.
+
+        Args:
+            rollout: The rollout object for the current task.
+            raw_result: The output from the agent's rollout method.
+
+        Returns:
+            The spans that are assumed to be added to the store.
+            This only serves as an estimation for logging purposes. For precise tracking, use the store directly.
+        """
+        store = self.get_store()
+
+        trace_spans: list[ReadableSpan] | list[Span] = []
+
+        # Case 0: result is None
+        if raw_result is None:
+            trace_spans = self._tracer.get_last_trace()
+
+        # Case 1: result is a float (final reward)
+        if isinstance(raw_result, float):
+            # Preserve the existing spans before another span is emitted
+            trace_spans = list(self._tracer.get_last_trace())
+            # This will emit another span to the tracer
+            reward_span = emit_reward(raw_result)
+            await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
+            trace_spans.append(reward_span)
+
+        if isinstance(raw_result, list):
+            # For rollout methods that return a list, we assume that the returned spans
+            # are the complete span set from the whole rollout
+            trace_spans = raw_result
+
+            # Case 2: result is a list of ReadableSpan (OpenTelemetry spans)
+            if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result):
+
+                if not isinstance(
+                    self._tracer, AgentOpsTracer
+                ):  # TODO: this should be replaced with general OpenTelemetry tracer in next version
+                    for span in raw_result:
+                        await store.add_otel_span(
+                            rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
+                        )
+                else:
+                    logger.warning(
+                        f"{self._log_prefix(rollout.rollout_id)} Tracer is already an OpenTelemetry tracer. "
+                        "The traces should have already been added to the store. "
+                        "No need to return anything from rollout."
+                    )
+
+            # Case 3: result is a list of Span (agentlightning spans)
+            elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result):
+                # Add the spans directly to the store
+                for span in raw_result:
+                    await store.add_span(cast(Span, span))
+                trace_spans = raw_result
+
+            # Left over cases for list
+            elif len(raw_result) == 0:
+                logger.warning(
+                    f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. "
+                    "Please check your rollout implementation."
+                )
+                trace_spans = raw_result
+
+            else:
+                types = [type(t).__name__ for t in raw_result][:10]
+                raise ValueError(
+                    f"Invalid raw result type. It's expected to be a list of ReadableSpan or Span, "
+                    f"but got: {', '.join(types)}..."
+                )
+
+        return trace_spans
+
+    async def _sleep_until_next_poll(self, event: Optional[Event] = None) -> None:
+        """Sleep until the next poll interval, with optional event-based interruption.
+
+        If an event is provided, the method will check it periodically (every 0.1s)
+        and return early if the event is set.
+
+        Args:
+            event: Optional Event object that can be used to interrupt the sleep.
+                If set during the sleep period, the method returns immediately.
+        """
+        if event is None:
+            await asyncio.sleep(self._poll_interval)
+            return
+        current_time = time.time()
+        next_time = current_time + self._poll_interval
+        while time.time() < next_time:
+            await asyncio.sleep(0.1)
+            if event.is_set():
+                return
+
+    async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: bool = False) -> None:
+        """Execute a single rollout implementation.
+
+        This is the core method that handles the execution of a single rollout,
+        including resource fetching, hook triggering, agent invocation, tracing,
+        and result processing.
+
+        Args:
+            next_rollout: The rollout to execute, containing input data, mode,
+                and resources information.
+            raise_on_exception: If True, exceptions during rollout execution will
+                be re-raised. If False, exceptions are logged but not propagated.
+        """
+        store = self.get_store()
+        agent = self.get_agent()
+
+        rollout_id = next_rollout.rollout_id
+
+        resources_id = next_rollout.resources_id
+        resources_update = None
+        if resources_id:
+            resources_update = await store.get_resources_by_id(resources_id)
+        else:
+            logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
+            resources_update = await store.get_latest_resources()
+        if not resources_update:
+            if raise_on_exception:
+                raise RuntimeError(f"{self._log_prefix(rollout_id)} Failed to fetch resources")
+            else:
+                logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
+                return
+
+        trace_spans: List[ReadableSpan] | List[Span] = []
+        has_exception: bool = False
+
+        try:
+            await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)
+
+            start_time = time.time()
+            with self._tracer.trace_context(
+                name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
+            ):
+                await self._trigger_hooks(
+                    hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
+                )
+
+                # NOTE: This is the most costly step in the whole function
+                # If the rollout method becomes unresponsive or timeouts, there is nothing we can do within the runner.
+                # We might need some mechanisms in execution strategy to restart the runner. But that's a future work.
+                if agent.is_async:
+                    rollout_method = (
+                        agent.training_rollout_async if next_rollout.mode == "train" else agent.validation_rollout_async
+                    )
+                    result = await rollout_method(
+                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
+                    )
+                else:
+                    rollout_method = (
+                        agent.training_rollout if next_rollout.mode == "train" else agent.validation_rollout
+                    )
+                    result = rollout_method(
+                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
+                    )
+
+                await self._trigger_hooks(
+                    hook_type="on_trace_end", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
+                )
+
+            trace_spans = await self._post_process_rollout_result(next_rollout, result)
+            last_reward = get_last_reward(trace_spans)
+
+            end_time = time.time()
+            logger.info(
+                f"{self._log_prefix(rollout_id)} Completed in "
+                f"{end_time - start_time:.2f}s. Collected {len(trace_spans)} span(s). "
+                f"Final reward: {last_reward}"
+            )
+
+        except Exception:
+            logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
+            has_exception = True
+
+            if raise_on_exception:
+                raise
+        finally:
+            try:
+                await self._trigger_hooks(
+                    hook_type="on_rollout_end", agent=agent, runner=self, rollout=next_rollout, spans=trace_spans
+                )
+            except Exception:
+                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+
+            if has_exception:
+                # possibly timed out and cancelled?
+                await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="failed")
+            else:
+                await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="succeeded")
+
+    async def iter(self, *, event: Optional[Event] = None) -> None:
+        """Run the runner, continuously iterating over tasks in the store.
+
+        This method polls the store for new rollouts and executes them until:
+        - The event is set (if provided)
+        - The max_tasks limit is reached (if configured)
+        - No more tasks are available
+
+        All exceptions during rollout execution are caught and logged but not
+        propagated, allowing the runner to continue processing subsequent tasks.
+
+        Args:
+            event: Optional Event object to signal the runner to stop. The runner
+                will check this event periodically and stop gracefully when set.
+        """
+        num_tasks_processed = 0
+        logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_tasks or 'unlimited'}).")
+        store = self.get_store()
+
+        while not (event is not None and event.is_set()) and (
+            self._max_tasks is None or num_tasks_processed < self._max_tasks
+        ):
+            # Retrieve the next rollout
+            next_rollout: Optional[RolloutV2] = None
+            while not (event is not None and event.is_set()):
+                logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
+                next_rollout = await store.dequeue_rollout()
+                if next_rollout is None:
+                    logger.debug(f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds.")
+                    await self._sleep_until_next_poll(event)
+                else:
+                    break
+
+            if next_rollout is None:
+                return
+
+            # 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()
+            )
+
+            # Execute the step
+            await self._step_impl(next_rollout)
+
+            num_tasks_processed += 1
+            if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
+                logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_tasks or 'unlimited'}")
+
+        logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[NamedResources] = None,
+        mode: Optional[RolloutMode] = None,
+        event: Optional[Event] = None,
+    ) -> None:
+        """Execute a single task directly, bypassing the task queue.
+
+        This method creates a new rollout for the given input and executes it
+        immediately. Unlike iter(), exceptions are propagated to the caller.
+
+        Args:
+            input: The task input to be processed by the agent.
+            resources: Optional named resources to be used for this specific task.
+                If provided, a new resources entry will be created in the store.
+                If not provided, the latest resources from the store will be used.
+            mode: Optional rollout mode ("train" or "validation"). If not provided,
+                the agent's default mode will be used.
+            event: Optional Event object to signal interruption (currently unused
+                but included for interface consistency).
+
+        Raises:
+            Exception: Any exception that occurs during rollout execution will be
+                re-raised to the caller.
+        """
+        store = self.get_store()
+
+        if resources is not None:
+            # TODO: move this to store.add_resources()
+            resources_id = "resource-" + str(uuid.uuid4())
+            await store.update_resources(resources_id=resources_id, resources=resources)
+        else:
+            resources_id = None
+
+        attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
+        await self._step_impl(attempted_rollout, raise_on_exception=True)

agentlightning/runner/base.py

@@ -0,0 +1,153 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Base runner interface for executing agent tasks.
+
+This module defines the abstract base class for all runner implementations
+in the agent-lightning framework. Runners are responsible for managing the
+execution lifecycle of agents and coordinating with the store.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
+
+from agentlightning.litagent import LitAgent
+from agentlightning.store.base import LightningStore
+from agentlightning.types import NamedResources, ParallelWorkerBase, RolloutMode
+
+if TYPE_CHECKING:
+    from agentlightning.execution.events import Event
+
+
+T_task = TypeVar("T_task")
+
+
+class BaseRunner(ParallelWorkerBase, Generic[T_task]):
+    """Base class for all runners.
+
+    This abstract base class defines the interface that all runner implementations
+    must follow. Runners are responsible for executing agent tasks, managing the
+    execution lifecycle, and coordinating with the store.
+    """
+
+    def init(self, agent: LitAgent[T_task], **kwargs: Any) -> None:
+        """Initialize the runner with the agent.
+
+        This method is called once during setup to configure the runner with
+        the agent it will execute.
+
+        Args:
+            agent: The LitAgent instance to be managed by this runner.
+            **kwargs: Additional initialization arguments specific to the runner implementation.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        """Initialize the runner for each worker with worker_id and store.
+
+        This method is called once per worker process in a distributed setup.
+        It provides the worker with its unique ID and the store instance for
+        task coordination.
+
+        Args:
+            worker_id: Unique identifier for this worker process.
+            store: The LightningStore instance for task coordination and data persistence.
+            **kwargs: Additional worker-specific initialization arguments.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    def run(self, *args: Any, **kwargs: Any) -> None:
+        """Undefined method - use iter() or step() instead.
+
+        This method is intentionally not implemented as the execution behavior
+        should be defined through iter() for continuous execution or step()
+        for single-task execution.
+
+        Args:
+            *args: Unused positional arguments.
+            **kwargs: Unused keyword arguments.
+
+        Raises:
+            RuntimeError: Always raised to indicate this method should not be used.
+        """
+        raise RuntimeError("The behavior of run() of Runner is undefined. Use iter() or step() instead.")
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Clean up runner resources and reset state.
+
+        This method is called once during shutdown to clean up any resources
+        allocated during initialization and reset the runner state.
+
+        Args:
+            *args: Additional teardown arguments.
+            **kwargs: Additional teardown keyword arguments.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Clean up worker-specific resources.
+
+        This method is called once per worker during shutdown to clean up
+        any resources specific to that worker.
+
+        Args:
+            worker_id: The unique identifier of the worker being torn down.
+            *args: Additional teardown arguments.
+            **kwargs: Additional teardown keyword arguments.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    async def iter(self, *, event: Optional[Event] = None) -> None:
+        """Run the runner, continuously iterating over tasks in the store.
+
+        This method runs in a loop, polling the store for new tasks and executing
+        them until interrupted by the event or when no more tasks are available.
+
+        Args:
+            event: Optional Event object that can be used to signal the runner
+                to stop gracefully. When set, the runner should finish its current
+                task and exit the iteration loop.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[NamedResources] = None,
+        mode: Optional[RolloutMode] = None,
+        event: Optional[Event] = None,
+    ) -> None:
+        """Execute a single task with the given input.
+
+        This method provides fine-grained control for executing individual tasks
+        directly, bypassing the store's task queue.
+
+        Args:
+            input: The task input to be processed by the agent.
+            resources: Optional named resources to be used for this specific task.
+                If not provided, the latest resources from the store will be used.
+            mode: Optional rollout mode (e.g., "train", "test"). If not provided,
+                the default mode will be used.
+            event: Optional Event object to signal interruption. When set, the
+                runner may abort the current execution.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()

agentlightning/runner/legacy.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 import json
 import logging
 import time
@@ -7,16 +9,22 @@ from typing import Any, Dict, List, Optional, cast
 
 from opentelemetry.sdk.trace import ReadableSpan
 
-from .adapter import TraceTripletAdapter
-from .client import AgentLightningClient
-from .litagent import LitAgent, is_v0_1_rollout_api
-from .tracer.base import BaseTracer
-from .types import ParallelWorkerBase, Rollout, RolloutRawResult, Triplet
+from agentlightning.adapter import TraceTripletAdapter
+from agentlightning.client import AgentLightningClient
+from agentlightning.litagent import LitAgent, is_v0_1_rollout_api
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import Rollout, RolloutRawResult, Triplet
+
+from .base import BaseRunner
 
 logger = logging.getLogger(__name__)
 
+__all__ = [
+    "AgentRunner",
+]
 
-class AgentRunner(ParallelWorkerBase):
+
+class AgentRunner(BaseRunner[Any]):
     """Manages the agent's execution loop and integrates with AgentOps.
 
     This class orchestrates the interaction between the agent (`LitAgent`) and
@@ -51,6 +59,19 @@ class AgentRunner(ParallelWorkerBase):
         self.worker_id = worker_id
         self.max_tasks = max_tasks
 
+    # These methods are overridden by BaseRunner, getting them back to old behavior.
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        self.worker_id = worker_id
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
     def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
         """Generates a standardized log prefix for the current worker."""
         if self.worker_id is not None:
@@ -131,7 +152,7 @@ class AgentRunner(ParallelWorkerBase):
             return result.model_copy(update=result_dict)
         return Rollout(**result_dict)
 
-    def run(self) -> bool:
+    def run(self) -> bool:  # type: ignore
         """Poll the task and rollout once synchronously."""
         self.agent.set_runner(self)  # Ensure the agent has a reference to this runner
 
@@ -193,7 +214,7 @@ class AgentRunner(ParallelWorkerBase):
 
         return True
 
-    def iter(self) -> int:
+    def iter(self) -> int:  # type: ignore
         """Executes the synchronous polling and rollout loop."""
         num_tasks_processed = 0
         logger.info(f"{self._log_prefix()} Started sync rollouts (max: {self.max_tasks or 'unlimited'}).")

agentlightning/store/base.py

@@ -6,7 +6,6 @@ from typing import Any, Dict, List, Literal, Optional, Sequence
 
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agentlightning.tracer import Span
 from agentlightning.types import (
     Attempt,
     AttemptedRollout,
@@ -16,6 +15,7 @@ from agentlightning.types import (
     RolloutConfig,
     RolloutStatus,
     RolloutV2,
+    Span,
     TaskInput,
 )
 

agentlightning/store/client_server.py

@@ -5,6 +5,9 @@ from __future__ import annotations
 import asyncio
 import logging
 import os
+import threading
+import time
+from contextlib import suppress
 from typing import Any, Dict, List, Literal, Optional, Sequence, Union
 
 import aiohttp
@@ -13,7 +16,6 @@ from fastapi import FastAPI
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel, Field
 
-from agentlightning.tracer import Span
 from agentlightning.types import (
     Attempt,
     AttemptedRollout,
@@ -23,6 +25,7 @@ from agentlightning.types import (
     RolloutConfig,
     RolloutStatus,
     RolloutV2,
+    Span,
     TaskInput,
 )
 
@@ -91,10 +94,12 @@ class LightningStoreServer(LightningStore):
         self.app: FastAPI | None = FastAPI(title="LightningStore Server")
         self._setup_routes()
         self._uvicorn_config: uvicorn.Config | None = uvicorn.Config(
-            self.app, host=self.host, port=self.port, log_level="info"
+            self.app, host="0.0.0.0", port=self.port, log_level="info"
         )
         self._uvicorn_server: uvicorn.Server | None = uvicorn.Server(self._uvicorn_config)
 
+        self._serving_thread: Optional[threading.Thread] = None
+
         # Process-awareness:
         # LightningStoreServer holds a plain Python object (self.store) in one process
         # (the process that runs uvicorn/FastAPI).
@@ -148,8 +153,25 @@ class LightningStoreServer(LightningStore):
         """
         assert self._uvicorn_server is not None
         logger.info(f"Starting server at {self.endpoint}")
-        asyncio.create_task(self._uvicorn_server.serve())
-        await asyncio.sleep(1)  # Allow time for server to start up.
+
+        uvicorn_server = self._uvicorn_server
+
+        def run_server_forever():
+            asyncio.run(uvicorn_server.serve())
+
+        self._serving_thread = threading.Thread(target=run_server_forever, daemon=True)
+        self._serving_thread.start()
+
+        # Wait for /health to be available
+        current_time = time.time()
+        while time.time() - current_time < 10:
+            async with aiohttp.ClientSession() as session:
+                with suppress(Exception):
+                    async with session.get(f"{self.endpoint}/health") as response:
+                        if response.status == 200:
+                            return
+            await asyncio.sleep(0.1)
+        raise RuntimeError("Server failed to start within the 10 seconds.")
 
     async def stop(self):
         """Gracefully stops the running FastAPI server.
@@ -160,7 +182,9 @@ class LightningStoreServer(LightningStore):
         if self._uvicorn_server.started:
             logger.info("Stopping server...")
             self._uvicorn_server.should_exit = True
-            await asyncio.sleep(1)  # Allow time for graceful shutdown.
+            if self._serving_thread is not None:
+                self._serving_thread.join(timeout=10)
+            self._serving_thread = None
             logger.info("Server stopped.")
 
     def _backend(self) -> LightningStore:
@@ -179,6 +203,10 @@ class LightningStoreServer(LightningStore):
         """Set up FastAPI routes for all store operations."""
         assert self.app is not None
 
+        @self.app.get("/health")
+        async def health():  # pyright: ignore[reportUnusedFunction]
+            return {"status": "ok"}
+
         @self.app.post("/start_rollout", response_model=AttemptedRollout)
         async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.store.start_rollout(
@@ -393,18 +421,55 @@ class LightningStoreClient(LightningStore):
 
     def __init__(self, server_address: str):
         self.server_address = server_address.rstrip("/")
-        self.session: Optional[aiohttp.ClientSession] = None
+        self._sessions: Dict[int, aiohttp.ClientSession] = {}  # id(loop) -> ClientSession
+        self._lock = threading.RLock()
 
     async def _get_session(self) -> aiohttp.ClientSession:
-        """Get or create aiohttp session."""
-        if self.session is None or self.session.closed:
-            self.session = aiohttp.ClientSession()
-        return self.session
+        # In the proxy process, FastAPI middleware calls
+        # client_store.get_next_span_sequence_id(...). With
+        # reuse_session=True, _get_session() creates and caches a
+        # single ClientSession bound to the uvicorn event loop.
+        #
+        # Later, the OpenTelemetry exporter (LightningSpanExporter)
+        # runs its flush on its own private event loop (in a different
+        # thread) and calls client_store.add_otel_span(...) ->
+        # client_store.add_span(...).
+        #
+        # If we reuse one session across all, the exporter tries to reuse the
+        # same cached ClientSession that was created on the uvicorn
+        # loop. aiohttp.ClientSession is not loop-agnostic or
+        # thread-safe. Using it from another loop can hang on the
+        # first request. That's why we need a map from loop to session.
+
+        loop = asyncio.get_running_loop()
+        key = id(loop)
+        with self._lock:
+            sess = self._sessions.get(key)
+            if sess is None or sess.closed:
+                sess = aiohttp.ClientSession()
+                self._sessions[key] = sess
+        return sess
 
     async def close(self):
         """Close the HTTP session."""
-        if self.session and not self.session.closed:
-            await self.session.close()
+        with self._lock:
+            sessions = list(self._sessions.values())
+            self._sessions.clear()
+
+        # close them on their own loops to avoid warnings
+        async def _close(sess: aiohttp.ClientSession):
+            if not sess.closed:
+                await sess.close()
+
+        # If called from one loop, best-effort close here.
+        for s in sessions:
+            try:
+                await _close(s)
+            except RuntimeError:
+                # If created on a different loop/thread, schedule a thread-safe close
+                # Fallback: close without awaiting (library tolerates it in practice),
+                # or keep a per-loop shutdown hook where they were created.
+                pass
 
     async def start_rollout(
         self,

agentlightning/store/memory.py

@@ -12,7 +12,6 @@ from typing import Any, Callable, Counter, Dict, List, Literal, Optional, Sequen
 
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agentlightning.tracer import Span
 from agentlightning.types import (
     Attempt,
     AttemptedRollout,
@@ -22,6 +21,7 @@ from agentlightning.types import (
     RolloutConfig,
     RolloutStatus,
     RolloutV2,
+    Span,
     TaskInput,
 )
 

agentlightning/store/threading.py

@@ -7,7 +7,6 @@ from typing import Any, Dict, List, Literal, Optional, Sequence
 
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agentlightning.tracer import Span
 from agentlightning.types import (
     Attempt,
     AttemptedRollout,
@@ -17,6 +16,7 @@ from agentlightning.types import (
     RolloutConfig,
     RolloutStatus,
     RolloutV2,
+    Span,
     TaskInput,
 )
 

agentlightning/tracer/__init__.py

@@ -2,6 +2,5 @@
 
 from .agentops import AgentOpsTracer
 from .base import BaseTracer
-from .types import Span
 
-__all__ = ["AgentOpsTracer", "BaseTracer", "Span"]
+__all__ = ["AgentOpsTracer", "BaseTracer"]

agentlightning/tracer/agentops.py

@@ -2,10 +2,12 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
+import threading
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, Awaitable, Iterator, List, Optional
 
 import agentops
 import agentops.sdk.core
@@ -15,6 +17,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 
 from agentlightning.instrumentation import instrument_all, uninstrument_all
 from agentlightning.instrumentation.agentops import AgentOpsServerManager
+from agentlightning.store.base import LightningStore
 
 from .base import BaseTracer
 
@@ -134,6 +137,8 @@ class AgentOpsTracer(BaseTracer):
         try:
             # new versions
             instance = agentops.sdk.core.tracer
+            # TODO: The span processor cannot be deleted once added.
+            # This might be a problem if the tracer is entered and exited multiple times.
             instance.provider.add_span_processor(self._lightning_span_processor)  # type: ignore
         except AttributeError:
             # old versions
@@ -148,12 +153,22 @@ class AgentOpsTracer(BaseTracer):
             logger.info(f"[Worker {worker_id}] Instrumentation removed.")
 
     @contextmanager
-    def trace_context(self, name: Optional[str] = None) -> Iterator[LightningSpanProcessor]:
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[LightningSpanProcessor]:
         """
         Starts a new tracing context. This should be used as a context manager.
 
         Args:
             name: Optional name for the tracing context.
+            store: Optional store to add the spans to.
+            rollout_id: Optional rollout ID to add the spans to.
+            attempt_id: Optional attempt ID to add the spans to.
 
         Yields:
             The LightningSpanProcessor instance to collect spans.
@@ -161,8 +176,15 @@ class AgentOpsTracer(BaseTracer):
         if not self._lightning_span_processor:
             raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
 
-        with self._lightning_span_processor:
-            yield self._lightning_span_processor
+        if store is not None and rollout_id is not None and attempt_id is not None:
+            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
+            with ctx as processor:
+                yield processor
+        elif store is None and rollout_id is None and attempt_id is None:
+            with self._lightning_span_processor:
+                yield self._lightning_span_processor
+        else:
+            raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
 
     def get_last_trace(self) -> List[ReadableSpan]:
         """
@@ -201,8 +223,29 @@ class AgentOpsTracer(BaseTracer):
 
 
 class LightningSpanProcessor(SpanProcessor):
+    def __init__(self):
+        self._spans: List[ReadableSpan] = []
 
-    _spans: List[ReadableSpan] = []
+        # Store related context and states
+        self._store: Optional[LightningStore] = None
+        self._rollout_id: Optional[str] = None
+        self._attempt_id: Optional[str] = None
+        self._lock = threading.Lock()
+
+        # private asyncio loop running in a daemon thread
+        self._loop_ready = threading.Event()
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
+        self._loop_thread.start()
+        self._loop_ready.wait()  # loop is ready
+
+    def _loop_runner(self):
+        loop = asyncio.new_event_loop()
+        self._loop = loop
+        asyncio.set_event_loop(loop)
+        self._loop_ready.set()
+        loop.run_forever()
+        loop.close()
 
     def __enter__(self):
         self._last_trace = None
@@ -210,7 +253,25 @@ class LightningSpanProcessor(SpanProcessor):
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
-        pass
+        self._store = None
+        self._rollout_id = None
+        self._attempt_id = None
+
+    def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
+        # submit to the dedicated loop and wait synchronously
+        if self._loop is None:
+            raise RuntimeError("Loop is not initialized. This should not happen.")
+        fut = asyncio.run_coroutine_threadsafe(coro, self._loop)  # type: ignore
+        return fut.result(timeout=timeout)  # raises on error  # type: ignore
+
+    def shutdown(self) -> None:
+        if self._loop:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._loop_thread.join(timeout=5)
+            self._loop = None
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
 
     def spans(self) -> List[ReadableSpan]:
         """
@@ -222,6 +283,22 @@ class LightningSpanProcessor(SpanProcessor):
         """
         return self._spans
 
+    def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
+        # simple context manager without nesting into asyncio
+        class _Ctx:
+            def __enter__(_):  # type: ignore
+                with self._lock:
+                    self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
+                    self._last_trace = None
+                    self._spans = []
+                return self
+
+            def __exit__(_, exc_type, exc, tb):  # type: ignore
+                with self._lock:
+                    self._store = self._rollout_id = self._attempt_id = None
+
+        return _Ctx()
+
     def on_end(self, span: ReadableSpan) -> None:
         """
         Process a span when it ends.
@@ -233,10 +310,15 @@ class LightningSpanProcessor(SpanProcessor):
         if not span.context or not span.context.trace_flags.sampled:
             return
 
+        if self._store and self._rollout_id and self._attempt_id:
+            try:
+                # Submit add_otel_span to the event loop and wait for it to complete
+                self._await_in_loop(
+                    self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
+                    timeout=5.0,
+                )
+            except Exception:
+                # log; on_end MUST NOT raise
+                logger.exception(f"Error adding span to store: {span.name}")
+
         self._spans.append(span)
-
-    def shutdown(self) -> None:
-        pass
-
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
-        return True

agentlightning/tracer/base.py

@@ -1,11 +1,15 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import logging
 from contextlib import contextmanager
-from typing import Any, Awaitable, Callable, Iterator, List, Optional
+from typing import Any, Awaitable, Callable, Iterator, List, Optional, cast
 
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agentlightning.types import ParallelWorkerBase
+from agentlightning.store.base import LightningStore
+from agentlightning.types import ParallelWorkerBase, SpanNames
+
+logger = logging.getLogger(__name__)
 
 
 class BaseTracer(ParallelWorkerBase):
@@ -44,7 +48,14 @@ class BaseTracer(ParallelWorkerBase):
     """
 
     @contextmanager
-    def trace_context(self, name: Optional[str] = None) -> Iterator[Any]:
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[Any]:
         """
         Starts a new tracing context. This should be used as a context manager.
 
@@ -53,8 +64,13 @@ class BaseTracer(ParallelWorkerBase):
         within the `with` block are collected and made available via
         `get_last_trace`.
 
+        If a store is provided, the spans will be added to the store when tracing.
+
         Args:
             name: The name for the root span of this trace context.
+            store: The store to add the spans to.
+            rollout_id: The rollout ID to add the spans to.
+            attempt_id: The attempt ID to add the spans to.
         """
         raise NotImplementedError()
 
@@ -67,6 +83,19 @@ class BaseTracer(ParallelWorkerBase):
         """
         raise NotImplementedError()
 
+    def get_last_reward(self) -> Optional[float]:
+        """
+        Retrieves the finalest reward from the most recent trace.
+        The behavior by default is to traverse the trace backward until the first reward span.
+        """
+        for span in reversed(self.get_last_trace()):
+            if span.name == SpanNames.REWARD.value and span.attributes:
+                reward = span.attributes.get("reward", None)
+                if not isinstance(reward, float):
+                    logger.error(f"Reward is not a number, got: {type(reward)}. This may cause undefined behaviors.")
+                return cast(float, reward)
+        return None
+
     def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single synchronous function.

agentlightning/tracer/http.py

@@ -76,7 +76,7 @@ class HttpTracer(BaseTracer):
         logger.info(f"[Worker {worker_id}] HttpTracer initialized.")
 
     @contextmanager
-    def trace_context(self, name: Optional[str] = None) -> Iterator[HTTPRecords]:
+    def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]:
         """
         Starts a new HTTP tracing context. This should be used as a context manager.
 

agentlightning/trainer.py

@@ -261,7 +261,7 @@ class Trainer(ParallelWorkerBase):
                 max_tasks=self.max_tasks,
                 worker_id=worker_id,
             )
-            loop.init_worker(worker_id)
+            loop.init_worker(worker_id)  # type: ignore
             if is_async:
                 num_processed = asyncio.run(loop.iter_async())
             else:

agentlightning/types/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .core import *
+from .tracer import *

agentlightning/types/core.py

@@ -2,11 +2,32 @@
 
 from __future__ import annotations
 
-from typing import Annotated, Any, Callable, Dict, Generic, List, Literal, Optional, Protocol, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel, Field, model_validator
 
+from .tracer import Span
+
+if TYPE_CHECKING:
+    from agentlightning.litagent import LitAgent
+    from agentlightning.runner.base import BaseRunner
+    from agentlightning.tracer.base import BaseTracer
+
 __all__ = [
     "Triplet",
     "Rollout",
@@ -14,8 +35,11 @@ __all__ = [
     "TaskInput",
     "TaskIfAny",
     "RolloutRawResult",
+    "RolloutRawResultV2",
+    "RolloutMode",
     "Resource",
     "LLM",
+    "ProxyLLM",
     "PromptTemplate",
     "ResourceUnion",
     "NamedResources",
@@ -29,6 +53,7 @@ __all__ = [
     "RolloutV2",
     "Attempt",
     "AttemptedRollout",
+    "Hook",
 ]
 
 T_co = TypeVar("T_co", covariant=True)
@@ -91,6 +116,8 @@ AttemptStatus = Literal[
     "timeout",  # the worker has been emitting new logs, but have been working on the task for too long
 ]
 
+RolloutMode = Literal["train", "val", "test"]
+
 
 class Attempt(BaseModel):
     """An attempt to execute a rollout. A rollout can have multiple attempts if retries are needed."""
@@ -132,7 +159,7 @@ class RolloutV2(BaseModel):
     start_time: float
     end_time: Optional[float] = None
 
-    mode: Optional[Literal["train", "val", "test"]] = None
+    mode: Optional[RolloutMode] = None
     resources_id: Optional[str] = None
 
     # Overall scheduling/running information
@@ -165,7 +192,7 @@ class Task(BaseModel):
     rollout_id: str
     input: TaskInput
 
-    mode: Optional[Literal["train", "val", "test"]] = None
+    mode: Optional[RolloutMode] = None
     resources_id: Optional[str] = None
 
     # Optional fields for tracking task lifecycle
@@ -184,6 +211,13 @@ class TaskIfAny(BaseModel):
 
 RolloutRawResult = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], Rollout]
 
+RolloutRawResultV2 = Union[
+    None,  # nothing (relies on tracer)
+    float,  # only final reward
+    List[ReadableSpan],  # constructed OTEL spans by user
+    List[Span],  # constructed Span objects by user
+]
+
 
 class Resource(BaseModel):
     """
@@ -210,6 +244,32 @@ class LLM(Resource):
     api_key: Optional[str] = None
     sampling_parameters: Dict[str, Any] = Field(default_factory=dict)
 
+    def base_url(self, *args: Any, **kwargs: Any) -> str:
+        return self.endpoint
+
+
+class ProxyLLM(LLM):
+    """Proxy LLM resource that is tailored by `llm_proxy.LLMProxy`."""
+
+    resource_type: Literal["proxy_llm"] = "proxy_llm"  # type: ignore
+
+    def base_url(self, rollout_id: str, attempt_id: str) -> str:
+        prefix = self.endpoint
+        if prefix.endswith("/"):
+            prefix = prefix[:-1]
+        if prefix.endswith("/v1"):
+            prefix = prefix[:-3]
+            has_v1 = True
+        else:
+            has_v1 = False
+        # Now the prefix should look like "http://localhost:11434"
+
+        # Append the rollout and attempt id to the prefix
+        prefix = prefix + f"/rollout/{rollout_id}/attempt/{attempt_id}"
+        if has_v1:
+            prefix = prefix + "/v1"
+        return prefix
+
 
 class PromptTemplate(Resource):
     """
@@ -228,7 +288,7 @@ class PromptTemplate(Resource):
 
 
 # Use discriminated union for proper deserialization
-ResourceUnion = Annotated[Union[LLM, PromptTemplate], Field(discriminator="resource_type")]
+ResourceUnion = Annotated[Union[LLM, ProxyLLM, PromptTemplate], Field(discriminator="resource_type")]
 NamedResources = Dict[str, ResourceUnion]
 """
 A dictionary-like class to hold named resources.
@@ -317,3 +377,70 @@ class Dataset(Protocol, Generic[T_co]):
     def __getitem__(self, index: int) -> T_co: ...
 
     def __len__(self) -> int: ...
+
+
+class Hook(ParallelWorkerBase):
+    """Base class for defining hooks in the agent runner's lifecycle."""
+
+    async def on_trace_start(
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: RolloutV2
+    ) -> None:
+        """Hook called immediately after the tracer enters the trace context but before the rollout begins.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            tracer: The :class:`BaseTracer` instance associated with the runner.
+            rollout: The :class:`RolloutV2` object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource setup. By default, this is a no-op.
+        """
+
+    async def on_trace_end(
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: RolloutV2
+    ) -> None:
+        """Hook called immediately after the rollout completes but before the tracer exits the trace context.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            tracer: The :class:`BaseTracer` instance associated with the runner.
+            rollout: The :class:`RolloutV2` object that has been processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource cleanup. By default, this is a no-op.
+        """
+
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: BaseRunner[Any], rollout: RolloutV2) -> None:
+        """Hook called immediately before a rollout *attempt* begins.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            rollout: The :class:`RolloutV2` object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as
+        logging, metric collection, or resource setup. By default, this is a
+        no-op.
+        """
+
+    async def on_rollout_end(
+        self,
+        *,
+        agent: LitAgent[Any],
+        runner: BaseRunner[Any],
+        rollout: RolloutV2,
+        spans: Union[List[ReadableSpan], List[Span]],
+    ) -> None:
+        """Hook called after a rollout *attempt* completes.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            rollout: The :class:`RolloutV2` object that has been processed.
+            spans: The spans that have been added to the store.
+
+        Subclasses can override this method for cleanup or additional
+        logging. By default, this is a no-op.
+        """

agentlightning/types/tracer.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+from enum import Enum
 from typing import Any, Dict, List, Optional, Sequence, Union
 
 from opentelemetry import trace as trace_api
@@ -12,6 +13,19 @@ from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.trace.status import Status as OtelStatus
 from pydantic import BaseModel
 
+__all__ = [
+    "AttributeValue",
+    "Attributes",
+    "TraceState",
+    "SpanContext",
+    "TraceStatus",
+    "Event",
+    "Link",
+    "Resource",
+    "Span",
+    "SpanNames",
+]
+
 
 def convert_timestamp(timestamp: Optional[int]) -> Optional[float]:
     """Convert timestamp from nanoseconds to seconds if needed.
@@ -229,3 +243,13 @@ class Span(BaseModel):
                 ],
             ),
         )
+
+
+class SpanNames(str, Enum):
+    """Standard span name values for AgentLightning.
+
+    Currently only reward spans are supported.
+    We will add more spans related to error handling in the future.
+    """
+
+    REWARD = "agentlightning.reward"

tests/common/__init__.py

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

tests/common/network.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import socket
+from contextlib import closing
+
+
+def get_free_port() -> int:
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("127.0.0.1", 0))
+        return int(s.getsockname()[1])

tests/common/tracer.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import agentops
+import agentops.sdk.core as agentops_core
 import opentelemetry.trace as trace_api
 
 
@@ -25,3 +27,9 @@ def clear_tracer_provider() -> None:
         if hasattr(trace_api._TRACER_PROVIDER_SET_ONCE, "_flag"):
             if trace_api._TRACER_PROVIDER_SET_ONCE._flag:  # type: ignore
                 trace_api._TRACER_PROVIDER_SET_ONCE._flag = False  # type: ignore
+
+
+def clear_agentops_init() -> None:
+    """Make agentops.init() runnable again."""
+    agentops.get_client().initialized = False
+    agentops_core.tracer._initialized = False

tests/common/vllm.py

@@ -0,0 +1,181 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+import subprocess
+import sys
+import time
+from typing import Any, Optional
+
+import httpx
+import openai
+
+VLLM_AVAILABLE = False
+VLLM_UNAVAILABLE_REASON = ""
+
+try:
+    import vllm
+    from vllm.engine.arg_utils import AsyncEngineArgs
+    from vllm.entrypoints.cli.serve import ServeSubcommand
+    from vllm.model_executor.model_loader import get_model_loader
+    from vllm.utils import FlexibleArgumentParser
+
+    VLLM_AVAILABLE = True  # type: ignore
+    VLLM_VERSION = tuple(int(v) for v in vllm.__version__.split("."))
+except ImportError as e:
+    AsyncEngineArgs = None
+    get_model_loader = None
+    FlexibleArgumentParser = None
+    ServeSubcommand = None
+    VLLM_VERSION = (0, 0, 0)  # type: ignore
+    VLLM_UNAVAILABLE_REASON = str(e)  # type: ignore
+
+
+class RemoteOpenAIServer:
+    """
+    A context manager for launching and interacting with a remote vLLM-based
+    OpenAI-compatible server instance.
+
+    This class handles:
+      - Preparing the environment and spawning the vLLM server process
+      - Ensuring that the requested model is downloaded before server startup
+      - Polling and health-checking the server until it is ready
+      - Providing helper methods to construct URLs for API calls
+      - Returning configured synchronous and asynchronous OpenAI clients
+        that can communicate with the launched server
+
+    Typical usage:
+        with RemoteOpenAIServer(vllm_serve_args, port, model) as server:
+            client = server.get_client()
+            response = client.chat.completions.create(...)
+
+    Attributes:
+        DUMMY_API_KEY (str): A placeholder API key for compatibility
+                             (vLLM does not require authentication).
+        host (str): Host address of the server (default: "localhost").
+        port (int): TCP port number for the server.
+        proc (subprocess.Popen): Handle to the launched server process.
+    """
+
+    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
+
+    def _start_server(self, model: str, vllm_serve_args: list[str], env_dict: Optional[dict[str, str]]) -> None:
+        """Subclasses override this method to customize server process launch"""
+        env = os.environ.copy()
+        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"  # safer CUDA init
+        if env_dict is not None:
+            env.update(env_dict)
+
+        if VLLM_VERSION >= (0, 10, 2):
+            # Supports return_token_ids
+            self.proc: subprocess.Popen[bytes] = subprocess.Popen(
+                ["vllm", "serve", model, *vllm_serve_args],
+                env=env,
+                stdout=sys.stdout,
+                stderr=sys.stderr,
+            )
+        else:
+            # Does not support return_token_ids
+            self.proc = subprocess.Popen(
+                ["python", "-m", "agentlightning.cli.vllm", "serve", model, *vllm_serve_args],
+                env=env,
+                stdout=sys.stdout,
+                stderr=sys.stderr,
+            )
+
+    def __init__(
+        self,
+        model: str,
+        vllm_serve_args: list[str],  # should not include the model name
+        env_dict: Optional[dict[str, str]] = None,
+        seed: Optional[int] = 0,
+        max_wait_seconds: Optional[float] = None,
+    ) -> None:
+        if (
+            not VLLM_AVAILABLE
+            or AsyncEngineArgs is None
+            or get_model_loader is None
+            or FlexibleArgumentParser is None
+            or ServeSubcommand is None
+        ):
+            raise ImportError("vLLM is not available: " + VLLM_UNAVAILABLE_REASON)
+
+        self.model = model
+
+        parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
+        subparsers = parser.add_subparsers(required=False, dest="subparser")
+        parser = ServeSubcommand().subparser_init(subparsers)  # pyright: ignore[reportUnknownMemberType]
+        args = parser.parse_args(["--model", model, *vllm_serve_args])
+        assert args is not None
+        self.host = str(args.host or "localhost")
+        self.port = int(args.port)
+
+        # download the model before starting the server to avoid timeout
+        is_local = os.path.isdir(model)
+        if not is_local:
+            engine_args = AsyncEngineArgs.from_cli_args(args)
+            model_config = engine_args.create_model_config()
+            load_config = engine_args.create_load_config()
+
+            model_loader = get_model_loader(load_config)
+            model_loader.download_model(model_config)
+
+        self._start_server(model, vllm_serve_args, env_dict)
+        max_wait_seconds = max_wait_seconds or 240
+        self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+        self.proc.terminate()
+        try:
+            self.proc.wait(8)
+        except subprocess.TimeoutExpired:
+            self.proc.kill()
+
+    def _poll(self) -> Optional[int]:
+        """Subclasses override this method to customize process polling"""
+        return self.proc.poll()
+
+    def _wait_for_server(self, *, url: str, timeout: float):
+        start = time.time()
+        client = httpx.Client()
+
+        while True:
+            try:
+                if client.get(url).status_code == 200:
+                    break
+            except Exception:
+                result = self._poll()
+                if result is not None and result != 0:
+                    raise RuntimeError("Server exited unexpectedly.") from None
+                time.sleep(0.5)
+                if time.time() - start > timeout:
+                    raise RuntimeError("Server failed to start in time.") from None
+
+    @property
+    def url_root(self) -> str:
+        return f"http://{self.host}:{self.port}"
+
+    def url_for(self, *parts: str) -> str:
+        return self.url_root + "/" + "/".join(parts)
+
+    def get_client(self, **kwargs: Any):
+        if "timeout" not in kwargs:
+            kwargs["timeout"] = 600
+        return openai.OpenAI(
+            base_url=self.url_for("v1"),
+            api_key=self.DUMMY_API_KEY,
+            max_retries=0,
+            **kwargs,
+        )
+
+    def get_async_client(self, **kwargs: Any):
+        if "timeout" not in kwargs:
+            kwargs["timeout"] = 600
+        return openai.AsyncOpenAI(
+            base_url=self.url_for("v1"),
+            api_key=self.DUMMY_API_KEY,
+            max_retries=0,
+            **kwargs,
+        )

tests/runner/__init__.py

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

tests/runner/test_agent_integration.py

@@ -0,0 +1,238 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import multiprocessing
+import os
+import time
+from multiprocessing.synchronize import Event as MpEvent
+from typing import Any, Dict, Optional, cast
+
+import openai
+import pytest
+
+from agentlightning.litagent import LitAgent
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.reward import emit_reward
+from agentlightning.runner import AgentRunnerV2
+from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.types.core import LLM, AttemptedRollout, NamedResources, RolloutV2
+
+from ..common.network import get_free_port
+from ..common.tracer import clear_tracer_provider
+from ..common.vllm import VLLM_AVAILABLE, RemoteOpenAIServer
+
+
+async def init_runner(
+    agent: LitAgent[Any],
+    *,
+    resources: Optional[Dict[str, LLM]] = None,
+) -> tuple[AgentRunnerV2[Any], InMemoryLightningStore]:
+    store = InMemoryLightningStore()
+    llm_resource: NamedResources = resources or {"llm": LLM(endpoint="http://localhost", model="dummy")}  # type: ignore[assignment]
+    await store.update_resources("default", llm_resource)
+
+    runner = AgentRunnerV2[Any](tracer=AgentOpsTracer(), poll_interval=0.01)
+    runner.init(agent)
+    runner.init_worker(worker_id=0, store=store)
+    return runner, store
+
+
+def teardown_runner(runner: AgentRunnerV2[Any]) -> None:
+    runner.teardown_worker(worker_id=0)
+    runner.teardown()
+
+
+@pytest.fixture(scope="module", autouse=True)
+def setup_module():
+    clear_tracer_provider()
+    yield
+
+
+@pytest.mark.asyncio
+async def test_runner_integration_basic_rollout() -> None:
+    class EchoAgent(LitAgent[str]):
+        async def validation_rollout_async(self, task: str, resources: Dict[str, Any], rollout: Any) -> None:
+            emit_reward(1.0)
+
+    agent = EchoAgent()
+    runner, store = await init_runner(agent)
+    try:
+        await runner.step("hello integration")
+    finally:
+        teardown_runner(runner)
+
+    rollouts = await store.query_rollouts()
+    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)
+    print(store.__dict__)
+    assert any(span.attributes.get("reward") == 1.0 for span in spans)
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")),
+    reason="OpenAI endpoint or key not configured",
+)
+async def test_runner_integration_with_openai() -> None:
+    class OpenAIAgent(LitAgent[str]):
+        async def validation_rollout_async(self, task: str, resources: NamedResources, rollout: RolloutV2) -> float:
+            llm = cast(LLM, resources["llm"])
+            client = openai.AsyncOpenAI(base_url=llm.endpoint, api_key=llm.api_key)
+            response = await client.chat.completions.create(
+                model=llm.model,
+                messages=[{"role": "user", "content": task}],
+            )
+            assert response.choices, "OpenAI response should contain choices"
+            return 0.0
+
+    base_url = os.environ["OPENAI_BASE_URL"]
+    api_key = os.environ["OPENAI_API_KEY"]
+    model = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+
+    agent = OpenAIAgent()
+    resources = {"llm": LLM(endpoint=base_url, model=model, api_key=api_key)}
+    runner, store = await init_runner(agent, resources=resources)
+    try:
+        await runner.step("Say hello in one word")
+    finally:
+        teardown_runner(runner)
+
+    rollouts = await store.query_rollouts()
+    assert rollouts and rollouts[0].status == "succeeded"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not (os.getenv("OPENAI_BASE_URL") and os.getenv("OPENAI_API_KEY")),
+    reason="OpenAI endpoint or key not configured",
+)
+async def test_runner_integration_with_litellm_proxy() -> None:
+    litellm = pytest.importorskip("litellm")
+
+    class LiteLLMAgent(LitAgent[str]):
+        def validation_rollout(self, task: str, resources: NamedResources, rollout: RolloutV2) -> float:
+            llm = cast(LLM, resources["llm"])
+            response = litellm.completion(
+                model=llm.model,
+                messages=[{"role": "user", "content": task}],
+            )
+            assert response.get("choices"), "litellm proxy should return choices"
+            return 0.0
+
+    agent = LiteLLMAgent()
+    resources = {"llm": LLM(endpoint="http://dummy", model="openai/gpt-4o-mini")}
+    runner, store = await init_runner(agent, resources=resources)
+    try:
+        await runner.step("Give me a short greeting")
+    finally:
+        teardown_runner(runner)
+
+    rollouts = await store.query_rollouts()
+    assert rollouts and rollouts[0].status == "succeeded"
+
+
+@pytest.fixture
+def server():
+    if not VLLM_AVAILABLE:
+        pytest.skip("vLLM is not available")
+    vllm_port = get_free_port()
+    with RemoteOpenAIServer(
+        model="Qwen/Qwen2.5-0.5B-Instruct",
+        vllm_serve_args=[
+            "--gpu-memory-utilization",
+            "0.7",
+            "--enable-auto-tool-choice",
+            "--tool-call-parser",
+            "hermes",
+            "--port",
+            str(vllm_port),
+        ],
+    ) as server:
+        yield server
+
+
+@pytest.mark.asyncio
+async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenAIServer) -> None:
+    torch = pytest.importorskip("torch")
+    if not torch.cuda.is_available():
+        pytest.skip("GPU not available")
+
+    class ProxyAgent(LitAgent[str]):
+        async def validation_rollout_async(self, task: str, resources: NamedResources, rollout: RolloutV2) -> float:
+            attempted_rollout = cast(AttemptedRollout, rollout)
+            llm_resource = cast(LLM, resources["llm"])
+            client = openai.AsyncOpenAI(
+                base_url=llm_resource.base_url(attempted_rollout.rollout_id, attempted_rollout.attempt.attempt_id),
+                api_key="dummy",
+            )
+            response = await client.chat.completions.create(
+                model=llm_resource.model,
+                messages=[{"role": "user", "content": task}],
+            )
+            assert response.choices, "Proxy should return at least one choice"
+            return 0.5
+
+    agent = ProxyAgent()
+    runner, store = await init_runner(agent)
+
+    server_store = LightningStoreServer(store=store, host="127.0.0.1", port=get_free_port())
+    await server_store.start()
+    client_store = LightningStoreClient(server_store.endpoint)
+
+    proxy = LLMProxy(
+        port=get_free_port(),
+        model_list=[
+            {
+                "model_name": "gpt-4o-arbitrary",
+                "litellm_params": {
+                    "model": "hosted_vllm/" + server.model,
+                    "api_base": server.url_for("v1"),
+                },
+            }
+        ],
+        store=client_store,
+    )
+
+    def run_proxy_server(proxy: LLMProxy, event: MpEvent):
+        clear_tracer_provider()
+        proxy.start()
+        event.set()
+        time.sleep(3600)  # Keep the server running
+
+    event = multiprocessing.Event()
+    process = multiprocessing.Process(target=run_proxy_server, args=(proxy, event))
+    process.start()
+    event.wait(timeout=30)
+
+    try:
+        await runner.step("Say hello to Agent Lightning", resources={"llm": proxy.as_resource()})
+
+        rollouts = await client_store.query_rollouts()
+        assert rollouts and rollouts[0].status == "succeeded"
+
+        spans = await client_store.query_spans(rollouts[0].rollout_id, "latest")
+        first_spans = [span for span in spans if span.sequence_id == 1]
+        assert len(first_spans) > 1
+        assert any("llm.hosted_vllm.choices" in span.attributes for span in first_spans)
+        assert any("llm.hosted_vllm.prompt_token_ids" in span.attributes for span in first_spans)
+        assert any("gen_ai.prompt.0.content" in span.attributes for span in first_spans)
+
+        second_spans = [span for span in spans if span.sequence_id == 2]
+        assert len(second_spans) == 1
+        assert second_spans[0].name == "openai.chat.completion"
+
+        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.reward"
+        assert last_spans[0].attributes.get("reward") == 0.5
+    finally:
+        teardown_runner(runner)
+        process.terminate()
+        await client_store.close()
+        await server_store.stop()
+        await asyncio.to_thread(process.join, timeout=1)
+        if process.is_alive():
+            process.kill()

tests/runner/test_agent_runner.py

@@ -0,0 +1,495 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import random
+from contextlib import contextmanager
+from typing import Any, Dict, Iterator, List, Optional, Sequence, cast
+
+import pytest
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
+from opentelemetry.trace import SpanContext, TraceFlags, TraceState
+from opentelemetry.trace.status import Status, StatusCode
+
+from agentlightning.execution.events import Event, ThreadingEvent
+from agentlightning.litagent import LitAgent
+from agentlightning.reward import emit_reward, get_last_reward
+from agentlightning.runner import AgentRunnerV2
+from agentlightning.runner.base import BaseRunner
+from agentlightning.store.base import LightningStore
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import Hook, RolloutV2, Span
+from agentlightning.types.core import LLM, PromptTemplate
+from agentlightning.types.tracer import SpanNames
+
+trace_api.set_tracer_provider(TracerProvider())
+
+
+def create_readable_span(name: str, attributes: Optional[Dict[str, Any]] = None) -> ReadableSpan:
+    trace_id = random.getrandbits(128)
+    span_id = random.getrandbits(64)
+    context = SpanContext(
+        trace_id=trace_id,
+        span_id=span_id,
+        is_remote=False,
+        trace_flags=TraceFlags(TraceFlags.SAMPLED),
+        trace_state=TraceState(),
+    )
+    status = Status(status_code=StatusCode.UNSET)
+    return ReadableSpan(
+        name=name,
+        context=context,
+        parent=None,
+        resource=Resource.create({}),
+        attributes=attributes or {},
+        events=(),
+        links=(),
+        status=status,
+    )
+
+
+def create_agent_span(
+    rollout_id: str,
+    attempt_id: str,
+    sequence_id: int,
+    name: str,
+    attributes: Optional[Dict[str, Any]] = None,
+) -> Span:
+    readable = create_readable_span(name, attributes)
+    return Span.from_opentelemetry(
+        readable,
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+    )
+
+
+class DummyTracer(BaseTracer):
+    def __init__(self) -> None:
+        super().__init__()
+        self._last_trace: List[ReadableSpan] = []
+        self._contexts: List[Dict[str, Any]] = []
+
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        self._last_trace.clear()
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        self._last_trace.clear()
+
+    def get_last_trace(self) -> List[ReadableSpan]:
+        return list(self._last_trace)
+
+    @contextmanager
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[List[ReadableSpan]]:
+        previous = self._contexts[-1] if self._contexts else None
+        current = {
+            "name": name,
+            "store": store,
+            "rollout_id": rollout_id,
+            "attempt_id": attempt_id,
+        }
+        self._contexts.append(current)
+        self._last_trace = []
+        try:
+            yield self._last_trace
+        finally:
+            self._contexts.pop()
+            if previous is None:
+                self._contexts = []
+
+    def record_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> ReadableSpan:
+        span = create_readable_span(name, attributes)
+        self._last_trace.append(span)
+        return span
+
+
+async def setup_runner(
+    agent: LitAgent[Any],
+    *,
+    tracer: Optional[DummyTracer] = None,
+    max_tasks: Optional[int] = None,
+    poll_interval: float = 0.01,
+    hooks: Sequence[Hook] = (),
+) -> tuple[AgentRunnerV2[Any], InMemoryLightningStore, DummyTracer]:
+    tracer = tracer or DummyTracer()
+    store = InMemoryLightningStore()
+    await store.update_resources("default", {"llm": LLM(endpoint="http://localhost", model="dummy")})
+
+    runner = AgentRunnerV2[Any](tracer=tracer, max_tasks=max_tasks, poll_interval=poll_interval)
+    runner.init(agent=agent, hooks=hooks)
+    runner.init_worker(worker_id=0, store=store)
+    return runner, store, tracer
+
+
+def teardown_runner(runner: AgentRunnerV2[Any]) -> None:
+    runner.teardown_worker(worker_id=0)
+    runner.teardown()
+
+
+async def assert_single_attempt_succeeded(store: InMemoryLightningStore) -> tuple[str, str]:
+    rollouts = await store.query_rollouts()
+    assert len(rollouts) == 1
+    rollout = rollouts[0]
+    attempts = await store.query_attempts(rollout.rollout_id)
+    assert attempts[-1].status == "succeeded"
+    return rollout.rollout_id, attempts[-1].attempt_id
+
+
+class RecordingHook(Hook):
+    def __init__(self) -> None:
+        super().__init__()
+        self.calls: List[str] = []
+        self.received_spans: Optional[List[ReadableSpan] | List[Span]] = None
+
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: BaseRunner[Any], rollout: RolloutV2) -> None:
+        self.calls.append("on_rollout_start")
+
+    async def on_trace_start(
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: RolloutV2
+    ) -> None:
+        self.calls.append("on_trace_start")
+
+    async def on_trace_end(
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: RolloutV2
+    ) -> None:
+        self.calls.append("on_trace_end")
+
+    async def on_rollout_end(
+        self,
+        *,
+        agent: LitAgent[Any],
+        runner: BaseRunner[Any],
+        rollout: RolloutV2,
+        spans: List[ReadableSpan] | List[Span],
+    ) -> None:
+        self.calls.append("on_rollout_end")
+        self.received_spans = spans
+
+
+@pytest.mark.asyncio
+async def test_step_records_spans_for_none_result() -> None:
+    tracer = DummyTracer()
+
+    class AsyncSpanAgent(LitAgent[Dict[str, Any]]):
+        async def validation_rollout_async(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> None:
+            span = tracer.record_span("work", {"task_id": task["task_id"]})
+            store = cast(AgentRunnerV2[Dict[str, Any]], self.runner).get_store()
+            await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, span)  # type: ignore[attr-defined]
+            return None
+
+    agent = AsyncSpanAgent()
+    runner, store, _ = await setup_runner(agent, tracer=tracer)
+    try:
+        await runner.step({"task_id": 1})
+    finally:
+        teardown_runner(runner)
+
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    assert [span.name for span in spans] == ["work"]
+    assert get_last_reward(spans) is None
+
+
+@pytest.mark.asyncio
+async def test_step_emits_reward_for_float_result() -> None:
+    class RewardAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> float:
+            return 0.75
+
+    agent = RewardAgent()
+    runner, store, _ = await setup_runner(agent)
+    try:
+        await runner.step({"prompt": "hello"})
+    finally:
+        teardown_runner(runner)
+
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    rewards = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
+    assert rewards == [0.75]
+
+
+@pytest.mark.asyncio
+async def test_step_handles_non_llm_resource() -> None:
+    class PromptAgent(LitAgent[str]):
+        def validation_rollout(self, task: str, resources: Dict[str, Any], rollout: Any) -> float:
+            template = resources["template"]
+            assert isinstance(template, PromptTemplate)
+            rendered = template.template.format(name=task)
+            assert task in rendered
+            return 0.1
+
+    agent = PromptAgent()
+    runner, store, _ = await setup_runner(agent)
+    try:
+        await store.update_resources(
+            "prompt-resource",
+            {"template": PromptTemplate(template="Hello {name}!", engine="f-string")},
+        )
+        await runner.step("Ada")
+    finally:
+        teardown_runner(runner)
+
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    rewards = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
+    assert rewards == [0.1]
+
+
+@pytest.mark.asyncio
+async def test_step_accepts_readable_span_list() -> None:
+    class ReadableSpanAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(
+            self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
+        ) -> List[ReadableSpan]:
+            return [create_readable_span(f"trace-{i}", {"idx": i}) for i in range(2)]
+
+    agent = ReadableSpanAgent()
+    runner, store, _ = await setup_runner(agent)
+    try:
+        await runner.step({"payload": True})
+    finally:
+        teardown_runner(runner)
+
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    assert [span.name for span in spans] == ["trace-0", "trace-1"]
+
+
+@pytest.mark.asyncio
+async def test_step_accepts_agent_span_list() -> None:
+    class AgentSpanAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> List[Span]:
+            return [
+                create_agent_span(rollout.rollout_id, rollout.attempt.attempt_id, 1, "custom-1", {"order": 1}),
+                create_agent_span(rollout.rollout_id, rollout.attempt.attempt_id, 2, "custom-2", {"order": 2}),
+            ]
+
+    agent = AgentSpanAgent()
+    runner, store, _ = await setup_runner(agent)
+    try:
+        await runner.step({"payload": False})
+    finally:
+        teardown_runner(runner)
+
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    assert [span.name for span in spans] == ["custom-1", "custom-2"]
+
+
+@pytest.mark.asyncio
+async def test_iter_respects_max_tasks() -> None:
+    class CountingAgent(LitAgent[Dict[str, Any]]):
+        def __init__(self) -> None:
+            super().__init__()
+            self.processed: List[int] = []
+
+        def training_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> float:
+            self.processed.append(task["idx"])
+            return 0.0
+
+    agent = CountingAgent()
+    runner, store, _ = await setup_runner(agent, max_tasks=2)
+
+    for idx in range(3):
+        await store.enqueue_rollout({"idx": idx}, mode="train")
+
+    try:
+        await asyncio.wait_for(runner.iter(), timeout=1)
+    finally:
+        teardown_runner(runner)
+
+    assert agent.processed == [0, 1]
+    rollouts = await store.query_rollouts()
+    statuses = {rollout.rollout_id: rollout.status for rollout in rollouts}
+    assert list(statuses.values()).count("succeeded") == 2
+
+
+@pytest.mark.asyncio
+async def test_iter_stops_when_event_is_set() -> None:
+    stop_event = ThreadingEvent()
+
+    class StoppableAgent(LitAgent[Dict[str, Any]]):
+        def __init__(self) -> None:
+            super().__init__()
+            self.processed: List[int] = []
+
+        async def training_rollout_async(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> None:
+            self.processed.append(task["idx"])
+            if len(self.processed) == 1:
+                stop_event.set()
+            await asyncio.sleep(0.05)
+            return None
+
+    agent = StoppableAgent()
+    runner, store, _ = await setup_runner(agent)
+
+    for idx in range(3):
+        await store.enqueue_rollout({"idx": idx}, mode="train")
+
+    iter_task = asyncio.create_task(runner.iter(event=stop_event))
+    try:
+        await asyncio.wait_for(asyncio.to_thread(stop_event.wait, timeout=1), timeout=2)
+        await asyncio.wait_for(iter_task, timeout=1)
+    finally:
+        teardown_runner(runner)
+
+    assert agent.processed == [0]
+    rollouts = await store.query_rollouts()
+    succeeded = [rollout for rollout in rollouts if rollout.status == "succeeded"]
+    assert len(succeeded) == 1
+
+
+@pytest.mark.asyncio
+async def test_iter_waits_when_queue_empty_calls_sleep(monkeypatch: pytest.MonkeyPatch) -> None:
+    stop_event = ThreadingEvent()
+
+    class IdleAgent(LitAgent[Any]):
+        def training_rollout(self, task: Any, resources: Any, rollout: Any) -> None:
+            return None
+
+    agent = IdleAgent()
+    runner, _, _ = await setup_runner(agent, poll_interval=0.01)
+
+    sleep_calls = 0
+
+    async def fake_sleep(event: Optional[Event] = None) -> None:
+        nonlocal sleep_calls
+        sleep_calls += 1
+        if event is not None:
+            event.set()
+
+    monkeypatch.setattr(runner, "_sleep_until_next_poll", fake_sleep)
+
+    try:
+        await runner.iter(event=stop_event)
+    finally:
+        teardown_runner(runner)
+
+    assert sleep_calls >= 1
+
+
+@pytest.mark.asyncio
+async def test_async_validation_rollout_used() -> None:
+    class AsyncValidationAgent(LitAgent[Dict[str, Any]]):
+        def __init__(self) -> None:
+            super().__init__()
+            self.validation_calls = 0
+            self.training_calls = 0
+
+        async def validation_rollout_async(
+            self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
+        ) -> float:
+            self.validation_calls += 1
+            return 0.0
+
+        async def training_rollout_async(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> float:
+            self.training_calls += 1
+            return 0.0
+
+    agent = AsyncValidationAgent()
+    runner, store, _ = await setup_runner(agent, max_tasks=1)
+    await store.enqueue_rollout({"idx": 1}, mode="val")
+
+    try:
+        await runner.iter()
+    finally:
+        teardown_runner(runner)
+
+    assert agent.validation_calls == 1
+    assert agent.training_calls == 0
+
+
+@pytest.mark.asyncio
+async def test_training_rollout_sync_used() -> None:
+    class SyncTrainingAgent(LitAgent[Dict[str, Any]]):
+        def __init__(self) -> None:
+            super().__init__()
+            self.training_calls = 0
+
+        def training_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> None:
+            self.training_calls += 1
+            return None
+
+    agent = SyncTrainingAgent()
+    runner, store, _ = await setup_runner(agent, max_tasks=1)
+    await store.enqueue_rollout({"idx": 99}, mode="train")
+
+    try:
+        await runner.iter()
+    finally:
+        teardown_runner(runner)
+
+    assert agent.training_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_step_handles_agent_exception_marks_attempt_failed() -> None:
+    class FailingAgent(LitAgent[Dict[str, Any]]):
+        def training_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> None:
+            raise RuntimeError("boom")
+
+    agent = FailingAgent()
+    runner, store, _ = await setup_runner(agent)
+    with pytest.raises(RuntimeError):
+        await runner.step({"task": "x"})
+
+    rollouts = await store.query_rollouts()
+    assert len(rollouts) == 1
+    attempts = await store.query_attempts(rollouts[0].rollout_id)
+    assert attempts[-1].status == "failed"
+    teardown_runner(runner)
+
+
+@pytest.mark.asyncio
+async def test_agent_emits_multiple_rewards() -> None:
+    class RewardListAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(
+            self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
+        ) -> List[ReadableSpan]:
+            return [emit_reward(0.2), emit_reward(0.6)]
+
+    agent = RewardListAgent()
+    runner, store, _ = await setup_runner(agent)
+    try:
+        await runner.step({"task": "reward"})
+    finally:
+        teardown_runner(runner)
+
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    reward_values = [span.attributes.get("reward") for span in spans if span.name == SpanNames.REWARD.value]
+    assert reward_values == [0.2, 0.6]
+
+
+@pytest.mark.asyncio
+async def test_hooks_triggered_in_order() -> None:
+    hook = RecordingHook()
+
+    class HookAgent(LitAgent[Dict[str, Any]]):
+        def validation_rollout(
+            self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
+        ) -> List[ReadableSpan]:
+            return [create_readable_span("hook-span")]
+
+    agent = HookAgent()
+    runner, store, _ = await setup_runner(agent, hooks=[hook])
+    try:
+        await runner.step({"task": "hook"})
+    finally:
+        teardown_runner(runner)
+
+    assert hook.calls == ["on_rollout_start", "on_trace_start", "on_trace_end", "on_rollout_end"]
+    assert hook.received_spans is not None
+    rollout_id, attempt_id = await assert_single_attempt_succeeded(store)
+    spans = await store.query_spans(rollout_id, attempt_id)
+    assert [span.name for span in spans] == ["hook-span"]

tests/store/dummy_store.py

@@ -5,7 +5,6 @@ from typing import Any, Dict, List, Literal, Optional, Sequence
 from opentelemetry.sdk.trace import ReadableSpan
 
 from agentlightning.store.base import UNSET, LightningStore
-from agentlightning.tracer import Span
 from agentlightning.types import (
     Attempt,
     AttemptedRollout,
@@ -14,6 +13,7 @@ from agentlightning.types import (
     ResourcesUpdate,
     RolloutStatus,
     RolloutV2,
+    Span,
     TaskInput,
 )
 

tests/store/test_client_server.py

@@ -15,8 +15,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 from agentlightning.store.base import UNSET
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
-from agentlightning.tracer import Span
-from agentlightning.tracer.types import Resource, TraceStatus
+from agentlightning.types import Resource, Span, TraceStatus
 
 
 def _get_free_port() -> int:

tests/store/test_memory.py

@@ -21,13 +21,13 @@ from unittest.mock import Mock
 import pytest
 
 from agentlightning.store.memory import InMemoryLightningStore
-from agentlightning.tracer import Span
 from agentlightning.types import (
     LLM,
     PromptTemplate,
     ResourcesUpdate,
     RolloutConfig,
     RolloutV2,
+    Span,
 )
 
 # Core CRUD Operations Tests

tests/store/test_threading.py

@@ -10,16 +10,18 @@ from opentelemetry.sdk.trace import ReadableSpan
 
 from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.threading import LightningStoreThreaded
-from agentlightning.tracer import Span
-from agentlightning.tracer.types import Resource, SpanContext, TraceStatus
 from agentlightning.types import (
     Attempt,
     AttemptedRollout,
     AttemptStatus,
     NamedResources,
+    Resource,
     ResourcesUpdate,
     RolloutV2,
+    Span,
+    SpanContext,
     TaskInput,
+    TraceStatus,
 )
 
 from .dummy_store import DummyLightningStore

tests/test_litagent_hooks.py

@@ -1,107 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from contextlib import contextmanager
-from typing import Any, Iterator, List, Optional, cast
-
-import pytest
-
-from agentlightning import LitAgent, ResourcesUpdate, Task
-from agentlightning.adapter import TraceTripletAdapter
-from agentlightning.client import AgentLightningClient
-from agentlightning.runner import AgentRunner
-from agentlightning.tracer import BaseTracer
-from agentlightning.types import Rollout
-
-
-class DummyTracer(BaseTracer):
-    @contextmanager
-    def trace_context(self, name: Optional[str] = None) -> Iterator[None]:
-        yield
-
-    def get_last_trace(self) -> List[Any]:
-        return []
-
-
-class DummyClient:
-    def __init__(self) -> None:
-        self.posted = None
-        self.polled = False
-
-    def poll_next_task(self) -> Optional[Task]:
-        if self.polled:
-            return None
-        self.polled = True
-        return Task(rollout_id="1", input={}, mode="train", resources_id=None)
-
-    def get_latest_resources(self) -> ResourcesUpdate:
-        return ResourcesUpdate(resources_id="r", resources={})
-
-    def post_rollout(self, rollout: Any) -> None:
-        self.posted = rollout
-
-
-class DummyAsyncClient:
-    def __init__(self) -> None:
-        self.posted = None
-        self.polled = False
-
-    async def poll_next_task_async(self) -> Optional[Task]:
-        if self.polled:
-            return None
-        self.polled = True
-        return Task(rollout_id="1", input={}, mode="train", resources_id=None)
-
-    async def get_latest_resources_async(self) -> ResourcesUpdate:
-        return ResourcesUpdate(resources_id="r", resources={})
-
-    async def post_rollout_async(self, rollout: Any) -> None:
-        self.posted = rollout
-
-
-class HookAgent(LitAgent[Any]):
-    def __init__(self):
-        super().__init__()
-        self.start_called = False
-        self.end_called = False
-        self.end_rollout: Rollout | None = None
-
-    def training_rollout(self, task: Any, resources: Any, rollout: Any) -> float:
-        return 0.5
-
-    async def training_rollout_async(self, task: Any, resources: Any, rollout: Any) -> float:
-        return 0.5
-
-    def on_rollout_start(self, task: Any, runner: Any, tracer: Any) -> None:
-        self.start_called = True
-        self.start_task = task
-
-    def on_rollout_end(self, task: Any, rollout: Any, runner: Any, tracer: Any) -> None:
-        self.end_called = True
-        self.end_rollout = rollout
-
-
-def test_runner_calls_hooks():
-    agent = HookAgent()
-    client = DummyClient()
-    tracer = DummyTracer()
-    runner = AgentRunner(agent, cast(AgentLightningClient, client), tracer, TraceTripletAdapter())
-
-    assert runner.run() is True
-    assert agent.start_called
-    assert agent.end_called
-    assert agent.end_rollout is not None
-    assert agent.end_rollout.final_reward == 0.5
-
-
-@pytest.mark.asyncio
-async def test_runner_calls_hooks_async():
-    agent = HookAgent()
-    client = DummyAsyncClient()
-    tracer = DummyTracer()
-    runner = AgentRunner(agent, cast(AgentLightningClient, client), tracer, TraceTripletAdapter())
-
-    assert await runner.run_async() is True
-    assert agent.start_called
-    assert agent.end_called
-    assert agent.end_rollout is not None
-    assert agent.end_rollout.final_reward == 0.5

tests/test_llm_proxy.py

@@ -16,26 +16,21 @@ There are some specific TODOs for each test function.
 import ast
 import asyncio
 import json
-import os
 import random
-import socket
-import subprocess
-import sys
-import time
-from contextlib import closing
-from typing import Any, List, Optional, cast
+from typing import Any, List, cast
 
 import anthropic
-import httpx
 import openai
 import pytest
 from opentelemetry.sdk.trace import ReadableSpan
 
 from agentlightning.llm_proxy import LightningSpanExporter, LLMProxy
 from agentlightning.store.memory import InMemoryLightningStore
-from agentlightning.tracer.types import Span
-from agentlightning.types import LLM
-from tests.tracer.utils import clear_tracer_provider
+from agentlightning.types import LLM, Span
+
+from .common.network import get_free_port
+from .common.tracer import clear_tracer_provider
+from .common.vllm import VLLM_VERSION, RemoteOpenAIServer
 
 try:
     import torch  # type: ignore
@@ -45,183 +40,6 @@ except Exception:
     GPU_AVAILABLE = False  # type: ignore
     pytest.skip(reason="GPU not available", allow_module_level=True)
 
-VLLM_AVAILABLE = False
-VLLM_UNAVAILABLE_REASON = ""
-
-try:
-    import vllm
-    from vllm.engine.arg_utils import AsyncEngineArgs
-    from vllm.entrypoints.cli.serve import ServeSubcommand
-    from vllm.model_executor.model_loader import get_model_loader
-    from vllm.utils import FlexibleArgumentParser
-
-    VLLM_AVAILABLE = True  # type: ignore
-    VLLM_VERSION = tuple(int(v) for v in vllm.__version__.split("."))
-except ImportError as e:
-    AsyncEngineArgs = None
-    get_model_loader = None
-    FlexibleArgumentParser = None
-    ServeSubcommand = None
-    VLLM_VERSION = (0, 0, 0)  # type: ignore
-    VLLM_UNAVAILABLE_REASON = str(e)  # type: ignore
-
-
-def _get_free_port() -> int:
-    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
-        s.bind(("127.0.0.1", 0))
-        return int(s.getsockname()[1])
-
-
-class RemoteOpenAIServer:
-    """
-    A context manager for launching and interacting with a remote vLLM-based
-    OpenAI-compatible server instance.
-
-    This class handles:
-      - Preparing the environment and spawning the vLLM server process
-      - Ensuring that the requested model is downloaded before server startup
-      - Polling and health-checking the server until it is ready
-      - Providing helper methods to construct URLs for API calls
-      - Returning configured synchronous and asynchronous OpenAI clients
-        that can communicate with the launched server
-
-    Typical usage:
-        with RemoteOpenAIServer(vllm_serve_args, port, model) as server:
-            client = server.get_client()
-            response = client.chat.completions.create(...)
-
-    Attributes:
-        DUMMY_API_KEY (str): A placeholder API key for compatibility
-                             (vLLM does not require authentication).
-        host (str): Host address of the server (default: "localhost").
-        port (int): TCP port number for the server.
-        proc (subprocess.Popen): Handle to the launched server process.
-    """
-
-    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
-
-    def _start_server(self, model: str, vllm_serve_args: list[str], env_dict: Optional[dict[str, str]]) -> None:
-        """Subclasses override this method to customize server process launch"""
-        env = os.environ.copy()
-        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"  # safer CUDA init
-        if env_dict is not None:
-            env.update(env_dict)
-
-        if VLLM_VERSION >= (0, 10, 2):
-            # Supports return_token_ids
-            self.proc: subprocess.Popen[bytes] = subprocess.Popen(
-                ["vllm", "serve", model, *vllm_serve_args],
-                env=env,
-                stdout=sys.stdout,
-                stderr=sys.stderr,
-            )
-        else:
-            # Does not support return_token_ids
-            self.proc = subprocess.Popen(
-                ["python", "-m", "agentlightning.cli.vllm", "serve", model, *vllm_serve_args],
-                env=env,
-                stdout=sys.stdout,
-                stderr=sys.stderr,
-            )
-
-    def __init__(
-        self,
-        model: str,
-        vllm_serve_args: list[str],  # should not include the model name
-        env_dict: Optional[dict[str, str]] = None,
-        seed: Optional[int] = 0,
-        max_wait_seconds: Optional[float] = None,
-    ) -> None:
-        if (
-            not VLLM_AVAILABLE
-            or AsyncEngineArgs is None
-            or get_model_loader is None
-            or FlexibleArgumentParser is None
-            or ServeSubcommand is None
-        ):
-            raise ImportError("vLLM is not available: " + VLLM_UNAVAILABLE_REASON)
-
-        self.model = model
-
-        parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
-        subparsers = parser.add_subparsers(required=False, dest="subparser")
-        parser = ServeSubcommand().subparser_init(subparsers)  # pyright: ignore[reportUnknownMemberType]
-        args = parser.parse_args(["--model", model, *vllm_serve_args])
-        assert args is not None
-        self.host = str(args.host or "localhost")
-        self.port = int(args.port)
-
-        # download the model before starting the server to avoid timeout
-        is_local = os.path.isdir(model)
-        if not is_local:
-            engine_args = AsyncEngineArgs.from_cli_args(args)
-            model_config = engine_args.create_model_config()
-            load_config = engine_args.create_load_config()
-
-            model_loader = get_model_loader(load_config)
-            model_loader.download_model(model_config)
-
-        self._start_server(model, vllm_serve_args, env_dict)
-        max_wait_seconds = max_wait_seconds or 240
-        self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
-        self.proc.terminate()
-        try:
-            self.proc.wait(8)
-        except subprocess.TimeoutExpired:
-            self.proc.kill()
-
-    def _poll(self) -> Optional[int]:
-        """Subclasses override this method to customize process polling"""
-        return self.proc.poll()
-
-    def _wait_for_server(self, *, url: str, timeout: float):
-        start = time.time()
-        client = httpx.Client()
-
-        while True:
-            try:
-                if client.get(url).status_code == 200:
-                    break
-            except Exception:
-                result = self._poll()
-                if result is not None and result != 0:
-                    raise RuntimeError("Server exited unexpectedly.") from None
-                time.sleep(0.5)
-                if time.time() - start > timeout:
-                    raise RuntimeError("Server failed to start in time.") from None
-
-    @property
-    def url_root(self) -> str:
-        return f"http://{self.host}:{self.port}"
-
-    def url_for(self, *parts: str) -> str:
-        return self.url_root + "/" + "/".join(parts)
-
-    def get_client(self, **kwargs: Any):
-        if "timeout" not in kwargs:
-            kwargs["timeout"] = 600
-        return openai.OpenAI(
-            base_url=self.url_for("v1"),
-            api_key=self.DUMMY_API_KEY,
-            max_retries=0,
-            **kwargs,
-        )
-
-    def get_async_client(self, **kwargs: Any):
-        if "timeout" not in kwargs:
-            kwargs["timeout"] = 600
-        return openai.AsyncOpenAI(
-            base_url=self.url_for("v1"),
-            api_key=self.DUMMY_API_KEY,
-            max_retries=0,
-            **kwargs,
-        )
-
 
 @pytest.fixture(scope="module")
 def qwen25_model():
@@ -234,7 +52,7 @@ def qwen25_model():
             "--tool-call-parser",
             "hermes",
             "--port",
-            str(_get_free_port()),
+            str(get_free_port()),
         ],
     ) as server:
         yield server
@@ -260,7 +78,7 @@ def setup_module():
 async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
     store = InMemoryLightningStore()
     proxy = LLMProxy(
-        port=_get_free_port(),
+        port=get_free_port(),
         model_list=[
             {
                 "model_name": "gpt-4o-arbitrary",
@@ -361,7 +179,7 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
 def _make_proxy_and_store(qwen25_model: RemoteOpenAIServer, *, retries: int = 0):
     store = InMemoryLightningStore()
     proxy = LLMProxy(
-        port=_get_free_port(),
+        port=get_free_port(),
         model_list=[
             {
                 "model_name": "gpt-4o-arbitrary",
@@ -542,6 +360,7 @@ async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
         proxy.stop()
 
 
+@pytest.mark.skip(reason="Streaming is not supported yet")
 @pytest.mark.asyncio
 async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
     proxy, store = _make_proxy_and_store(qwen25_model)

tests/test_otel.py

@@ -0,0 +1,406 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# pyright: reportPrivateUsage=false
+
+"""Unit tests for LightningSpanProcessor."""
+
+import asyncio
+import multiprocessing
+import pickle
+import threading
+import time
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import SpanContext, TraceFlags
+
+from agentlightning.store.base import LightningStore
+from agentlightning.tracer.agentops import LightningSpanProcessor
+
+
+def create_span(name: str, sampled: bool = True, with_context: bool = True) -> MagicMock:
+    """Helper to create mock spans with different properties."""
+    span = MagicMock(spec=ReadableSpan)
+    span.name = name
+    if with_context:
+        span.context = SpanContext(
+            trace_id=hash(name) % (2**64),
+            span_id=hash(name) % (2**64),
+            is_remote=False,
+            trace_flags=TraceFlags(0x01 if sampled else 0x00),
+        )
+    else:
+        span.context = None
+    return span
+
+
+def create_mock_store() -> MagicMock:
+    """Helper to create a mock LightningStore."""
+    store = MagicMock(spec=LightningStore)
+    store.add_otel_span = AsyncMock(return_value=None)
+    return store
+
+
+def test_initialization_and_shutdown():
+    """Test processor lifecycle: initialization, loop thread, and shutdown."""
+    processor = LightningSpanProcessor()
+
+    # Verify initialization
+    assert processor._spans == []
+    assert processor._store is None
+    assert processor._rollout_id is None
+    assert processor._attempt_id is None
+
+    # Verify loop thread is running correctly
+    assert processor._loop is not None
+    assert processor._loop.is_running()
+    assert processor._loop_thread.is_alive()
+    assert processor._loop_thread.daemon is True
+    assert processor._loop_thread.name == "otel-loop"
+
+    # Verify shutdown stops everything
+    thread = processor._loop_thread
+    processor.shutdown()
+
+    time.sleep(0.1)  # Give thread time to stop
+    assert not thread.is_alive()
+    assert processor._loop is None
+
+    # Verify double shutdown is safe (idempotent)
+    processor.shutdown()  # Should not raise
+    assert processor._loop is None
+
+
+def test_span_collection_with_filtering():
+    """Test that spans are collected with proper filtering."""
+    processor = LightningSpanProcessor()
+
+    sampled_span = create_span("sampled", sampled=True)
+    unsampled_span = create_span("unsampled", sampled=False)
+    no_context_span = create_span("no_context", with_context=False)
+
+    # Process different types of spans
+    processor.on_end(sampled_span)
+    processor.on_end(unsampled_span)
+    processor.on_end(no_context_span)
+
+    # Only sampled span with context should be collected
+    collected = processor.spans()
+    assert len(collected) == 1
+    assert collected[0] == sampled_span
+
+    processor.shutdown()
+
+
+def test_context_managers_clear_state():
+    """Test that both context managers properly manage state."""
+    processor = LightningSpanProcessor()
+    store = create_mock_store()
+
+    # Add a span first
+    span1 = create_span("span1")
+    processor.on_end(span1)
+    assert len(processor.spans()) == 1
+
+    # Test basic context manager: __enter__ clears spans
+    with processor:
+        assert len(processor.spans()) == 0
+        span2 = create_span("span2")
+        processor.on_end(span2)
+        assert len(processor.spans()) == 1
+
+    # After exit, store context should be cleared
+    assert processor._store is None
+    assert processor._rollout_id is None
+    assert processor._attempt_id is None
+
+    # Test with_context: sets and clears store context
+    with processor.with_context(store=store, rollout_id="r1", attempt_id="a1") as proc:
+        assert proc is processor
+        assert processor._store is store
+        assert processor._rollout_id == "r1"
+        assert processor._attempt_id == "a1"
+
+    # After exit, context should be cleared
+    assert processor._store is None
+    assert processor._rollout_id is None
+    assert processor._attempt_id is None
+
+    processor.shutdown()
+
+
+def test_store_integration_complete():
+    """Test all store integration scenarios: writes, errors, timeout, thread verification."""
+    processor = LightningSpanProcessor()
+    store = create_mock_store()
+
+    # Test 1: Successful store writes
+    with processor.with_context(store=store, rollout_id="r1", attempt_id="a1"):
+        span1 = create_span("span1")
+        span2 = create_span("span2")
+        processor.on_end(span1)
+        processor.on_end(span2)
+
+    # Verify both spans written to store
+    assert store.add_otel_span.call_count == 2
+    assert len(processor.spans()) == 2
+
+    # Verify call arguments
+    calls = store.add_otel_span.call_args_list
+    assert calls[0][0] == ("r1", "a1", span1)
+    assert calls[1][0] == ("r1", "a1", span2)
+
+    # Test 2: Store write errors are caught and don't crash
+    store.add_otel_span.reset_mock()
+    store.add_otel_span.side_effect = RuntimeError("Store failure")
+
+    with processor.with_context(store=store, rollout_id="r2", attempt_id="a2"):
+        span3 = create_span("span3")
+        processor.on_end(span3)  # Should not raise
+
+    # Span still collected despite store error
+    assert len(processor.spans()) == 1
+    store.add_otel_span.side_effect = None
+
+    # Test 3: Verify writes happen in loop thread
+    execution_thread = None
+
+    async def track_thread(*args: Any, **kwargs: Any):
+        nonlocal execution_thread
+        execution_thread = threading.current_thread()
+
+    store.add_otel_span = AsyncMock(side_effect=track_thread)
+
+    with processor.with_context(store=store, rollout_id="r3", attempt_id="a3"):
+        span4 = create_span("span4")
+        processor.on_end(span4)
+
+    assert execution_thread is not None
+    assert execution_thread.name == "otel-loop"
+
+    # Test 4: Unsampled and no-context spans don't trigger store writes
+    store.add_otel_span.reset_mock()
+    store.add_otel_span = AsyncMock()
+
+    with processor.with_context(store=store, rollout_id="r4", attempt_id="a4"):
+        unsampled = create_span("unsampled", sampled=False)
+        no_ctx = create_span("no_ctx", with_context=False)
+        processor.on_end(unsampled)
+        processor.on_end(no_ctx)
+
+    # Store should not be called for filtered spans
+    store.add_otel_span.assert_not_called()
+
+    processor.shutdown()
+
+
+def test_event_loop_operations():
+    """Test _await_in_loop: execution, exceptions, and timeout."""
+    processor = LightningSpanProcessor()
+
+    # Test 1: Successful coroutine execution
+    async def successful_coro():
+        await asyncio.sleep(0.01)
+        return "success"
+
+    result = processor._await_in_loop(successful_coro())
+    assert result == "success"
+
+    # Test 2: Exception propagation
+    async def failing_coro():
+        raise ValueError("Expected error")
+
+    with pytest.raises(ValueError, match="Expected error"):
+        processor._await_in_loop(failing_coro())
+
+    # Test 3: Timeout handling
+    async def slow_coro():
+        await asyncio.sleep(10)
+        return "too_slow"
+
+    with pytest.raises(Exception):  # Should timeout
+        processor._await_in_loop(slow_coro(), timeout=0.1)
+
+    processor.shutdown()
+
+
+def test_concurrent_access():
+    """Test thread-safe concurrent span processing."""
+    processor = LightningSpanProcessor()
+    store = create_mock_store()
+
+    num_threads = 10
+    spans_per_thread = 5
+    barrier = threading.Barrier(num_threads)
+
+    def process_spans(thread_id: int) -> None:
+        barrier.wait()  # Synchronize all threads to start together
+        for i in range(spans_per_thread):
+            span = create_span(f"thread{thread_id}_span{i}")
+            processor.on_end(span)
+
+    with processor.with_context(store=store, rollout_id="concurrent", attempt_id="test"):
+        threads = [threading.Thread(target=process_spans, args=(i,)) for i in range(num_threads)]
+
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+    # All spans should be collected and written
+    assert len(processor.spans()) == num_threads * spans_per_thread
+    assert store.add_otel_span.call_count == num_threads * spans_per_thread
+
+    processor.shutdown()
+
+
+def test_multiprocessing_behavior():
+    """Test processor behavior across process boundaries."""
+
+    # Test 1: Creating new processor in subprocess works
+    def subprocess_task(result_queue: "multiprocessing.Queue[tuple[str, Any]]") -> None:
+        try:
+            processor = LightningSpanProcessor()
+
+            # Verify processor works in new process
+            assert processor._loop is not None
+            assert processor._loop_thread.is_alive()
+
+            span = create_span("subprocess_span")
+            processor.on_end(span)
+
+            result_queue.put(("success", len(processor.spans())))
+            processor.shutdown()
+        except Exception as e:
+            result_queue.put(("error", str(e)))
+
+    result_queue: multiprocessing.Queue[tuple[str, Any]] = multiprocessing.Queue()
+    process = multiprocessing.Process(target=subprocess_task, args=(result_queue,))
+    process.start()
+    process.join(timeout=5)
+
+    assert not process.is_alive()
+    status, value = result_queue.get(timeout=1)
+    assert status == "success"
+    assert value == 1
+
+    # Test 2: Processor cannot be pickled (threads aren't picklable)
+    processor = LightningSpanProcessor()
+
+    with pytest.raises((TypeError, AttributeError)):
+        pickle.dumps(processor)
+
+    processor.shutdown()
+
+
+def test_edge_cases():
+    """Test edge cases and error conditions."""
+    processor = LightningSpanProcessor()
+
+    # Test 1: force_flush always returns True
+    assert processor.force_flush() is True
+    assert processor.force_flush(timeout_millis=5000) is True
+
+    # Test 2: Calling _await_in_loop after shutdown raises
+    processor.shutdown()
+
+    async def dummy_coro():
+        return "test"
+
+    with pytest.raises(RuntimeError, match="Loop is not initialized"):
+        processor._await_in_loop(dummy_coro())
+
+    # Test 3: Verify shutdown thread join timeout is respected
+    processor2 = LightningSpanProcessor()
+
+    # Mock thread.join to verify timeout parameter
+    original_join = processor2._loop_thread.join
+    join_timeout: float | None = None
+
+    def mock_join(timeout: float | None = None) -> None:
+        nonlocal join_timeout
+        join_timeout = timeout
+        return original_join(timeout=timeout)
+
+    processor2._loop_thread.join = mock_join
+    processor2.shutdown()
+
+    assert join_timeout == 5  # Should pass 5 second timeout
+
+
+def test_store_write_timeout():
+    """Test that slow store writes respect timeout."""
+    processor = LightningSpanProcessor()
+    store = create_mock_store()
+
+    # Create a slow async function that exceeds timeout
+    async def slow_write(*args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(10)
+
+    store.add_otel_span = AsyncMock(side_effect=slow_write)
+
+    with processor.with_context(store=store, rollout_id="r1", attempt_id="a1"):
+        span = create_span("test_span")
+        # Should not raise - timeout is caught in on_end
+        processor.on_end(span)
+
+    # Span still collected despite timeout
+    # Note: The timeout in on_end is 5.0 seconds, but the error is caught
+    assert len(processor.spans()) == 1
+
+    processor.shutdown()
+
+
+def test_multiple_processors_in_same_process():
+    """Test that multiple processors can coexist in the same process."""
+    processor1 = LightningSpanProcessor()
+    processor2 = LightningSpanProcessor()
+
+    # Both should have independent loops and threads
+    assert processor1._loop is not processor2._loop
+    assert processor1._loop_thread is not processor2._loop_thread
+    assert processor1._loop is not None and processor1._loop.is_running()
+    assert processor2._loop is not None and processor2._loop.is_running()
+
+    # Both should work independently
+    span1 = create_span("p1_span")
+    span2 = create_span("p2_span")
+
+    processor1.on_end(span1)
+    processor2.on_end(span2)
+
+    assert len(processor1.spans()) == 1
+    assert len(processor2.spans()) == 1
+    assert processor1.spans()[0] == span1
+    assert processor2.spans()[0] == span2
+
+    processor1.shutdown()
+    processor2.shutdown()
+
+
+def test_context_manager_reusability():
+    """Test that context managers can be entered and exited multiple times."""
+    processor = LightningSpanProcessor()
+    store = create_mock_store()
+
+    # First usage
+    with processor.with_context(store=store, rollout_id="r1", attempt_id="a1"):
+        span1 = create_span("span1")
+        processor.on_end(span1)
+
+    assert len(processor.spans()) == 1
+    assert processor._store is None
+
+    # Second usage - should work fine
+    with processor.with_context(store=store, rollout_id="r2", attempt_id="a2"):
+        span2 = create_span("span2")
+        processor.on_end(span2)
+
+    # Spans should be cleared from first context, only span2 present
+    assert len(processor.spans()) == 1
+    assert processor.spans()[0].name == "span2"
+    assert processor._store is None
+
+    processor.shutdown()

tests/test_tracer.py

@@ -22,6 +22,7 @@ import asyncio
 import difflib
 import inspect
 import json
+import multiprocessing
 import os
 import pprint
 import re
@@ -68,7 +69,7 @@ from agentlightning.tracer.agentops import AgentOpsTracer, LightningSpanProcesso
 from agentlightning.tracer.http import HttpTracer
 from agentlightning.types import Triplet
 
-from .tracer.utils import clear_tracer_provider
+from .common.tracer import clear_agentops_init, clear_tracer_provider
 
 USE_OPENAI = os.environ.get("USE_OPENAI", "false").lower() == "true"
 if USE_OPENAI:
@@ -785,14 +786,33 @@ def create_prompt_caches() -> None:
         run_all()
 
 
-@pytest.fixture(scope="module", autouse=True)
-def setup_module():
-    clear_tracer_provider()
-    yield
+@pytest.mark.parametrize("agent_func_name", [f.__name__ for f in iterate_over_agents()], ids=str)
+def test_run_with_agentops_tracer(agent_func_name: str):
+    """AgentOps tracer tests are notoriously problematic and does not work well with other tests."""
+    if agent_func_name in ["openai_agents_sdk_mcp_tool_use", "agent_autogen_mcp"]:
+        pytest.skip("Async MCP server is problematic with AgentOps tracer in multiprocessing mode.")
+
+    ctx = multiprocessing.get_context("spawn")
+    proc = ctx.Process(target=_test_run_with_agentops_tracer_impl, args=(agent_func_name,))
+    proc.start()
+    proc.join(10)
+
+    if proc.is_alive():
+        proc.terminate()
+        proc.join(5)
+        if proc.is_alive():
+            proc.kill()
+
+        assert False, "Child process hung. Check test output for details."
+
+    assert proc.exitcode == 0, (
+        f"Child process for {agent_func_name!r} failed with exit code {proc.exitcode}. "
+        "Check child traceback in test output."
+    )
 
 
-@pytest.mark.parametrize("agent_func", list(iterate_over_agents()), ids=lambda f: f.__name__)
-def test_run_with_agentops_tracer(agent_func):
+def _test_run_with_agentops_tracer_impl(agent_func_name: str):
+    agent_func = next(f for f in iterate_over_agents() if f.__name__ == agent_func_name)
     tracer = AgentOpsTracer()
     tracer.init()
     tracer.init_worker(0)
@@ -800,9 +820,7 @@ def test_run_with_agentops_tracer(agent_func):
     global _langchain_callback_handler
     _langchain_callback_handler = tracer.get_langchain_callback_handler()
 
-    loop = asyncio.new_event_loop()
     try:
-        asyncio.set_event_loop(loop)
         tracer.trace_run(
             run_one,
             agent_func,
@@ -841,8 +859,6 @@ def test_run_with_agentops_tracer(agent_func):
     finally:
         tracer.teardown_worker(0)
         tracer.teardown()
-        loop.close()
-        asyncio.set_event_loop(None)
 
 
 @pytest.mark.parametrize("agent_func", list(iterate_over_agents()), ids=lambda f: f.__name__)