chore: import upstream snapshot with attribution

tests/utils/__init__.py

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

tests/utils/test_metrics.py

@@ -0,0 +1,433 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+import logging
+import sys
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import pytest
+
+import agentlightning.utils.metrics as metrics_module
+from agentlightning.utils.metrics import ConsoleMetricsBackend, MetricsBackend, MultiMetricsBackend
+
+from ..common.prometheus_stub import make_prometheus_stub
+
+pytestmark = pytest.mark.utils
+
+
+def test_validate_labels_reports_missing_label_with_metric_name() -> None:
+    labels = {"method": "GET", "status": "200"}
+    result = metrics_module._validate_labels("counter", "requests", labels, ("method", "status"))
+    assert result == (("method", "GET"), ("status", "200"))
+
+    with pytest.raises(ValueError) as excinfo:
+        metrics_module._validate_labels("counter", "requests", {"method": "POST"}, ("method", "status"))
+
+    message = str(excinfo.value)
+    assert "Counter 'requests'" in message
+    assert "'status' is required" in message
+
+
+def test_normalize_label_names_preserves_original_order() -> None:
+    assert metrics_module._normalize_label_names(["b", "a", "b"]) == ("b", "a", "b")
+    assert metrics_module._normalize_label_names(None) == ()
+
+
+def test_console_backend_logs_counters_with_registration_order() -> None:
+    backend = ConsoleMetricsBackend(log_interval_seconds=0.0)
+    line = backend._log_counter(
+        "rate",
+        {"group1": "a", "group2": "b"},
+        timestamps=[0.0, 1.0],
+        amounts=[1.0, 1.0],
+        snapshot_time=5.0,
+    )
+
+    assert line == "rate{group1=a,group2=b}=0.40/s"
+
+
+def test_console_backend_logs_histograms_with_human_units() -> None:
+    backend = ConsoleMetricsBackend()
+    line = backend._log_histogram(
+        "latency",
+        {"group1": "a", "group2": "b"},
+        values=[0.00395, 0.0168, 3.5],
+        buckets=(0.5,),
+        snapshot_time=1.0,
+    )
+
+    assert line is not None and line.startswith("latency{group1=a,group2=b}=")
+    payload = line.rsplit("=", 1)[1]
+    p50, p95, p99 = payload.split(",", 2)
+    assert p50.endswith("ms")
+    assert p95.endswith("s")
+    assert p99.endswith("s")
+
+
+def test_console_backend_respects_group_level_limit():
+    backend = ConsoleMetricsBackend(group_level=2, log_interval_seconds=0.0)
+    truncated = backend._truncate_labels_for_logging({"group1": "a", "group2": "b", "group3": "c"}, backend.group_level)
+    line = backend._log_counter("metric", truncated, [0.0, 2.0], [1.0, 1.0], snapshot_time=5.0)
+
+    assert line == "metric{group1=a,group2=b}=0.40/s"
+
+
+def test_console_backend_log_uses_logger(caplog: pytest.LogCaptureFixture) -> None:
+    backend = ConsoleMetricsBackend()
+    caplog.set_level(logging.INFO, logger="agentlightning.utils.metrics")
+    backend._log("hello metrics")
+    assert any(record.message == "hello metrics" for record in caplog.records)
+
+
+class _RecordingBackend(MetricsBackend):
+    def __init__(self) -> None:
+        self.calls: List[Tuple[str, Tuple[Any, ...]]] = []
+
+    def register_counter(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        group_level: Optional[int] = None,
+    ) -> None:
+        self.calls.append(("register_counter", (name, tuple(label_names or ()), group_level)))
+
+    def register_histogram(
+        self,
+        name: str,
+        label_names: Optional[Sequence[str]] = None,
+        buckets: Optional[Sequence[float]] = None,
+        group_level: Optional[int] = None,
+    ) -> None:
+        self.calls.append(("register_histogram", (name, tuple(label_names or ()), tuple(buckets or ()), group_level)))
+
+    async def inc_counter(self, name: str, amount: float = 1.0, labels: Optional[Dict[str, str]] = None) -> None:
+        self.calls.append(("inc_counter", (name, amount, labels or {})))
+
+    async def observe_histogram(self, name: str, value: float, labels: Optional[Dict[str, str]] = None) -> None:
+        self.calls.append(("observe_histogram", (name, value, labels or {})))
+
+
+@pytest.mark.asyncio
+async def test_multi_metrics_backend_fans_out_calls() -> None:
+    backend_a = _RecordingBackend()
+    backend_b = _RecordingBackend()
+    multi = MultiMetricsBackend([backend_a, backend_b])
+
+    multi.register_counter("hits", label_names=["method"])
+    multi.register_histogram("latency", label_names=["method"], buckets=[0.1, 1.0])
+    await multi.inc_counter("hits", amount=2.5, labels={"method": "GET"})
+    await multi.observe_histogram("latency", value=0.4, labels={"method": "GET"})
+
+    expected_calls: List[Tuple[str, Tuple[Any, ...]]] = [
+        ("register_counter", ("hits", ("method",), None)),
+        ("register_histogram", ("latency", ("method",), (0.1, 1.0), None)),
+        ("inc_counter", ("hits", 2.5, {"method": "GET"})),
+        ("observe_histogram", ("latency", 0.4, {"method": "GET"})),
+    ]
+
+    assert backend_a.calls == expected_calls
+    assert backend_b.calls == expected_calls
+
+
+@pytest.mark.asyncio
+async def test_prometheus_backend_binds_stubbed_prometheus(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+
+    backend = metrics_module.PrometheusMetricsBackend()
+    backend.register_counter("hits", ["method"])
+    backend.register_histogram("latency", ["method"], buckets=[0.5])
+
+    await backend.inc_counter("hits", amount=2.0, labels={"method": "GET"})
+    await backend.observe_histogram("latency", value=0.1, labels={"method": "GET"})
+
+    counter_instance = stub.counter_instances[0]
+    histogram_instance = stub.histogram_instances[0]
+    assert counter_instance.children[(("method", "GET"),)].value == 2.0
+    assert histogram_instance.children[(("method", "GET"),)].values == [0.1]
+
+
+def test_prometheus_backend_normalizes_metric_names(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+
+    backend = metrics_module.PrometheusMetricsBackend()
+    backend.register_counter("api.v1.hits")
+    backend.register_histogram("latency.v1")
+
+    assert stub.counter_instances[0].name == "api_v1_hits"
+    assert stub.histogram_instances[0].name == "latency_v1"
+
+
+def test_prometheus_backend_detects_normalized_name_conflicts(monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = make_prometheus_stub()
+    monkeypatch.setitem(sys.modules, "prometheus_client", stub)
+
+    backend = metrics_module.PrometheusMetricsBackend()
+    backend.register_counter("api.v1.hits")
+    with pytest.raises(ValueError):
+        backend.register_histogram("api_v1_hits")
+
+
+def _split_segments(log_lines: Sequence[str]) -> List[str]:
+    segments: List[str] = []
+    for line in log_lines:
+        segments.extend(part.strip() for part in line.split("  "))
+    return [seg for seg in segments if seg]
+
+
+@pytest.mark.asyncio
+async def test_console_backend_sliding_window_rate_and_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=5.0, log_interval_seconds=0.0, group_level=2)
+    backend.register_counter("requests", ["group", "path"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0, 6.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    labels = {"group": "api", "path": "/list"}
+    await backend.inc_counter("requests", labels=labels)
+    await backend.inc_counter("requests", labels=labels)
+    await backend.inc_counter("requests", labels=labels)
+
+    segments = [seg for seg in _split_segments(logged) if seg.startswith("requests{group=api,path=/list}")]
+    assert len(segments) == 3
+
+    latest = segments[-1]
+    rate_str = latest.rsplit("=", 1)[1].rstrip("/s")
+    rate = float(rate_str)
+    assert abs(rate - 0.40) < 1e-2
+
+
+def _duration_to_seconds(payload: str) -> float:
+    if payload.endswith("ms"):
+        return float(payload[:-2]) / 1_000
+    if payload.endswith("µs"):
+        return float(payload[:-2]) / 1_000_000
+    if payload.endswith("ns"):
+        return float(payload[:-2]) / 1_000_000_000
+    if payload.endswith("s"):
+        return float(payload[:-1])
+    raise AssertionError(f"Unknown duration format: {payload}")
+
+
+@pytest.mark.asyncio
+async def test_console_backend_histogram_quantiles_and_group_depth(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=10.0, log_interval_seconds=0.0, group_level=2)
+    backend.register_histogram("latency", ["service", "endpoint", "status"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0, 2.0, 3.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    svc_a_labels = {"service": "svcA", "endpoint": "/search", "status": "200"}
+    svc_b_labels = {"service": "svcB", "endpoint": "/chat", "status": "500"}
+
+    await backend.observe_histogram("latency", value=0.01, labels=svc_a_labels)
+    await backend.observe_histogram("latency", value=0.02, labels=svc_a_labels)
+    await backend.observe_histogram("latency", value=0.03, labels=svc_a_labels)
+    await backend.observe_histogram("latency", value=2.0, labels=svc_b_labels)
+
+    svc_a_segments = [
+        seg for seg in _split_segments(logged) if seg.startswith("latency{service=svcA,endpoint=/search}")
+    ]
+    assert svc_a_segments, "expected log entries for service A"
+    latest = svc_a_segments[-1]
+    assert "status" not in latest
+
+    payload = latest.rsplit("=", 1)[1]
+    p50_str, p95_str, p99_str = payload.split(",", 2)
+    p50 = _duration_to_seconds(p50_str)
+    p95 = _duration_to_seconds(p95_str)
+    p99 = _duration_to_seconds(p99_str)
+
+    assert abs(p50 - 0.02) < 1e-3
+    assert abs(p95 - 0.029) < 5e-3
+    assert abs(p99 - 0.0298) < 5e-3
+
+
+@pytest.mark.asyncio
+async def test_console_backend_logs_all_metric_groups(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0)
+    backend.register_counter("requests", ["group"])
+    backend.register_counter("errors", ["group"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.inc_counter("requests", labels={"group": "api"})
+    await backend.inc_counter("errors", labels={"group": "api"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests=" in last_line
+    assert "errors=" in last_line
+
+
+def test_console_backend_snapshot_logs_single_line() -> None:
+    backend = ConsoleMetricsBackend()
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("counter", {"g": "1"}, [0.0, 1.0], [1.0, 1.0]),
+        ],
+        [
+            ("latency", {"g": "1"}, [0.1, 0.2], (0.5,)),
+        ],
+        snapshot_time=1.5,
+    )
+
+    assert logged and "counter=" in logged[0] and "latency=" in logged[0]
+
+
+def test_console_backend_snapshot_entries_are_sorted() -> None:
+    backend = ConsoleMetricsBackend(group_level=1)
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("zeta", {"a": "1"}, [0.0], [1.0]),
+            ("alpha", {"b": "2"}, [0.1], [2.0]),
+        ],
+        [],
+        snapshot_time=1.0,
+    )
+
+    assert logged, "expected log output"
+    line = logged[-1]
+    assert line.startswith("alpha{b=2}")
+    assert "  zeta{a=1}" in line
+
+
+def test_console_backend_group_level_none_aggregates_all_label_groups() -> None:
+    backend = ConsoleMetricsBackend(group_level=None)
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("requests", {"group": "api", "path": "/a"}, [0.0], [1.0]),
+            ("requests", {"group": "api", "path": "/b"}, [0.5], [2.0]),
+        ],
+        [
+            ("latency", {"group": "api", "path": "/a"}, [0.1], (0.5,)),
+            ("latency", {"group": "api", "path": "/b"}, [0.2], (0.5,)),
+        ],
+        snapshot_time=1.0,
+    )
+
+    assert logged, "expected log output"
+    line = logged[-1]
+    assert line.count("requests=") == 1
+    assert line.count("latency=") == 1
+
+
+def test_console_backend_group_level_positive_only_aggregates_prefix() -> None:
+    backend = ConsoleMetricsBackend(group_level=1)
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    backend._log_snapshot(
+        [
+            ("requests", {"a": "x", "b": "y"}, [0.0], [1.0]),
+            ("requests", {"a": "x", "b": "z"}, [0.1], [2.0]),
+            ("requests", {"a": "w", "b": "y"}, [0.2], [3.0]),
+        ],
+        [],
+        snapshot_time=1.0,
+    )
+
+    assert logged, "expected log output"
+    line = logged[-1]
+    # With group_level=1 we should see aggregation by the first declared label key.
+    assert line.count("requests{a=w}") == 1
+    assert line.count("requests{a=x}") == 1
+
+
+@pytest.mark.asyncio
+async def test_console_backend_logs_preserve_registration_label_order(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0, group_level=3)
+    backend.register_counter("requests", ["path", "method", "status"])
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.inc_counter("requests", labels={"status": "200", "path": "/search", "method": "GET"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests{path=/search,method=GET,status=200}" in last_line
+
+
+@pytest.mark.asyncio
+async def test_console_backend_metric_specific_group_level_applies(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0)
+    backend.register_counter("requests", ["service", "endpoint", "status"], group_level=2)
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.inc_counter("requests", labels={"status": "200", "endpoint": "/search", "service": "svcA"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests{service=svcA,endpoint=/search}" in last_line
+    assert "status" not in last_line
+
+
+@pytest.mark.asyncio
+async def test_console_backend_global_group_level_overrides_metric(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0, group_level=1)
+    backend.register_counter("requests", ["service", "endpoint"], group_level=2)
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.inc_counter("requests", labels={"service": "svcA", "endpoint": "/search"})
+
+    assert logged, "expected log output"
+    last_line = logged[-1]
+    assert "requests{service=svcA}" in last_line
+    assert "endpoint" not in last_line
+
+
+@pytest.mark.asyncio
+async def test_console_backend_histogram_metric_group_level(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend = ConsoleMetricsBackend(window_seconds=None, log_interval_seconds=0.0)
+    backend.register_histogram("latency", ["service", "endpoint"], group_level=1)
+
+    logged: List[str] = []
+    backend._log = logged.append  # type: ignore[assignment]
+
+    times = iter([0.0, 1.0, 2.0])
+    monkeypatch.setattr(metrics_module.time, "time", lambda: next(times))
+
+    await backend.observe_histogram("latency", value=0.01, labels={"service": "svcA", "endpoint": "/search"})
+    await backend.observe_histogram("latency", value=0.02, labels={"service": "svcA", "endpoint": "/chat"})
+
+    assert logged, "expected log output"
+    latest = logged[-1]
+    assert "latency{service=svcA}" in latest
+    assert "endpoint" not in latest

tests/utils/test_otel.py

@@ -0,0 +1,735 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+import logging
+from types import SimpleNamespace
+from typing import Any, Dict, List, cast
+
+import opentelemetry.trace as trace_api
+import pytest
+from opentelemetry.sdk.trace import ReadableSpan, SynchronousMultiSpanProcessor
+from opentelemetry.semconv.attributes import exception_attributes
+from opentelemetry.trace import TraceFlags
+from pydantic import ValidationError
+
+from agentlightning.semconv import LightningSpanAttributes, LinkPydanticModel
+from agentlightning.types.tracer import Span
+from agentlightning.utils import otel
+from agentlightning.utils.otel import (
+    check_attributes_sanity,
+    extract_links_from_attributes,
+    extract_tags_from_attributes,
+    filter_and_unflatten_attributes,
+    filter_attributes,
+    flatten_attributes,
+    format_exception_attributes,
+    full_qualified_name,
+    get_tracer,
+    get_tracer_provider,
+    make_link_attributes,
+    make_tag_attributes,
+    query_linked_spans,
+    sanitize_attribute_value,
+    sanitize_attributes,
+    sanitize_list_attribute_sanity,
+    unflatten_attributes,
+)
+
+pytestmark = pytest.mark.utils
+
+
+def _span_context(trace_id_hex: str, span_id_hex: str) -> trace_api.SpanContext:
+    return trace_api.SpanContext(
+        trace_id=int(trace_id_hex, 16),
+        span_id=int(span_id_hex, 16),
+        is_remote=False,
+        trace_flags=TraceFlags(TraceFlags.SAMPLED),
+        trace_state=trace_api.TraceState(),
+    )
+
+
+def test_flatten_simple_nested_dict_and_list() -> None:
+    data = {"a": {"b": 1, "c": [2, 3]}}
+    result = flatten_attributes(data)
+    assert result == {
+        "a.b": 1,
+        "a.c": [2, 3],
+    }
+
+
+def test_flatten_simple_nested_dict_and_list_with_leaf_expansion() -> None:
+    data = {"a": {"b": 1, "c": [2, 3]}}
+    result = flatten_attributes(data, expand_leaf_lists=True)
+    assert result == {
+        "a.b": 1,
+        "a.c.0": 2,
+        "a.c.1": 3,
+    }
+
+
+def test_full_qualified_name_handles_builtin_and_custom_classes() -> None:
+    class LocalClass:
+        pass
+
+    builtin_result = full_qualified_name(int)
+    local_result = full_qualified_name(LocalClass)
+
+    assert builtin_result == "int"
+    assert local_result.endswith("LocalClass")
+    assert local_result.startswith("tests.utils.test_otel")
+
+
+def test_flatten_empty_dict() -> None:
+    data: Dict[str, Any] = {}
+    assert flatten_attributes(data) == {}
+
+
+def test_flatten_empty_list() -> None:
+    data: List[Any] = []
+    # No elements -> no keys
+    assert flatten_attributes(data) == {}
+
+
+def test_flatten_root_list_of_primitives() -> None:
+    data = [10, 20, 30]
+    result = flatten_attributes(data)
+    assert result == {
+        "0": 10,
+        "1": 20,
+        "2": 30,
+    }
+
+
+def test_flatten_nested_lists_and_dicts() -> None:
+    data: Dict[str, Any] = {
+        "users": [
+            {"name": "Alice", "tags": ["admin", "staff"]},
+            {"name": "Bob", "tags": []},
+        ]
+    }
+    result = flatten_attributes(data)
+    assert result == {
+        "users.0.name": "Alice",
+        "users.0.tags": ["admin", "staff"],
+        "users.1.name": "Bob",
+        # Empty leaf lists remain implicit
+    }
+
+
+def test_flatten_leaf_lists_can_stay_compact() -> None:
+    data = {"tags": ["fast", "reliable"]}
+    result = flatten_attributes(data, expand_leaf_lists=False)
+    assert result == {"tags": ["fast", "reliable"]}
+
+
+def test_flatten_non_leaf_lists_still_expand_when_leaf_expansion_disabled() -> None:
+    data = {"users": [{"name": "Alice"}, {"name": "Bob"}]}
+    result = flatten_attributes(data, expand_leaf_lists=False)
+    assert result == {
+        "users.0.name": "Alice",
+        "users.1.name": "Bob",
+    }
+
+
+def test_flatten_leaf_lists_with_mixed_types_warns_and_expands(caplog: pytest.LogCaptureFixture) -> None:
+    data = {"values": [1, "two"]}
+    caplog.set_level(logging.WARNING, logger="agentlightning.utils.otel")
+
+    result = flatten_attributes(data, expand_leaf_lists=False)
+
+    assert result == {
+        "values.0": 1,
+        "values.1": "two",
+    }
+    assert "mixed primitive types" in caplog.text
+
+
+def test_flatten_leaf_lists_with_mixed_numbers_warns_and_expands(caplog: pytest.LogCaptureFixture) -> None:
+    data = {"values": [1, 2.0, 3]}
+    caplog.set_level(logging.WARNING, logger="agentlightning.utils.otel")
+
+    result = flatten_attributes(data, expand_leaf_lists=False)
+
+    assert result == {
+        "values.0": 1,
+        "values.1": 2.0,
+        "values.2": 3,
+    }
+    assert "mixed primitive types" in caplog.text
+
+
+def test_flatten_mixed_types_and_none() -> None:
+    data = {
+        "a": True,
+        "b": None,
+        "c": 3.14,
+        "d": "hello",
+        "e": {"f": False},
+    }
+    result = flatten_attributes(data)
+    assert result == {
+        "a": True,
+        "b": None,
+        "c": 3.14,
+        "d": "hello",
+        "e.f": False,
+    }
+
+
+def test_flatten_non_string_key_raises_value_error() -> None:
+    data = {
+        "a": {
+            1: "bad",  # non-string key inside nested dict
+        }
+    }
+    with pytest.raises(ValueError) as excinfo:
+        flatten_attributes(data)
+
+    msg = str(excinfo.value)
+    assert "Only string keys are supported in dictionaries" in msg
+    # Ensure the offending key is mentioned
+    assert "'1'" in msg
+    assert "type <class 'int'>" in msg
+
+
+def test_flatten_root_primitive_is_allowed() -> None:
+    # Even though the type hint says Dict/List, function behavior supports primitives.
+    data = 42
+    result = flatten_attributes(data)  # type: ignore[arg-type]
+    assert result == {"": 42}
+
+
+def test_unflatten_simple_nested_dict() -> None:
+    flat = {
+        "a.b": 1,
+        "a.c": 2,
+    }
+    result = unflatten_attributes(flat)
+    assert result == {"a": {"b": 1, "c": 2}}
+
+
+def test_unflatten_consecutive_numeric_keys_to_list() -> None:
+    flat = {
+        "a.0": "x",
+        "a.1": "y",
+        "a.2": "z",
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "a": ["x", "y", "z"],
+    }
+
+
+def test_unflatten_non_consecutive_numeric_keys_stays_dict() -> None:
+    flat = {
+        "a.0": "first",
+        "a.2": "third",
+    }
+    result = unflatten_attributes(flat)
+    # Keys are numeric but not consecutive -> remains dict
+    assert result == {
+        "a": {
+            "0": "first",
+            "2": "third",
+        }
+    }
+
+
+def test_unflatten_mixed_numeric_and_non_numeric_keys_stays_dict() -> None:
+    flat = {
+        "a.0": "zero",
+        "a.foo": "bar",
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "a": {
+            "0": "zero",
+            "foo": "bar",
+        }
+    }
+
+
+def test_unflatten_root_list_from_numeric_keys() -> None:
+    flat = {
+        "0": "a",
+        "1": "b",
+        "2": "c",
+    }
+    result = unflatten_attributes(flat)
+    # Root dict with all numeric keys 0..n-1 becomes list
+    assert result == ["a", "b", "c"]
+
+
+def test_unflatten_empty_flat_dict_returns_empty_dict() -> None:
+    flat: Dict[str, Any] = {}
+    result = unflatten_attributes(flat)
+    assert result == {}
+
+
+def test_unflatten_nested_lists_and_dicts() -> None:
+    flat = {
+        "users.0.name": "Alice",
+        "users.0.tags.0": "admin",
+        "users.0.tags.1": "staff",
+        "users.1.name": "Bob",
+        "users.1.tags.0": "guest",
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "users": [
+            {"name": "Alice", "tags": ["admin", "staff"]},
+            {"name": "Bob", "tags": ["guest"]},
+        ]
+    }
+
+
+def test_unflatten_list_of_lists() -> None:
+    flat = {
+        "a.0.0": 1,
+        "a.0.1": 2,
+        "a.1.0": 3,
+    }
+    result = unflatten_attributes(flat)
+    assert result == {
+        "a": [
+            [1, 2],
+            [3],
+        ]
+    }
+
+
+def test_unflatten_conflicting_primitive_and_nested_path_prefers_nested() -> None:
+    # "a" is first set to a primitive, then to a nested dict via "a.b"
+    flat = {
+        "a": 1,
+        "a.b": 2,
+    }
+    result = unflatten_attributes(flat)
+    # Primitive is overwritten by nested dict structure
+    assert result == {"a": {"b": 2}}
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        {"a": {"b": 1, "c": [2, 3]}},
+        {"x": [1, 2, {"y": 3}]},
+        {"root": [{"k": "v"}, {"k": "w"}]},
+        [{"name": "Alice"}, {"name": "Bob", "scores": [10, 20]}],
+    ],
+)
+def test_round_trip_flatten_then_unflatten_preserves_structure(value: Dict[str, Any] | List[Any]) -> None:
+    flat = flatten_attributes(value)  # type: ignore[arg-type]
+    reconstructed = unflatten_attributes(flat)
+    assert reconstructed == value
+
+
+@pytest.mark.parametrize(
+    "flat",
+    [
+        {"a.b": 1, "a.c": 2},
+        {"0": "x", "1": "y"},
+        {
+            "users.0.name": "Alice",
+            "users.1.name": "Bob",
+        },
+    ],
+)
+def test_round_trip_unflatten_then_flatten_preserves_flat_structure(flat: Dict[str, Any]) -> None:
+    nested = unflatten_attributes(flat)
+    re_flat = flatten_attributes(nested)
+    # Order of items in dict shouldn't matter
+    assert re_flat == flat
+
+
+def test_round_trip_with_empty_list_information_loss_is_expected() -> None:
+    """This documents the corner case: empty list flattens to {},
+    which unflattens back to {} (empty dict), losing the distinction.
+    """
+    data: List[Any] = []
+    flat = flatten_attributes(data)
+    assert flat == {}
+    reconstructed = unflatten_attributes(flat)
+    assert reconstructed == {}
+    assert reconstructed != data  # explicit documentation of the behavior
+
+
+def test_sanitize_attribute_value_handles_primitives_and_lists() -> None:
+    assert sanitize_attribute_value("text") == "text"
+    assert sanitize_attribute_value(42) == 42
+    assert sanitize_attribute_value(3.14) == 3.14
+    assert sanitize_attribute_value(True) is True
+    assert sanitize_attribute_value([True, 2]) == [1, 2]
+
+
+def test_sanitize_attribute_value_falls_back_to_json_for_hybrid_lists(caplog: pytest.LogCaptureFixture) -> None:
+    caplog.set_level(logging.WARNING, logger="agentlightning.utils.otel")
+    result = sanitize_attribute_value([1, "two"])
+    assert result == json.dumps([1, "two"])
+    assert "Failed to sanitize list attribute" in caplog.text
+
+
+def test_sanitize_attribute_value_serializes_non_serializable_objects_when_forced() -> None:
+    class Unserializable:
+        def __str__(self) -> str:
+            return "forced-serialization"
+
+    result = sanitize_attribute_value(Unserializable())
+    assert json.loads(cast(str, result)) == "forced-serialization"
+
+
+def test_sanitize_attribute_value_respects_force_flag() -> None:
+    class Unserializable:
+        pass
+
+    with pytest.raises(ValueError, match="Object must be JSON serializable"):
+        sanitize_attribute_value(Unserializable(), force=False)
+
+
+def test_sanitize_attributes_returns_clean_values() -> None:
+    attributes = {
+        "name": "agent",
+        "enabled": True,
+        "scores": [1, True],
+        "flags": [True, False],
+    }
+
+    result = sanitize_attributes(attributes)
+    assert result["name"] == "agent"
+    assert result["enabled"] is True
+    assert result["scores"] == [1, 1]
+    assert result["flags"] == [True, False]
+
+
+def test_sanitize_attributes_serializes_non_serializable_values_by_default() -> None:
+    class CustomValue:
+        def __str__(self) -> str:
+            return "custom-payload"
+
+    attributes = {"payload": CustomValue()}
+
+    result = sanitize_attributes(attributes)
+    assert json.loads(cast(str, result["payload"])) == "custom-payload"
+
+
+def test_sanitize_attributes_respects_force_flag() -> None:
+    class CustomValue:
+        pass
+
+    attributes = {"payload": CustomValue()}
+
+    with pytest.raises(ValueError, match="Failed to sanitize attribute 'payload'"):
+        sanitize_attributes(attributes, force=False)
+
+
+def test_sanitize_attributes_exposes_key_in_error() -> None:
+    class Bad:
+        pass
+
+    attributes = {"ok": "value", "bad": Bad()}
+    with pytest.raises(ValueError, match="Failed to sanitize attribute 'bad'"):
+        sanitize_attributes(attributes, force=False)
+
+
+def test_format_exception_attributes_captures_metadata() -> None:
+    try:
+        raise RuntimeError("boom")
+    except RuntimeError as err:
+        attrs = format_exception_attributes(err)
+
+    assert attrs[exception_attributes.EXCEPTION_TYPE] == "RuntimeError"
+    assert attrs[exception_attributes.EXCEPTION_MESSAGE] == "boom"
+    assert attrs[exception_attributes.EXCEPTION_ESCAPED] is True
+    assert exception_attributes.EXCEPTION_STACKTRACE in attrs
+    assert "RuntimeError: boom" in attrs[exception_attributes.EXCEPTION_STACKTRACE]  # type: ignore
+
+
+def test_sanitize_list_attribute_sanity_supports_primitive_lists() -> None:
+    assert sanitize_list_attribute_sanity(["a", "b"]) == ["a", "b"]
+    assert sanitize_list_attribute_sanity([True, False]) == [True, False]
+    assert sanitize_list_attribute_sanity([1, False]) == [1, 0]
+    assert sanitize_list_attribute_sanity([1.0, 2, True]) == [1.0, 2.0, 1.0]
+
+
+def test_sanitize_list_attribute_sanity_rejects_mixed_types() -> None:
+    with pytest.raises(ValueError, match="List must contain only one type of primitive values"):
+        sanitize_list_attribute_sanity([1, "two"])
+
+
+def test_check_attributes_sanity_accepts_valid_payload() -> None:
+    attributes: Dict[str, Any] = {
+        "name": "agent",
+        "count": 3,
+        "ratio": 0.5,
+        "enabled": False,
+        "flags": [True, False],
+        "scores": [1, True],
+    }
+
+    check_attributes_sanity(attributes)
+
+
+def test_check_attributes_sanity_requires_string_keys() -> None:
+    with pytest.raises(ValueError, match="Attribute key must be a string"):
+        check_attributes_sanity({1: "value"})  # type: ignore[arg-type]
+
+
+def test_check_attributes_sanity_wraps_list_errors() -> None:
+    with pytest.raises(ValueError, match="Failed to sanitize list attribute 'mixed'"):
+        check_attributes_sanity({"mixed": [1, "two"]})
+
+
+def test_check_attributes_sanity_rejects_non_primitive_values() -> None:
+    with pytest.raises(ValueError, match="Attribute value must be a string"):
+        check_attributes_sanity({"bad": {"nested": "value"}})
+
+
+def test_make_and_extract_link_attributes_round_trip() -> None:
+    flattened = make_link_attributes(
+        {
+            "gen_ai.response.id": "response-123",
+            "span_id": "abcd1234abcd1234",
+        }
+    )
+    assert flattened == {
+        f"{LightningSpanAttributes.LINK.value}.0.key_match": "gen_ai.response.id",
+        f"{LightningSpanAttributes.LINK.value}.0.value_match": "response-123",
+        f"{LightningSpanAttributes.LINK.value}.1.key_match": "span_id",
+        f"{LightningSpanAttributes.LINK.value}.1.value_match": "abcd1234abcd1234",
+    }
+
+    extracted = extract_links_from_attributes(flattened)
+    assert [link.model_dump() for link in extracted] == [
+        {"key_match": "gen_ai.response.id", "value_match": "response-123"},
+        {"key_match": "span_id", "value_match": "abcd1234abcd1234"},
+    ]
+
+
+def test_make_link_attributes_rejects_non_string_values() -> None:
+    with pytest.raises(ValueError) as excinfo:
+        make_link_attributes({"span_id": 123})  # type: ignore
+
+    assert "Link value must be a string" in str(excinfo.value)
+
+
+def test_make_tag_attributes_and_extract_round_trip() -> None:
+    flattened = make_tag_attributes(["fast", "reliable"])
+    assert flattened == {
+        f"{LightningSpanAttributes.TAG.value}.0": "fast",
+        f"{LightningSpanAttributes.TAG.value}.1": "reliable",
+    }
+
+    assert extract_tags_from_attributes(flattened) == ["fast", "reliable"]
+
+
+def test_extract_tags_from_attributes_rejects_non_strings() -> None:
+    attributes = {
+        f"{LightningSpanAttributes.TAG.value}.0": 1,
+    }
+
+    with pytest.raises(ValidationError):
+        extract_tags_from_attributes(attributes)
+
+
+def test_filter_attributes_keeps_exact_matches_and_children() -> None:
+    attributes = {
+        "agentlightning.link": "root",
+        "agentlightning.link.0.key_match": "trace_id",
+        "agentlightning.other": "discard",
+        "agentlightning.link_extra": "different_prefix",
+    }
+
+    filtered = filter_attributes(attributes, LightningSpanAttributes.LINK.value)
+    assert filtered == {
+        "agentlightning.link": "root",
+        "agentlightning.link.0.key_match": "trace_id",
+    }
+
+
+def test_filter_and_unflatten_attributes_strips_prefix_and_rebuilds_nested_structure() -> None:
+    attributes = {
+        f"{LightningSpanAttributes.LINK.value}.0.key_match": "trace_id",
+        f"{LightningSpanAttributes.LINK.value}.0.value_match": "aaa",
+        f"{LightningSpanAttributes.LINK.value}.1.key_match": "span_id",
+        f"{LightningSpanAttributes.LINK.value}.1.value_match": "bbb",
+    }
+
+    result = filter_and_unflatten_attributes(attributes, LightningSpanAttributes.LINK.value)
+    assert result == [
+        {"key_match": "trace_id", "value_match": "aaa"},
+        {"key_match": "span_id", "value_match": "bbb"},
+    ]
+
+
+def test_filter_and_unflatten_attributes_rejects_exact_prefix_key() -> None:
+    attributes = {LightningSpanAttributes.LINK.value: "invalid"}
+
+    with pytest.raises(ValueError):
+        filter_and_unflatten_attributes(attributes, LightningSpanAttributes.LINK.value)
+
+
+def test_query_linked_spans_matches_trace_id_on_readable_span() -> None:
+    readable_span = ReadableSpan(
+        name="upstream",
+        context=_span_context("a" * 32, "b" * 16),
+        attributes={},
+    )
+
+    assert readable_span.context is not None
+    links = [
+        LinkPydanticModel(key_match="trace_id", value_match=trace_api.format_trace_id(readable_span.context.trace_id)),
+    ]
+
+    matches = query_linked_spans([readable_span], links)
+    assert matches == [readable_span]
+
+
+def test_query_linked_spans_matches_custom_span_attributes() -> None:
+    custom_span = Span.from_attributes(
+        attributes={"gen_ai.response.id": "response-123", "custom": "needle"},
+        trace_id="c" * 32,
+        span_id="d" * 16,
+    )
+
+    links = [
+        LinkPydanticModel(key_match="gen_ai.response.id", value_match="response-123"),
+        LinkPydanticModel(key_match="custom", value_match="needle"),
+    ]
+
+    matches = query_linked_spans([custom_span], links)
+    assert matches == [custom_span]
+
+
+def test_query_linked_spans_excludes_span_with_mismatched_span_id() -> None:
+    span = Span.from_attributes(
+        attributes={"marker": "x"},
+        trace_id="e" * 32,
+        span_id="f" * 16,
+    )
+
+    links = [
+        LinkPydanticModel(key_match="span_id", value_match="deadbeefdeadbeef"),
+        LinkPydanticModel(key_match="marker", value_match="x"),
+    ]
+
+    assert query_linked_spans([span], links) == []
+
+
+def test_query_linked_spans_requires_all_links_to_match() -> None:
+    span = Span.from_attributes(
+        attributes={"marker": "x", "other": "y"},
+        trace_id="1" * 32,
+        span_id="2" * 16,
+    )
+
+    links = [
+        LinkPydanticModel(key_match="marker", value_match="x"),
+        LinkPydanticModel(key_match="other", value_match="z"),
+    ]
+
+    assert query_linked_spans([span], links) == []
+
+
+def test_query_linked_spans_handles_readable_span_without_context() -> None:
+    readable_span = ReadableSpan(name="orphan", context=None, attributes={"marker": "x"})
+
+    links = [LinkPydanticModel(key_match="marker", value_match="x")]
+
+    matches = query_linked_spans([readable_span], links)
+    assert matches == [readable_span]
+
+
+def test_get_tracer_provider_raises_when_tracer_uninitialized(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", None, raising=False)
+
+    with pytest.raises(RuntimeError):
+        get_tracer_provider(inspect=False)
+
+
+def test_get_tracer_provider_logs_when_provider_not_sdk(
+    caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    sentinel_provider = object()
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", sentinel_provider, raising=False)
+    monkeypatch.setattr(otel, "otel_get_tracer_provider", lambda: sentinel_provider)
+
+    caplog.set_level(logging.ERROR, logger=otel.logger.name)
+
+    returned = get_tracer_provider(inspect=False)
+
+    assert returned is sentinel_provider
+    assert any("Tracer provider is expected" in rec.getMessage() for rec in caplog.records)
+
+
+def test_get_tracer_delegates_to_active_span_processor(monkeypatch: pytest.MonkeyPatch) -> None:
+    class DummyProvider:
+        def __init__(self) -> None:
+            self.calls: List[str] = []
+
+        def get_tracer(self, name: str) -> str:
+            self.calls.append(name)
+            return f"tracer:{name}"
+
+    provider = DummyProvider()
+
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", object(), raising=False)
+    monkeypatch.setattr(otel, "get_tracer_provider", lambda inspect=True: provider)
+
+    tracer = get_tracer()
+
+    assert tracer == "tracer:agentlightning"
+    assert provider.calls == ["agentlightning"]
+
+
+def test_get_tracer_without_active_span_processor_builds_isolated_tracer(monkeypatch: pytest.MonkeyPatch) -> None:
+    provider = SimpleNamespace(
+        sampler="sampler",
+        resource="resource",
+        id_generator="id_gen",
+    )
+
+    created: Dict[str, Any] = {}
+
+    class DummyTracer:
+        def __init__(
+            self,
+            sampler: Any,
+            resource: Any,
+            span_processor: SynchronousMultiSpanProcessor,
+            id_generator: Any,
+            instrumentation_info: Any,
+            span_limits: Any,
+            instrumentation_scope: Any,
+        ) -> None:
+            created["args"] = (
+                sampler,
+                resource,
+                span_processor,
+                id_generator,
+                instrumentation_info,
+                span_limits,
+                instrumentation_scope,
+            )
+
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", object(), raising=False)
+    monkeypatch.setattr(otel, "get_tracer_provider", lambda inspect=True: provider)
+    monkeypatch.setattr(otel, "Tracer", DummyTracer)
+
+    tracer = get_tracer(use_active_span_processor=False)
+
+    assert isinstance(created.get("args"), tuple)
+    assert created["args"][0] == "sampler"
+    assert created["args"][1] == "resource"
+    assert isinstance(created["args"][2], SynchronousMultiSpanProcessor)
+    assert created["args"][3] == "id_gen"
+    assert created["args"][4].name == "agentlightning"
+    assert tracer is not None
+
+
+def test_get_tracer_raises_when_provider_not_initialized(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(trace_api, "_TRACER_PROVIDER", None, raising=False)
+
+    with pytest.raises(RuntimeError):
+        get_tracer()

tests/utils/test_otlp.py

@@ -0,0 +1,453 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+import gzip
+from types import SimpleNamespace
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, cast
+
+import pytest
+from fastapi import Request
+from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
+    ExportTraceServiceRequest,
+    ExportTraceServiceResponse,
+)
+from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
+from opentelemetry.proto.resource.v1.resource_pb2 import Resource as ProtoResource
+from opentelemetry.proto.trace.v1.trace_pb2 import Span as ProtoSpan
+from opentelemetry.proto.trace.v1.trace_pb2 import Status as ProtoStatus
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExportResult
+from starlette.types import Message, Scope
+
+from agentlightning.semconv import LightningResourceAttributes
+from agentlightning.utils import otlp
+
+BASE_TIME_NANOS = 1_700_000_000_000_000_000
+EVENT_TIME_OFFSET = 3_000_000_000
+EVENT_TIME_SECONDS = (BASE_TIME_NANOS + EVENT_TIME_OFFSET) / 1_000_000_000
+EXTRA_EVENT_TIME_OFFSET = 4_000_000_000
+EXTRA_EVENT_TIME_SECONDS = (BASE_TIME_NANOS + EXTRA_EVENT_TIME_OFFSET) / 1_000_000_000
+
+pytestmark = pytest.mark.utils
+
+
+class _StubStore:
+    def __init__(self) -> None:
+        self.sequence_calls: List[tuple[str, str]] = []
+        self.next_value = 1
+
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        self.sequence_calls.extend(rollout_attempt_ids)
+        allocations: List[int] = []
+        for _ in rollout_attempt_ids:
+            allocations.append(self.next_value)
+            self.next_value += 1
+        return allocations
+
+
+def _make_request(
+    body: bytes,
+    *,
+    content_type: str = otlp.PROTOBUF_CT,
+    content_encoding: Optional[str] = None,
+    accept_encoding: Optional[str] = None,
+) -> Request:
+    headers: List[tuple[bytes, bytes]] = [(b"content-type", content_type.encode())]
+    if content_encoding:
+        headers.append((b"content-encoding", content_encoding.encode()))
+    if accept_encoding:
+        headers.append((b"accept-encoding", accept_encoding.encode()))
+
+    scope: Scope = {
+        "type": "http",
+        "asgi": {"version": "3.0", "spec_version": "2.3"},
+        "method": "POST",
+        "path": "/v1/test",
+        "headers": headers,
+    }
+
+    body_sent = False
+
+    async def receive() -> Message:
+        nonlocal body_sent
+        if body_sent:
+            return {"type": "http.request", "body": b"", "more_body": False}
+        body_sent = True
+        return {"type": "http.request", "body": body, "more_body": False}
+
+    return Request(scope, receive)
+
+
+def _set_any_value(av: AnyValue, value: object) -> None:
+    if isinstance(value, bool):
+        av.bool_value = value
+    elif isinstance(value, int):
+        av.int_value = value
+    elif isinstance(value, float):
+        av.double_value = value
+    elif isinstance(value, bytes):
+        av.bytes_value = value
+    elif isinstance(value, list):
+        for item in cast(List[Any], value):
+            _set_any_value(av.array_value.values.add(), item)
+    elif isinstance(value, dict):
+        for key, item in cast(Dict[str, Any], value).items():
+            kv = av.kvlist_value.values.add()
+            kv.key = key
+            _set_any_value(kv.value, item)
+    else:
+        av.string_value = str(value)
+
+
+def _add_attribute(attrs: Iterable[KeyValue], key: str, value: object) -> None:
+    kv = attrs.add()  # type: ignore
+    kv.key = key
+    _set_any_value(kv.value, value)  # type: ignore
+
+
+def _build_span_request() -> ExportTraceServiceRequest:
+    request = ExportTraceServiceRequest()
+    resource_spans = request.resource_spans.add()
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "resource-rollout")
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "resource-attempt")
+    _add_attribute(
+        resource_spans.resource.attributes,
+        LightningResourceAttributes.SPAN_SEQUENCE_ID.value,
+        "5",
+    )
+    resource_spans.schema_url = "https://example/schema"
+
+    scope_spans = resource_spans.scope_spans.add()
+    span = scope_spans.spans.add()
+    span.trace_id = bytes.fromhex("01" * 16)
+    span.span_id = bytes.fromhex("02" * 8)
+    span.parent_span_id = bytes.fromhex("03" * 8)
+    span.name = "test-span"
+    span.start_time_unix_nano = BASE_TIME_NANOS
+    span.end_time_unix_nano = BASE_TIME_NANOS + 2_000_000_000
+    span.status.code = ProtoStatus.STATUS_CODE_ERROR
+    span.status.message = "boom"
+
+    _add_attribute(span.attributes, "foo", "bar")
+    _add_attribute(span.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "span-rollout")
+    _add_attribute(span.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "span-attempt")
+    _add_attribute(span.attributes, LightningResourceAttributes.SPAN_SEQUENCE_ID.value, "7")
+
+    event = span.events.add()
+    event.name = "event"
+    event.time_unix_nano = BASE_TIME_NANOS + EVENT_TIME_OFFSET
+    _add_attribute(event.attributes, "event-attr", 9)
+
+    link = span.links.add()
+    link.trace_id = bytes.fromhex("04" * 16)
+    link.span_id = bytes.fromhex("05" * 8)
+    _add_attribute(link.attributes, "link-attr", True)
+
+    return request
+
+
+@pytest.mark.asyncio
+async def test_handle_otlp_export_success_with_gzip_response() -> None:
+    request_msg = _build_span_request()
+    body = request_msg.SerializeToString()
+    request = _make_request(
+        body,
+        accept_encoding="gzip;q=0.9,br",
+    )
+
+    received: List[ExportTraceServiceRequest] = []
+
+    async def callback(message: ExportTraceServiceRequest) -> None:
+        received.append(message)
+
+    response = await otlp.handle_otlp_export(
+        request,
+        ExportTraceServiceRequest,
+        ExportTraceServiceResponse,
+        callback,
+        signal_name="traces",
+    )
+
+    assert response.status_code == 200
+    assert received and received[0].SerializeToString() == body
+    assert response.headers["Content-Encoding"] == "gzip"
+    assert gzip.decompress(response.body) == ExportTraceServiceResponse().SerializeToString()
+
+
+@pytest.mark.asyncio
+async def test_handle_otlp_export_rejects_invalid_content_type() -> None:
+    request = _make_request(b"{}", content_type="application/json")
+
+    response = await otlp.handle_otlp_export(
+        request,
+        ExportTraceServiceRequest,
+        ExportTraceServiceResponse,
+        None,
+        signal_name="traces",
+    )
+
+    assert response.status_code == 400
+    status = otlp.Status()  # type: ignore[attr-defined]
+    status.ParseFromString(response.body)  # type: ignore
+    assert "Unsupported Content-Type" in status.message
+
+
+@pytest.mark.asyncio
+async def test_handle_otlp_export_rejects_bad_payload() -> None:
+    request = _make_request(b"not-a-proto")
+
+    response = await otlp.handle_otlp_export(
+        request,
+        ExportTraceServiceRequest,
+        ExportTraceServiceResponse,
+        None,
+        signal_name="traces",
+    )
+
+    assert response.status_code == 400
+    status = otlp.Status()  # type: ignore[attr-defined]
+    status.ParseFromString(response.body)  # type: ignore
+    assert "Unable to parse" in status.message
+
+
+@pytest.mark.asyncio
+async def test_handle_otlp_export_accepts_gzip_body() -> None:
+    request_msg = ExportTraceServiceRequest()
+    request_msg.resource_spans.add()
+    gz_body = gzip.compress(request_msg.SerializeToString())
+    request = _make_request(gz_body, content_encoding="gzip")
+
+    response = await otlp.handle_otlp_export(
+        request,
+        ExportTraceServiceRequest,
+        ExportTraceServiceResponse,
+        None,
+        signal_name="traces",
+    )
+
+    assert response.status_code == 200
+
+
+@pytest.mark.asyncio
+async def test_spans_from_proto_prefers_span_level_metadata() -> None:
+    store = _StubStore()
+    request = _build_span_request()
+
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
+
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.rollout_id == "span-rollout"
+    assert span.attempt_id == "span-attempt"
+    assert span.sequence_id == 7
+    assert span.status.status_code == "ERROR"
+    assert span.events[0].timestamp == pytest.approx(EVENT_TIME_SECONDS)  # type: ignore
+    assert span.links[0].context.trace_id == "0404" * 8
+    assert span.links[0].attributes == {"link-attr": True}
+    assert span.resource.attributes[LightningResourceAttributes.ROLLOUT_ID.value] == "resource-rollout"
+    assert span.resource.schema_url == "https://example/schema"
+    assert not store.sequence_calls
+
+
+@pytest.mark.asyncio
+async def test_spans_from_proto_requests_sequence_ids_when_missing() -> None:
+    store = _StubStore()
+    request = ExportTraceServiceRequest()
+    resource_spans = request.resource_spans.add()
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "r1")
+    _add_attribute(resource_spans.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a1")
+
+    scope_span = resource_spans.scope_spans.add()
+    span = scope_span.spans.add()
+    span.trace_id = b""  # exercise default ids
+    span.span_id = b""
+    span.name = "needs-seq"
+
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
+
+    assert len(spans) == 1
+    assert spans[0].sequence_id == 1
+    assert store.sequence_calls == [("r1", "a1")]
+
+
+@pytest.mark.asyncio
+async def test_spans_from_proto_bulk_issues_for_mixed_rollouts() -> None:
+    store = _StubStore()
+    request = ExportTraceServiceRequest()
+
+    resource_first = request.resource_spans.add()
+    _add_attribute(resource_first.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "r1")
+    _add_attribute(resource_first.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a-default")
+
+    scope_first = resource_first.scope_spans.add()
+    span_missing = scope_first.spans.add()
+    span_missing.trace_id = bytes.fromhex("11" * 16)
+    span_missing.span_id = bytes.fromhex("22" * 8)
+    span_missing.name = "missing-seq"
+    _add_attribute(span_missing.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a1")
+
+    span_negative = scope_first.spans.add()
+    span_negative.trace_id = bytes.fromhex("33" * 16)
+    span_negative.span_id = bytes.fromhex("44" * 8)
+    span_negative.name = "negative-seq"
+    _add_attribute(span_negative.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "a2")
+    _add_attribute(span_negative.attributes, LightningResourceAttributes.SPAN_SEQUENCE_ID.value, "-5")
+
+    resource_second = request.resource_spans.add()
+    _add_attribute(resource_second.resource.attributes, LightningResourceAttributes.ROLLOUT_ID.value, "r2")
+    _add_attribute(resource_second.resource.attributes, LightningResourceAttributes.ATTEMPT_ID.value, "b1")
+
+    scope_second = resource_second.scope_spans.add()
+    span_second = scope_second.spans.add()
+    span_second.trace_id = bytes.fromhex("55" * 16)
+    span_second.span_id = bytes.fromhex("66" * 8)
+    span_second.name = "other-rollout"
+
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
+    assert [span.name for span in spans] == ["missing-seq", "negative-seq", "other-rollout"]
+    assert [span.sequence_id for span in spans] == [1, 2, 3]
+    assert store.sequence_calls == [("r1", "a1"), ("r1", "a2"), ("r2", "b1")]
+
+
+@pytest.mark.asyncio
+async def test_spans_from_proto_skips_spans_without_ids() -> None:
+    store = _StubStore()
+    request = ExportTraceServiceRequest()
+    request.resource_spans.add()  # missing rollout and attempt
+
+    spans = await otlp.spans_from_proto(request, store.get_many_span_sequence_ids)
+
+    assert spans == []
+    assert store.sequence_calls == []
+
+
+def test_normalize_sequence_id_handles_bad_values(caplog: pytest.LogCaptureFixture) -> None:
+    caplog.set_level("WARNING")
+    assert otlp._normalize_sequence_id("not-int") is None
+    assert any("Invalid sequence_id" in record.message for record in caplog.records)
+
+
+def test_any_value_to_python_full_roundtrip() -> None:
+    av = AnyValue()
+    _set_any_value(
+        av,
+        {
+            "s": "hello",
+            "b": True,
+            "i": 5,
+            "d": 1.5,
+            "arr": ["x", 2],
+            "nested": {"k": b"\x01"},
+        },
+    )
+
+    result = otlp._any_value_to_python(av)
+    assert result == {
+        "s": "hello",
+        "b": True,
+        "i": 5,
+        "d": 1.5,
+        "arr": ["x", 2],
+        "nested": {"k": "01"},
+    }
+
+
+def test_kv_list_to_dict_converts_values() -> None:
+    resource = ProtoResource()
+    _add_attribute(resource.attributes, "num", 10)
+    _add_attribute(resource.attributes, "flag", False)
+
+    converted = otlp._kv_list_to_dict(resource.attributes)
+    assert converted == {"num": 10, "flag": False}
+
+
+def test_bytes_to_hex_helpers() -> None:
+    assert otlp._bytes_to_trace_id_hex(b"") == "0" * 32
+    assert otlp._bytes_to_span_id_hex(b"") == "0" * 16
+    assert otlp._bytes_to_trace_id_hex(b"\xff") == "ff".rjust(32, "0")
+    assert otlp._bytes_to_span_id_hex(b"\xaa") == "aa".rjust(16, "0")
+
+
+def test_events_and_links_from_proto() -> None:
+    span = ProtoSpan()
+    event = span.events.add()
+    event.name = "evt"
+    event.time_unix_nano = BASE_TIME_NANOS + EXTRA_EVENT_TIME_OFFSET
+    _add_attribute(event.attributes, "alpha", "beta")
+
+    link = span.links.add()
+    link.trace_id = bytes.fromhex("06" * 16)
+    link.span_id = bytes.fromhex("07" * 8)
+    _add_attribute(link.attributes, "delta", 1)
+
+    events = otlp._events_from_proto(span)
+    links = otlp._links_from_proto(span)
+
+    assert events[0].timestamp == pytest.approx(EXTRA_EVENT_TIME_SECONDS)  # type: ignore
+    assert events[0].attributes == {"alpha": "beta"}
+    assert links[0].context.trace_id == "0606" * 8
+    assert links[0].attributes == {"delta": 1}
+
+
+def test_resource_from_proto() -> None:
+    resource = ProtoResource()
+    _add_attribute(resource.attributes, "key", "value")
+    result = otlp._resource_from_proto(resource, schema_url="https://example/schema")
+    assert result.attributes == {"key": "value"}
+    assert result.schema_url == "https://example/schema"
+
+
+def test_maybe_gzip_response_parses_quality_values() -> None:
+    request = SimpleNamespace(headers={"Accept-Encoding": "br, gzip;q=0.1"})
+    payload = b"payload"
+    compressed, headers = otlp._maybe_gzip_response(cast(Request, request), payload)
+
+    assert headers == {"Content-Encoding": "gzip"}
+    assert gzip.decompress(compressed) == payload
+
+
+def test_bad_request_response_matches_request_encoding() -> None:
+    request = SimpleNamespace(headers={"Accept-Encoding": "gzip"})
+    response = otlp._bad_request_response(cast(Request, request), "error")
+
+    assert response.status_code == 400
+    assert response.media_type == otlp.PROTOBUF_CT
+    status = otlp.Status()  # type: ignore[attr-defined]
+    status.ParseFromString(gzip.decompress(response.body))
+    assert status.message == "error"
+
+
+class _DummyReadableSpan:
+    def __init__(self) -> None:
+        self._resource = Resource.create({"existing": "value"})
+
+
+def test_lightning_store_otlp_exporter_overrides_resources(monkeypatch: pytest.MonkeyPatch) -> None:
+    exporter = otlp.LightningStoreOTLPExporter(endpoint="http://collector")
+
+    captured_spans: List[List[_DummyReadableSpan]] = []
+
+    def fake_export(self: otlp.LightningStoreOTLPExporter, spans: List[_DummyReadableSpan]) -> SpanExportResult:
+        captured_spans.append(spans)
+        return SpanExportResult.SUCCESS
+
+    monkeypatch.setattr(otlp.OTLPSpanExporter, "export", fake_export, raising=False)
+
+    exporter.enable_store_otlp("http://store", "rollout", "attempt")
+    span = _DummyReadableSpan()
+
+    result = exporter.export([cast(ReadableSpan, span)])
+
+    assert result == SpanExportResult.SUCCESS
+    assert captured_spans
+    attributes = captured_spans[0][0]._resource.attributes  # type: ignore[attr-defined]
+    assert attributes[LightningResourceAttributes.ROLLOUT_ID.value] == "rollout"
+    assert attributes[LightningResourceAttributes.ATTEMPT_ID.value] == "attempt"
+
+    exporter.disable_store_otlp()
+    assert exporter._rollout_id is None
+    assert exporter._attempt_id is None

tests/utils/test_system_snapshot.py

@@ -0,0 +1,140 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any, Callable, Optional
+
+import pytest
+
+from agentlightning.utils import system_snapshot
+
+pytestmark = pytest.mark.utils
+
+
+def _patch_system_snapshot(
+    monkeypatch: pytest.MonkeyPatch,
+    include_gpu: bool = False,
+    cpu_percent_fn: Optional[Callable[[float], float]] = None,
+) -> Optional[SimpleNamespace]:
+    monkeypatch.setattr(system_snapshot.platform, "processor", lambda: "test-cpu")
+    monkeypatch.setattr(system_snapshot.platform, "platform", lambda: "test-platform")
+    monkeypatch.setattr(system_snapshot.socket, "gethostname", lambda: "test-host")
+
+    def fake_cpu_count(logical: bool = True) -> int:
+        return 4 if logical else 2
+
+    monkeypatch.setattr(system_snapshot.psutil, "cpu_count", fake_cpu_count)
+
+    def default_cpu_percent(interval: float) -> float:
+        return 33.3
+
+    monkeypatch.setattr(
+        system_snapshot.psutil,
+        "cpu_percent",
+        cpu_percent_fn or default_cpu_percent,
+    )
+
+    vm = SimpleNamespace(used=5 * (2**30), total=10 * (2**30), percent=50.0)
+    monkeypatch.setattr(system_snapshot.psutil, "virtual_memory", lambda: vm)
+
+    du = SimpleNamespace(used=2 * (2**30), total=8 * (2**30), percent=25.0)
+    monkeypatch.setattr(system_snapshot.psutil, "disk_usage", lambda _: du)  # type: ignore
+
+    net = SimpleNamespace(bytes_sent=4 * (2**20), bytes_recv=6 * (2**20))
+    monkeypatch.setattr(system_snapshot.psutil, "net_io_counters", lambda: net)
+
+    if include_gpu:
+        dummy_gpu = SimpleNamespace(
+            name="Test GPU",
+            utilization=37.5,
+            memory_used=1024,
+            memory_total=4096,
+            temperature=68,
+        )
+        monkeypatch.setattr(
+            system_snapshot.GPUStatCollection,
+            "new_query",
+            lambda *args, **kwargs: SimpleNamespace(gpus=[dummy_gpu]),  # type: ignore
+        )
+
+        return dummy_gpu
+    return None
+
+
+def test_system_snapshot_excludes_gpu_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
+    def fail_gpu_query(*args: Any, **kwargs: Any) -> Any:
+        raise AssertionError("GPU stats should not be queried when include_gpu is False")
+
+    monkeypatch.setattr(system_snapshot.GPUStatCollection, "new_query", fail_gpu_query)
+
+    _patch_system_snapshot(monkeypatch)
+
+    snapshot = system_snapshot.system_snapshot()
+
+    assert snapshot["cpu_name"] == "test-cpu"
+    assert snapshot["cpu_cores"] == 2
+    assert snapshot["cpu_threads"] == 4
+    assert snapshot["cpu_usage_pct"] == 33.3
+    assert snapshot["mem_used_gb"] == 5.0
+    assert snapshot["mem_total_gb"] == 10.0
+    assert snapshot["mem_pct"] == 50.0
+    assert snapshot["disk_used_gb"] == 2.0
+    assert snapshot["disk_total_gb"] == 8.0
+    assert snapshot["disk_pct"] == 25.0
+    assert snapshot["bytes_sent_mb"] == 4.0
+    assert snapshot["bytes_recv_mb"] == 6.0
+    assert snapshot["host"] == "test-host"
+    assert snapshot["os"] == "test-platform"
+    assert "gpus" not in snapshot
+
+
+def test_system_snapshot_includes_gpus_when_requested(monkeypatch: pytest.MonkeyPatch) -> None:
+    dummy_gpu = _patch_system_snapshot(monkeypatch, include_gpu=True)
+    assert dummy_gpu is not None
+
+    snapshot = system_snapshot.system_snapshot(include_gpu=True)
+
+    assert snapshot["gpus"] == [
+        {
+            "gpu": dummy_gpu.name,
+            "util_pct": dummy_gpu.utilization,
+            "mem_used_mb": dummy_gpu.memory_used,
+            "mem_total_mb": dummy_gpu.memory_total,
+            "temp_c": dummy_gpu.temperature,
+        }
+    ]
+
+
+def test_sanity_check() -> None:
+    snapshot = system_snapshot.system_snapshot()
+    assert snapshot is not None
+
+    snapshot = system_snapshot.system_snapshot(include_gpu=True)
+    assert snapshot is not None
+
+    try:
+        import torch  # type: ignore
+
+        GPU_AVAILABLE = torch.cuda.is_available()
+    except Exception:
+        GPU_AVAILABLE = False  # type: ignore
+
+    if GPU_AVAILABLE:
+        assert snapshot["gpus"] is not None
+        assert len(snapshot["gpus"]) > 0
+
+
+def test_system_snapshot_cpu_percent_samples_immediately(monkeypatch: pytest.MonkeyPatch) -> None:
+    called_intervals: list[float] = []
+
+    def record_cpu_percent(interval: float) -> float:
+        called_intervals.append(interval)
+        return 77.7
+
+    _patch_system_snapshot(monkeypatch, cpu_percent_fn=record_cpu_percent)
+
+    snapshot = system_snapshot.system_snapshot()
+
+    assert snapshot["cpu_usage_pct"] == 77.7
+    assert called_intervals == [0.0]