diff --git a/AGENTS.md b/AGENTS.md index 5a86d70..f033ded 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,6 +38,10 @@ Key variables: | `HARNESS_CONFIG_FILE` | Path to YAML config file (overrides env) | | `HARNESS_GEN_AI_PAYLOAD_CAPTURE_ENABLED` | Capture LLM prompt/response payloads (default: off) | | `HARNESS_GEN_AI_PAYLOAD_EVALUATION_ENABLED` | Run control plugins on GenAI spans | +| `HARNESS_SPAN_ATTRIBUTE_FLATTEN_ENABLED` | Flatten dict values from `set_span_attribute` into dot-notation keys (default: on; only `false` disables) | +| `HARNESS_SPAN_ATTRIBUTE_FLATTEN_MAX_DEPTH` | Max nesting depth when flattening dicts (default: 3) | +| `HARNESS_SPAN_ATTRIBUTE_FLATTEN_MAX_LEAVES` | Max leaf attributes emitted per dict (default: 32) | +| `HARNESS_SPAN_ATTRIBUTE_FLATTEN_RAW_JSON` | Also keep the original key as a JSON string when flattening (default: off) | ### Instrumentation opt-in (`HARNESS_` or `HA_` prefix; `HARNESS_` wins — no `AT_`/`TA_` aliases) diff --git a/README.md b/README.md index 5589c9b..382c5e7 100644 --- a/README.md +++ b/README.md @@ -168,6 +168,37 @@ For duplicate keys, the last write wins; instrumentation can overwrite a custome if it writes the same key later. Because the span already exists, these attributes cannot influence head sampling. +#### Dictionary values + +Dictionaries are also accepted and are flattened into dot-notation attributes so each +leaf stays individually queryable in the backend: + +```python +set_span_attribute("agent", {"action": "generate", "model": {"name": "gemini-2.0"}}) +# exported as: agent.action="generate", agent.model.name="gemini-2.0" +``` + +Nothing is serialized when you call the helper — the dict is flattened at span end, just +before export. Flattening rules: + +| Case | Result | +|---|---| +| `str` / `bool` / `int` / `float` leaf | kept as the native OTel type | +| `None` leaf | skipped | +| Any other object | `str(value)` | +| List of same-typed scalars | OTel array attribute at that key | +| List of dicts or mixed types | JSON string at that key | +| Nesting deeper than the configured max depth (default 3) | JSON string at the depth limit | +| Flattened key already set on the span | skipped — the explicit value wins | + +At most `HARNESS_SPAN_ATTRIBUTE_FLATTEN_MAX_LEAVES` leaf attributes are emitted per +dictionary (default 32); the rest are dropped with a debug log. Override depth via +`HARNESS_SPAN_ATTRIBUTE_FLATTEN_MAX_DEPTH` (default 3). The original key (`agent`) is not set unless +`HARNESS_SPAN_ATTRIBUTE_FLATTEN_RAW_JSON=true`, which additionally stores the whole dict +as JSON there. Flattening is **enabled by default**; set +`HARNESS_SPAN_ATTRIBUTE_FLATTEN_ENABLED=false` to turn it off, in which case OTel rejects +dictionary values as it did before. + ## Plugins The SDK loads extensions via [setuptools entry points](https://setuptools.pypa.io/en/latest/userguide/entry_point.html). Each plugin has a **name** (the entry-point key). Names are listed in config or environment variables; only installed plugins are loaded, in the order you configure. diff --git a/src/harness_sdk/agent_init.py b/src/harness_sdk/agent_init.py index c1f51f1..ad3fa16 100644 --- a/src/harness_sdk/agent_init.py +++ b/src/harness_sdk/agent_init.py @@ -18,6 +18,8 @@ from harness_sdk import constants from harness_sdk.config import config_pb2 from harness_sdk.env import is_env_flag_enabled +from harness_sdk.flatten_dict_registry import is_flatten_enabled +from harness_sdk.flatten_dict_span_processor import FlattenDictSpanProcessor from harness_sdk.otlp_reporting import ( compression_type_to_otlp_grpc, compression_type_to_otlp_http, @@ -97,8 +99,10 @@ def register_processor(self, processor) -> None: def set_console_span_processor(self) -> None: console_span_exporter = ConsoleSpanExporter( service_name=self._config.config.service_name) - simple_export_span_processor = SimpleSpanProcessor(console_span_exporter) - trace.get_tracer_provider().add_span_processor(simple_export_span_processor) + processor = SimpleSpanProcessor(console_span_exporter) + if is_flatten_enabled(): + processor = FlattenDictSpanProcessor(processor) + trace.get_tracer_provider().add_span_processor(processor) def init_exporter(self, trace_reporter_type): exporter_type = '' diff --git a/src/harness_sdk/flatten_dict_registry.py b/src/harness_sdk/flatten_dict_registry.py new file mode 100644 index 0000000..8364172 --- /dev/null +++ b/src/harness_sdk/flatten_dict_registry.py @@ -0,0 +1,122 @@ +"""Thread-safe hand-off of dict span attributes from the hot path to span end. + +``set_span_attribute("agent", {...})`` cannot go through +``span.set_attribute``: OTel rejects mapping values outright. Serializing on +the caller's thread would put JSON encoding on the application hot path, so +dict values are parked here untouched and flattened later by +``FlattenDictSpanProcessor.on_end``. + +Entries are keyed by span context rather than object identity because +``Span.end()`` hands ``on_end`` a fresh ``ReadableSpan`` snapshot, not the +recording ``Span`` the enrichment helper saw. +""" +import threading +from collections import OrderedDict + +from harness_sdk.custom_logger import get_custom_logger +from harness_sdk.env import get_env_value, is_env_flag_enabled + +logger = get_custom_logger(__name__) + +FLATTEN_ENABLED_ENV = "SPAN_ATTRIBUTE_FLATTEN_ENABLED" +FLATTEN_RAW_JSON_ENV = "SPAN_ATTRIBUTE_FLATTEN_RAW_JSON" +FLATTEN_MAX_DEPTH_ENV = "SPAN_ATTRIBUTE_FLATTEN_MAX_DEPTH" +FLATTEN_MAX_LEAVES_ENV = "SPAN_ATTRIBUTE_FLATTEN_MAX_LEAVES" + +DEFAULT_MAX_DEPTH = 3 +DEFAULT_MAX_LEAVES = 32 + +# Bound on spans holding pending dicts. A span that is never ended would +# otherwise leak its entry forever; oldest entries are evicted instead. +_MAX_TRACKED_SPANS = 2048 + + +def is_flatten_enabled(): + """Dict flattening is on by default; only explicit ``false`` disables it.""" + value = get_env_value(FLATTEN_ENABLED_ENV) + if value is None: + return True + return value.strip().lower() != "false" + + +def is_raw_json_enabled(): + """Opt in to additionally keeping the original key as a JSON string.""" + return is_env_flag_enabled(FLATTEN_RAW_JSON_ENV) + + +def _positive_int_env(env_key, default): + raw = (get_env_value(env_key) or "").strip() + if not raw: + return default + try: + value = int(raw) + return value if value > 0 else default + except ValueError: + return default + + +def get_flatten_max_depth(): + return _positive_int_env(FLATTEN_MAX_DEPTH_ENV, DEFAULT_MAX_DEPTH) + + +def get_flatten_max_leaves(): + return _positive_int_env(FLATTEN_MAX_LEAVES_ENV, DEFAULT_MAX_LEAVES) + + +def _span_key(span): + get_context = getattr(span, "get_span_context", None) + if get_context is None: + return None + context = get_context() + if context is None or not context.trace_id: + return None + return (context.trace_id, context.span_id) + + +class FlattenDictRegistry: + """Maps span identity to the dict attributes awaiting flattening.""" + + def __init__(self, max_tracked_spans=_MAX_TRACKED_SPANS): + self._lock = threading.Lock() + self._pending = OrderedDict() + self._max_tracked_spans = max_tracked_spans + + def register(self, span, key, value): + """Park ``value`` under ``key`` for ``span``; last write wins.""" + span_key = _span_key(span) + if span_key is None: + return + with self._lock: + attributes = self._pending.get(span_key) + if attributes is None: + attributes = OrderedDict() + self._pending[span_key] = attributes + attributes[key] = value + self._pending.move_to_end(span_key) + while len(self._pending) > self._max_tracked_spans: + evicted, _ = self._pending.popitem(last=False) + logger.debug( + "Flatten: evicted pending dict attributes for span %s " + "(registry limit %s reached)", + evicted, + self._max_tracked_spans, + ) + + def pop(self, span): + """Remove and return the pending dict attributes for ``span``.""" + span_key = _span_key(span) + if span_key is None: + return {} + with self._lock: + return self._pending.pop(span_key, {}) + + def clear(self): + with self._lock: + self._pending.clear() + + +_REGISTRY = FlattenDictRegistry() + + +def get_registry(): + return _REGISTRY diff --git a/src/harness_sdk/flatten_dict_span_processor.py b/src/harness_sdk/flatten_dict_span_processor.py new file mode 100644 index 0000000..23c6ed8 --- /dev/null +++ b/src/harness_sdk/flatten_dict_span_processor.py @@ -0,0 +1,148 @@ +"""Span processor that expands dict span attributes into dot-notation keys. + +``set_span_attribute("agent", {"action": "generate"})`` parks the dict in +``flatten_dict_registry`` instead of handing it to OTel, which would reject it. +This processor drains the registry at ``on_end`` and writes +``agent.action=generate`` so the backend gets individually queryable +attributes instead of an opaque JSON blob. + +It must be the outermost processor: downstream scrubbing and exclusion logic +matches on attribute keys, so the flattened keys have to exist before those +run. Mutating the ended span works the same way ``GenAiPayloadScrubSpanProcessor`` +relies on: ``ReadableSpan.attributes`` is a read-only view over the attribute +store the concrete SDK ``Span`` still owns at ``on_end`` time. +""" +import json +from typing import Mapping + +from opentelemetry.sdk.trace import SpanProcessor + +from harness_sdk.custom_logger import get_custom_logger +from harness_sdk.flatten_dict_registry import ( + get_registry, + get_flatten_max_depth, + get_flatten_max_leaves, + is_raw_json_enabled, +) + +logger = get_custom_logger(__name__) + +_SCALAR_TYPES = (bool, int, float, str) + + +def _is_scalar(value): + return isinstance(value, _SCALAR_TYPES) + + +def _scalar_kind(value): + # bool is a subclass of int, but OTel treats them as distinct array types. + if isinstance(value, bool): + return bool + if isinstance(value, int): + return int + if isinstance(value, float): + return float + return str + + +def _to_json(value): + try: + return json.dumps(value, default=str) + except (TypeError, ValueError): + return str(value) + + +def _sequence_leaf(value): + """Homogeneous scalar sequences stay arrays; anything else becomes JSON.""" + items = tuple(value) + if not items: + return items + kinds = {_scalar_kind(item) for item in items if _is_scalar(item)} + if len(kinds) == 1 and all(_is_scalar(item) for item in items): + return items + return _to_json(value) + + +def _leaf_value(value): + """Convert a non-mapping value to something OTel accepts, or None to skip.""" + if value is None: + return None + if _is_scalar(value): + return value + if isinstance(value, (list, tuple, set, frozenset)): + return _sequence_leaf(value) + return str(value) + + +def _collect(prefix, mapping, depth, flattened, max_depth, max_leaves): + """Walk ``mapping`` into ``flattened``; returns False once the cap is hit.""" + for key, value in mapping.items(): + if len(flattened) >= max_leaves: + return False + flat_key = f"{prefix}.{key}" + if isinstance(value, Mapping): + if depth < max_depth: + if not _collect(flat_key, value, depth + 1, flattened, max_depth, max_leaves): + return False + else: + flattened[flat_key] = _to_json(value) + continue + leaf = _leaf_value(value) + if leaf is not None: + flattened[flat_key] = leaf + return True + + +class FlattenDictSpanProcessor(SpanProcessor): + """Flattens registered dict attributes onto the span before export.""" + + def __init__(self, processor): + self._processor = processor + + def on_start(self, span, parent_context=None): + self._processor.on_start(span, parent_context) + + def on_end(self, span): + pending = get_registry().pop(span) + if pending: + try: + self._flatten(span, pending) + except Exception as err: # pylint: disable=W0703 + logger.debug( + "Flatten: failed to flatten dict attributes on span %s: %s", + getattr(span, "name", None), + err, + ) + self._processor.on_end(span) + + @staticmethod + def _flatten(span, pending): + attributes = getattr(span, "_attributes", None) + if attributes is None: + return + raw_json = is_raw_json_enabled() + max_depth = get_flatten_max_depth() + max_leaves = get_flatten_max_leaves() + for root_key, value in pending.items(): + flattened = {} + if not _collect(root_key, value, 1, flattened, max_depth, max_leaves): + logger.debug( + "Flatten: dict attribute %r on span %s exceeded %s leaf " + "attributes; remaining entries dropped", + root_key, + getattr(span, "name", None), + max_leaves, + ) + for flat_key, leaf in flattened.items(): + # An explicitly set attribute always wins over a flattened one. + if flat_key in attributes: + continue + attributes[flat_key] = leaf + if raw_json and root_key not in attributes: + attributes[root_key] = _to_json(value) + + def force_flush(self, timeout_millis=30000): + return self._processor.force_flush(timeout_millis) + + def shutdown(self): + return self._processor.shutdown() diff --git a/src/harness_sdk/plugins/builtin/pipeline.py b/src/harness_sdk/plugins/builtin/pipeline.py index e665500..fcd18d2 100644 --- a/src/harness_sdk/plugins/builtin/pipeline.py +++ b/src/harness_sdk/plugins/builtin/pipeline.py @@ -9,6 +9,8 @@ from harness_sdk.env import is_env_flag_enabled from harness_sdk.excluded_by_attribute_span_processor import ExcludeByAttributeSpanProcessor from harness_sdk.db_control_span_processor import DbControlSpanProcessor +from harness_sdk.flatten_dict_registry import is_flatten_enabled +from harness_sdk.flatten_dict_span_processor import FlattenDictSpanProcessor from harness_sdk.gen_ai_payload_scrub_span_processor import GenAiPayloadScrubSpanProcessor logger = get_custom_logger(__name__) @@ -43,10 +45,14 @@ def create_span_processors(self, config: Any) -> List[SpanProcessor]: excluded_value="nospan", ) db_control_processor = DbControlSpanProcessor(filter_processor) - # Outermost: scrub GenAI payload attributes before any other on_end - # logic (control evaluation, filtering, batching) sees the span. + # Scrub GenAI payload attributes before control evaluation, filtering + # and batching see the span. scrub_processor = GenAiPayloadScrubSpanProcessor(db_control_processor) - return [scrub_processor] + if not is_flatten_enabled(): + return [scrub_processor] + # Outermost: dict attributes must be expanded into their flat keys + # before scrubbing and exclusion match on attribute keys. + return [FlattenDictSpanProcessor(scrub_processor)] def shutdown(self) -> None: pass diff --git a/src/harness_sdk/span_enrichment.py b/src/harness_sdk/span_enrichment.py index ae79761..8209404 100644 --- a/src/harness_sdk/span_enrichment.py +++ b/src/harness_sdk/span_enrichment.py @@ -1,22 +1,35 @@ """Public helpers for enriching the current OpenTelemetry span.""" -from typing import Mapping +from typing import Any, Mapping, Union from opentelemetry import trace from opentelemetry.util.types import AttributeValue +from harness_sdk.flatten_dict_registry import get_registry, is_flatten_enabled -def set_span_attribute(key: str, value: AttributeValue) -> None: +EnrichmentValue = Union[AttributeValue, Mapping[str, Any]] + + +def set_span_attribute(key: str, value: EnrichmentValue) -> None: """Set one attribute on the current recording span.""" span = trace.get_current_span() if span.is_recording(): - span.set_attribute(key, value) + _set(span, key, value) -def set_span_attributes(attributes: Mapping[str, AttributeValue]) -> None: +def set_span_attributes(attributes: Mapping[str, EnrichmentValue]) -> None: """Set attributes on the current recording span.""" span = trace.get_current_span() if not span.is_recording(): return for key, value in attributes.items(): - span.set_attribute(key, value) + _set(span, key, value) + + +def _set(span, key: str, value: EnrichmentValue) -> None: + # Dict values are parked for FlattenDictSpanProcessor to expand into + # dot-notation keys at span end; nothing is serialized on this thread. + if isinstance(value, Mapping) and is_flatten_enabled(): + get_registry().register(span, key, value) + return + span.set_attribute(key, value) diff --git a/test/flatten_dict_span_processor_test.py b/test/flatten_dict_span_processor_test.py new file mode 100644 index 0000000..e4ec0a8 --- /dev/null +++ b/test/flatten_dict_span_processor_test.py @@ -0,0 +1,292 @@ +"""Tests for FlattenDictSpanProcessor and the dict attribute registry.""" +import json + +import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.trace import INVALID_SPAN + +from harness_sdk.flatten_dict_registry import ( + DEFAULT_MAX_LEAVES, + FLATTEN_ENABLED_ENV, + FLATTEN_MAX_DEPTH_ENV, + FLATTEN_MAX_LEAVES_ENV, + FLATTEN_RAW_JSON_ENV, + FlattenDictRegistry, + get_flatten_max_depth, + get_flatten_max_leaves, + get_registry, + is_flatten_enabled, + is_raw_json_enabled, +) +from harness_sdk.flatten_dict_span_processor import FlattenDictSpanProcessor + + +class RecordingProcessor: + """Innermost processor capturing what reaches export.""" + + def __init__(self): + self.started = [] + self.ended = [] + self.flushed = False + self.shutdown_called = False + + def on_start(self, span, parent_context=None): + self.started.append(span) + + def on_end(self, span): + self.ended.append(span) + + def force_flush(self, timeout_millis=30000): + self.flushed = True + return True + + def shutdown(self): + self.shutdown_called = True + + +@pytest.fixture(autouse=True) +def _clean_registry(): + get_registry().clear() + yield + get_registry().clear() + + +@pytest.fixture(name="downstream") +def _downstream(): + return RecordingProcessor() + + +def _ended_attributes(downstream, attributes): + """Run a span through the processor and return its exported attributes.""" + processor = FlattenDictSpanProcessor(downstream) + provider = TracerProvider() + provider.add_span_processor(processor) + with provider.get_tracer(__name__).start_as_current_span("span") as span: + for key, value in attributes.items(): + get_registry().register(span, key, value) + assert len(downstream.ended) == 1 + return dict(downstream.ended[0].attributes or {}) + + +def test_flat_dict_becomes_dot_notation_keys(downstream): + attributes = _ended_attributes( + downstream, {"agent": {"action": "generate", "name": "devops"}} + ) + + assert attributes["agent.action"] == "generate" + assert attributes["agent.name"] == "devops" + assert "agent" not in attributes + + +def test_native_types_are_preserved(downstream): + attributes = _ended_attributes( + downstream, + {"agent": {"name": "devops", "active": True, "retries": 3, "score": 1.5}}, + ) + + assert attributes["agent.name"] == "devops" + assert attributes["agent.active"] is True + assert attributes["agent.retries"] == 3 + assert attributes["agent.score"] == 1.5 + + +def test_nesting_up_to_max_depth_is_flattened(downstream): + attributes = _ended_attributes( + downstream, {"agent": {"model": {"provider": {"name": "vertex"}}}} + ) + + assert attributes["agent.model.provider.name"] == "vertex" + + +def test_dict_beyond_max_depth_is_json_encoded(downstream): + attributes = _ended_attributes( + downstream, {"agent": {"a": {"b": {"c": {"d": 1}}}}} + ) + + assert json.loads(attributes["agent.a.b.c"]) == {"d": 1} + + +def test_none_leaf_is_skipped(downstream): + attributes = _ended_attributes( + downstream, {"agent": {"name": "devops", "parent": None}} + ) + + assert attributes["agent.name"] == "devops" + assert "agent.parent" not in attributes + + +def test_non_scalar_leaf_falls_back_to_string(downstream): + class Model: + def __str__(self): + return "gemini-2.0" + + attributes = _ended_attributes(downstream, {"agent": {"model": Model()}}) + + assert attributes["agent.model"] == "gemini-2.0" + + +def test_scalar_list_becomes_array_attribute(downstream): + attributes = _ended_attributes( + downstream, {"agent": {"tools": ["search", "shell"]}} + ) + + assert attributes["agent.tools"] == ("search", "shell") + + +def test_mixed_and_dict_lists_are_json_encoded(downstream): + attributes = _ended_attributes( + downstream, + { + "agent": { + "mixed": ["search", 2], + "steps": [{"name": "plan"}, {"name": "act"}], + } + }, + ) + + assert json.loads(attributes["agent.mixed"]) == ["search", 2] + assert json.loads(attributes["agent.steps"]) == [{"name": "plan"}, {"name": "act"}] + + +def test_explicit_attribute_wins_over_flattened_key(downstream): + processor = FlattenDictSpanProcessor(downstream) + provider = TracerProvider() + provider.add_span_processor(processor) + with provider.get_tracer(__name__).start_as_current_span("span") as span: + span.set_attribute("agent.name", "explicit") + get_registry().register(span, "agent", {"name": "flattened"}) + + assert downstream.ended[0].attributes["agent.name"] == "explicit" + + +def test_leaf_cap_stops_flattening(downstream): + oversized = {f"key{index}": index for index in range(DEFAULT_MAX_LEAVES + 10)} + + attributes = _ended_attributes(downstream, {"agent": oversized}) + + flattened = [key for key in attributes if key.startswith("agent.")] + assert len(flattened) == DEFAULT_MAX_LEAVES + + +def test_raw_json_flag_also_sets_original_key(downstream, monkeypatch): + monkeypatch.setenv(f"HARNESS_{FLATTEN_RAW_JSON_ENV}", "true") + + attributes = _ended_attributes(downstream, {"agent": {"name": "devops"}}) + + assert attributes["agent.name"] == "devops" + assert json.loads(attributes["agent"]) == {"name": "devops"} + + +def test_span_without_pending_dicts_passes_through(downstream): + processor = FlattenDictSpanProcessor(downstream) + provider = TracerProvider() + provider.add_span_processor(processor) + with provider.get_tracer(__name__).start_as_current_span("span") as span: + span.set_attribute("agent.name", "devops") + + assert len(downstream.started) == 1 + assert dict(downstream.ended[0].attributes) == {"agent.name": "devops"} + + +def test_lifecycle_calls_delegate_to_wrapped_processor(downstream): + processor = FlattenDictSpanProcessor(downstream) + + assert processor.force_flush() is True + processor.shutdown() + + assert downstream.flushed + assert downstream.shutdown_called + + +def test_registry_pop_is_one_shot(): + registry = FlattenDictRegistry() + with TracerProvider().get_tracer(__name__).start_as_current_span("span") as span: + registry.register(span, "agent", {"name": "devops"}) + + assert registry.pop(span) == {"agent": {"name": "devops"}} + assert registry.pop(span) == {} + + +def test_registry_last_write_wins_per_key(): + registry = FlattenDictRegistry() + with TracerProvider().get_tracer(__name__).start_as_current_span("span") as span: + registry.register(span, "agent", {"name": "first"}) + registry.register(span, "agent", {"name": "second"}) + + assert registry.pop(span) == {"agent": {"name": "second"}} + + +def test_registry_evicts_oldest_spans_over_limit(): + registry = FlattenDictRegistry(max_tracked_spans=1) + tracer = TracerProvider().get_tracer(__name__) + with tracer.start_as_current_span("first") as first: + registry.register(first, "agent", {"name": "first"}) + with tracer.start_as_current_span("second") as second: + registry.register(second, "agent", {"name": "second"}) + + assert registry.pop(first) == {} + assert registry.pop(second) == {"agent": {"name": "second"}} + + +def test_registry_ignores_non_recording_span(): + registry = FlattenDictRegistry() + + registry.register(INVALID_SPAN, "agent", {"name": "devops"}) + + assert registry.pop(INVALID_SPAN) == {} + + +def test_flatten_enabled_defaults_to_true(monkeypatch): + for prefix in ("HARNESS_", "HA_", "AT_", "TA_"): + monkeypatch.delenv(f"{prefix}{FLATTEN_ENABLED_ENV}", raising=False) + monkeypatch.delenv(f"{prefix}{FLATTEN_RAW_JSON_ENV}", raising=False) + + assert is_flatten_enabled() is True + assert is_raw_json_enabled() is False + + +def test_flatten_disabled_only_when_explicitly_false(monkeypatch): + monkeypatch.setenv(f"HARNESS_{FLATTEN_ENABLED_ENV}", "false") + assert is_flatten_enabled() is False + + monkeypatch.setenv(f"HARNESS_{FLATTEN_ENABLED_ENV}", "FALSE") + assert is_flatten_enabled() is False + + monkeypatch.setenv(f"HARNESS_{FLATTEN_ENABLED_ENV}", "true") + assert is_flatten_enabled() is True + + +def test_flatten_stays_enabled_for_non_false_values(monkeypatch): + for value in ("", "1", "yes", "on", "garbage"): + monkeypatch.setenv(f"HARNESS_{FLATTEN_ENABLED_ENV}", value) + assert is_flatten_enabled() is True + + +def test_flatten_limits_default(monkeypatch): + for prefix in ("HARNESS_", "HA_", "AT_", "TA_"): + monkeypatch.delenv(f"{prefix}{FLATTEN_MAX_DEPTH_ENV}", raising=False) + monkeypatch.delenv(f"{prefix}{FLATTEN_MAX_LEAVES_ENV}", raising=False) + assert get_flatten_max_depth() == 3 + assert get_flatten_max_leaves() == 32 + + +def test_flatten_limits_from_env(monkeypatch, downstream): + monkeypatch.setenv(f"HARNESS_{FLATTEN_MAX_DEPTH_ENV}", "2") + + assert get_flatten_max_depth() == 2 + + depth_attrs = _ended_attributes( + downstream, + {"agent": {"l1": {"l2": {"l3": {"l4": "deep"}}}}}, + ) + assert depth_attrs["agent.l1.l2"] == json.dumps({"l3": {"l4": "deep"}}) + + +def test_flatten_leaf_limit_from_env(monkeypatch, downstream): + monkeypatch.setenv(f"HARNESS_{FLATTEN_MAX_LEAVES_ENV}", "5") + + leaf_attrs = _ended_attributes( + downstream, + {"agent": {f"key{index}": index for index in range(10)}}, + ) + assert len([key for key in leaf_attrs if key.startswith("agent.")]) == 5 diff --git a/test/plugins/builtin/test_pipeline.py b/test/plugins/builtin/test_pipeline.py index f7674cc..ec3f660 100644 --- a/test/plugins/builtin/test_pipeline.py +++ b/test/plugins/builtin/test_pipeline.py @@ -2,11 +2,13 @@ from harness_sdk.config.config import Config from harness_sdk.db_control_span_processor import DbControlSpanProcessor from harness_sdk.excluded_by_attribute_span_processor import ExcludeByAttributeSpanProcessor +from harness_sdk.flatten_dict_registry import FLATTEN_ENABLED_ENV +from harness_sdk.flatten_dict_span_processor import FlattenDictSpanProcessor from harness_sdk.gen_ai_payload_scrub_span_processor import GenAiPayloadScrubSpanProcessor from harness_sdk.plugins.builtin.pipeline import BuiltinPipelinePlugin -def test_scrub_processor_wraps_chain_as_outermost_layer(monkeypatch): +def _build_processors(monkeypatch): # Force the real OTLP-export branch (skip the console-exporter early-return) # so the full processor chain gets assembled. monkeypatch.delenv("HA_ENABLE_CONSOLE_SPAN_EXPORTER", raising=False) @@ -14,8 +16,23 @@ def test_scrub_processor_wraps_chain_as_outermost_layer(monkeypatch): config = Config() plugin = BuiltinPipelinePlugin() plugin.on_init(config) + return plugin.create_span_processors(config) + + +def test_flatten_processor_wraps_chain_as_outermost_layer(monkeypatch): + processors = _build_processors(monkeypatch) + + assert len(processors) == 1 + flatten_processor = processors[0] + assert isinstance(flatten_processor, FlattenDictSpanProcessor) + # pylint: disable=protected-access + assert isinstance(flatten_processor._processor, GenAiPayloadScrubSpanProcessor) + + +def test_scrub_processor_wraps_chain_as_outermost_layer(monkeypatch): + monkeypatch.setenv(f"HARNESS_{FLATTEN_ENABLED_ENV}", "false") - processors = plugin.create_span_processors(config) + processors = _build_processors(monkeypatch) assert len(processors) == 1 scrub_processor = processors[0] diff --git a/test/span_enrichment_test.py b/test/span_enrichment_test.py index 530328c..8178837 100644 --- a/test/span_enrichment_test.py +++ b/test/span_enrichment_test.py @@ -1,14 +1,24 @@ import asyncio +import pytest from opentelemetry.sdk.trace import TracerProvider from harness_sdk import set_span_attribute, set_span_attributes +from harness_sdk.flatten_dict_registry import FLATTEN_ENABLED_ENV, get_registry +from harness_sdk.flatten_dict_span_processor import FlattenDictSpanProcessor def _tracer(): return TracerProvider().get_tracer(__name__) +@pytest.fixture(autouse=True) +def _clean_registry(): + get_registry().clear() + yield + get_registry().clear() + + def test_set_span_attribute_updates_current_recording_span(): with _tracer().start_as_current_span("parent") as span: set_span_attribute("request.client.name", "acme") @@ -73,3 +83,56 @@ async def enrich(): asyncio.run(enrich()) assert span.attributes["agent.action.type"] == "search" + + +def test_dict_value_is_deferred_not_set_on_span(): + with _tracer().start_as_current_span("parent") as span: + set_span_attribute("agent", {"action": "generate"}) + + assert not span.attributes + assert get_registry().pop(span) == {"agent": {"action": "generate"}} + + +def test_dict_value_is_flattened_at_span_end(): + class Capture: + def __init__(self): + self.ended = [] + + def on_start(self, span, parent_context=None): + pass + + def on_end(self, span): + self.ended.append(span) + + def force_flush(self, timeout_millis=30000): + return True + + def shutdown(self): + pass + + capture = Capture() + provider = TracerProvider() + provider.add_span_processor(FlattenDictSpanProcessor(capture)) + with provider.get_tracer(__name__).start_as_current_span("parent"): + set_span_attributes({ + "agent": {"action": "generate", "name": "devops"}, + "custom.retry.count": 2, + }) + + attributes = dict(capture.ended[0].attributes) + assert attributes == { + "custom.retry.count": 2, + "agent.action": "generate", + "agent.name": "devops", + } + + +def test_dict_value_is_rejected_when_flatten_disabled(monkeypatch): + monkeypatch.setenv(f"HARNESS_{FLATTEN_ENABLED_ENV}", "false") + + with _tracer().start_as_current_span("parent") as span: + set_span_attribute("agent", {"action": "generate"}) + + # OTel drops the unsupported mapping value itself; nothing is deferred. + assert not span.attributes + assert get_registry().pop(span) == {}