# Copyright (c) Microsoft. All rights reserved. # pyright: reportPrivateUsage=false from __future__ import annotations import asyncio import multiprocessing import pickle import threading import time from multiprocessing.connection import Connection from typing import Any, Dict, List, Optional from unittest.mock import AsyncMock, MagicMock import pytest import uvicorn from fastapi import FastAPI, Request from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( ExportTraceServiceRequest, ExportTraceServiceResponse, ) from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.trace import SpanContext, TraceFlags from portpicker import pick_unused_port from agentlightning.reward import emit_reward, find_reward_spans, get_reward_value, reward from agentlightning.store.base import LightningStore from agentlightning.tracer.agentops import LightningSpanProcessor from agentlightning.tracer.otel import OtelTracer from agentlightning.types import TraceStatus from agentlightning.types.tracer import Span from agentlightning.utils import otlp from ..common.tracer import clear_tracer_provider def create_span(name: str, sampled: bool = True, with_context: bool = True) -> ReadableSpan: """Helper to create mock spans with different properties.""" return ReadableSpan( name=name, 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), ) if with_context else None ), ) async def add_otel_span( rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None ) -> Span: span = Span.from_opentelemetry( readable_span, rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id or 0 ) return span def create_mock_store(otlp_supported: bool = False) -> MagicMock: """Helper to create a mock LightningStore.""" store = MagicMock(spec=LightningStore) store.add_otel_span = AsyncMock(side_effect=add_otel_span) store.capabilities = {"otlp_traces": otlp_supported} store.otlp_traces_endpoint.return_value = "http://store/v1/traces" return store @pytest.fixture(params=[False, True], ids=["store-no-otlp", "store-otlp"]) def store_supports_otlp(request: pytest.FixtureRequest) -> bool: return bool(request.param) @pytest.fixture def otlp_server(store_supports_otlp: bool): if not store_supports_otlp: yield None return app = FastAPI() received: List[ExportTraceServiceRequest] = [] @app.post("/v1/traces") async def _export_traces(request: Request): # type: ignore async def capture(message: ExportTraceServiceRequest) -> None: received.append(message) return await otlp.handle_otlp_export( request, ExportTraceServiceRequest, ExportTraceServiceResponse, capture, signal_name="traces", ) port = pick_unused_port() config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error") server = uvicorn.Server(config) thread = threading.Thread(target=server.run, daemon=True) thread.start() timeout = time.time() + 5 while not getattr(server, "started", False): if time.time() > timeout: raise RuntimeError("OTLP test server failed to start") if not thread.is_alive(): raise RuntimeError("OTLP test server thread exited before startup") time.sleep(0.01) try: yield {"url": f"http://127.0.0.1:{port}/v1/traces", "received": received} finally: server.should_exit = True thread.join(timeout=5) @pytest.fixture def store(store_supports_otlp: bool, otlp_server: Optional[Dict[str, Any]]) -> MagicMock: mock_store = create_mock_store(store_supports_otlp) if store_supports_otlp: assert otlp_server is not None mock_store.otlp_traces_endpoint.return_value = otlp_server["url"] return mock_store @pytest.fixture def otel_tracer(): clear_tracer_provider() tracer = OtelTracer() with tracer.lifespan(): yield tracer clear_tracer_provider() def test_otel_tracer_create_span_records_fields(otel_tracer: OtelTracer) -> None: timestamp = 2000.123 status = TraceStatus(status_code="ERROR", description="boom") span = otel_tracer.create_span( "otel-span", attributes={"foo": "bar"}, timestamp=timestamp, status=status, ) assert span.name == "otel-span" assert span.attributes["foo"] == "bar" assert span.start_time == pytest.approx(timestamp) # type: ignore assert span.end_time == pytest.approx(timestamp) # type: ignore assert span.status.status_code == "ERROR" assert span.status.description == "boom" def test_otel_tracer_operation_context_records_span(otel_tracer: OtelTracer) -> None: with pytest.raises(RuntimeError): with otel_tracer.operation_context("otel-operation", attributes={"initial": "value"}) as ctx: ctx.record_attributes({"custom": "attr"}) raise RuntimeError("otel failure") recorded_span = ctx.get_recorded_span() # type: ignore assert recorded_span.name == "otel-operation" assert recorded_span.attributes["initial"] == "value" assert recorded_span.attributes["custom"] == "attr" assert recorded_span.status.status_code == "ERROR" assert recorded_span.status.description is not None assert "otel failure" in recorded_span.status.description assert recorded_span.start_time is not None assert recorded_span.end_time is not None 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 assert processor._loop is None assert processor._loop_thread is None # Start the loop processor._ensure_loop() # Verify loop thread is running correctly assert processor._loop is not None assert processor._loop.is_running() assert processor._loop_thread is not None 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].name == "sampled" processor.shutdown() def test_context_managers_clear_state(store: MagicMock): """Test that both context managers properly manage state.""" processor = LightningSpanProcessor() # 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(store: MagicMock): """Test all store integration scenarios: writes, errors, timeout, thread verification.""" processor = LightningSpanProcessor() # 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(store: MagicMock): """Test thread-safe concurrent span processing.""" processor = LightningSpanProcessor() 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(store_supports_otlp: bool): """Test processor behavior across process boundaries.""" # Test 1: Creating new processor in subprocess works def subprocess_task(result_queue: "multiprocessing.Queue[tuple[str, Any]]", supports_otlp: bool) -> None: try: processor = LightningSpanProcessor() span = create_span("subprocess_span") processor.on_end(span) # Loop is not used assert processor._loop is None assert processor._loop_thread is None with processor.with_context(store=create_mock_store(supports_otlp), rollout_id="r1", attempt_id="a1"): processor.on_end(span) # Verify processor works in new process assert processor._loop is not None assert processor._loop_thread is not None assert processor._loop_thread.is_alive() 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, store_supports_otlp)) process.start() process.join(timeout=5) assert not process.is_alive() status, value = result_queue.get(timeout=1) assert status == "success", f"Subprocess failed: {status}, {value}" 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_store_write_timeout(store: MagicMock): """Test that slow store writes respect timeout.""" processor = LightningSpanProcessor() # Create a slow async function that exceeds timeout async def slow_write(*args: Any, **kwargs: Any) -> None: await asyncio.sleep(15) 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() processor1._ensure_loop() processor2._ensure_loop() # 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].name == "p1_span" assert processor2.spans()[0].name == "p2_span" processor1.shutdown() processor2.shutdown() def test_context_manager_reusability(store: MagicMock): """Test that context managers can be entered and exited multiple times.""" processor = LightningSpanProcessor() # 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() def _otel_reward_subprocess(mode: str, conn: Connection[tuple[str, Any]]) -> None: asyncio.run(_otel_reward_subprocess_async(mode, conn)) async def _otel_reward_subprocess_async(mode: str, conn: Connection[tuple[str, Any]]) -> None: tracer: OtelTracer | None = None try: clear_tracer_provider() tracer = OtelTracer() tracer.init_worker(0) if mode == "decorator": expected_reward = 2.5 @reward def compute_reward() -> float: return expected_reward async with tracer.trace_context(name="reward-decorator"): returned = compute_reward() if returned != expected_reward: raise AssertionError(f"Expected reward {expected_reward}, got {returned}") elif mode == "emit": expected_reward = 4.5 async with tracer.trace_context(name="reward-emit"): emit_reward(expected_reward) else: raise ValueError(f"Unsupported mode: {mode}") spans = tracer.get_last_trace() total_spans = len(spans) reward_spans = find_reward_spans(spans) if not reward_spans: raise AssertionError("No reward spans captured by OtelTracer.") actual_reward = get_reward_value(reward_spans[-1]) if actual_reward is None: raise AssertionError("Captured reward span does not contain a reward value.") conn.send( ( "success", { "mode": mode, "reward": actual_reward, "expected_reward": expected_reward, "total_spans": total_spans, "reward_spans": len(reward_spans), }, ) ) except Exception as exc: # pragma: no cover - defensive path try: conn.send(("error", repr(exc))) except Exception: pass finally: if tracer is not None: tracer.teardown_worker(0) conn.close() def _run_otel_reward_test(mode: str) -> dict[str, Any]: ctx = multiprocessing.get_context("spawn") parent_conn, child_conn = ctx.Pipe(duplex=False) proc = ctx.Process(target=_otel_reward_subprocess, args=(mode, child_conn)) proc.start() proc.join(15.0) if proc.is_alive(): proc.terminate() proc.join(5) pytest.fail(f"Subprocess for mode {mode!r} hung.") try: if not parent_conn.poll(1.0): pytest.fail(f"No result from subprocess for mode {mode!r}: timed out") status, payload = parent_conn.recv() finally: parent_conn.close() if status != "success": pytest.fail(f"Subprocess for mode {mode!r} failed: {payload}") assert proc.exitcode == 0, f"Subprocess for mode {mode!r} exited with code {proc.exitcode}" return payload def test_otel_tracer_captures_reward_from_decorator(): payload = _run_otel_reward_test("decorator") assert payload["reward_spans"] == 1 assert payload["total_spans"] == 1 assert payload["reward"] == pytest.approx(payload["expected_reward"]) # pyright: ignore[reportUnknownMemberType] assert payload["expected_reward"] == pytest.approx(2.5) # pyright: ignore[reportUnknownMemberType] def test_otel_tracer_captures_reward_from_emit_reward(): payload = _run_otel_reward_test("emit") assert payload["reward_spans"] == 1 assert payload["total_spans"] == 1 assert payload["reward"] == pytest.approx(payload["expected_reward"]) # pyright: ignore[reportUnknownMemberType] assert payload["expected_reward"] == pytest.approx(4.5) # pyright: ignore[reportUnknownMemberType]