From bb60f28c14b82de690f5ac00231100573a15bb7f Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Mon, 3 Aug 2026 16:37:26 +0530 Subject: [PATCH 1/3] feat: flatten dict values passed to set_span_attribute MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Passing a dict to set_span_attribute silently dropped the value: OTel rejects mapping attribute values outright. Callers had to json.dumps by hand, which left the backend with an opaque string that cannot be queried or aggregated per field. Dict values are now accepted and expanded into dot-notation attributes (agent.action=generate) so every leaf stays individually queryable. Nothing is serialized on the caller's thread: the dict is parked in a thread-safe registry keyed by span context and flattened by a new FlattenDictSpanProcessor at on_end. That processor sits outermost in the pipeline so GenAI scrubbing and attribute-based exclusion match on the flattened keys. Bounded by design — depth 3, 32 leaf attributes per dict, explicit attributes win on key collision — and disabled with HARNESS_SPAN_ATTRIBUTE_FLATTEN_ENABLED=false. Co-authored-by: Cursor --- AGENTS.md | 2 + README.md | 29 ++ src/harness_sdk/agent_init.py | 8 +- src/harness_sdk/flatten_dict_registry.py | 97 +++++++ .../flatten_dict_span_processor.py | 147 ++++++++++ src/harness_sdk/plugins/builtin/pipeline.py | 12 +- src/harness_sdk/span_enrichment.py | 23 +- test/flatten_dict_span_processor_test.py | 251 ++++++++++++++++++ test/plugins/builtin/test_pipeline.py | 21 +- test/span_enrichment_test.py | 63 +++++ 10 files changed, 641 insertions(+), 12 deletions(-) create mode 100644 src/harness_sdk/flatten_dict_registry.py create mode 100644 src/harness_sdk/flatten_dict_span_processor.py create mode 100644 test/flatten_dict_span_processor_test.py diff --git a/AGENTS.md b/AGENTS.md index 5a86d70..edeac56 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,6 +38,8 @@ 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) | +| `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..93736a3 100644 --- a/README.md +++ b/README.md @@ -168,6 +168,35 @@ 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 3 levels | JSON string at the depth-3 key | +| Flattened key already set on the span | skipped — the explicit value wins | + +At most 32 leaf attributes are emitted per dictionary; the rest are dropped with a debug +log. The original key (`agent`) is not set unless +`HARNESS_SPAN_ATTRIBUTE_FLATTEN_RAW_JSON=true`, which additionally stores the whole dict +as JSON there. Set `HARNESS_SPAN_ATTRIBUTE_FLATTEN_ENABLED=false` to turn flattening 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..50ab4b9 --- /dev/null +++ b/src/harness_sdk/flatten_dict_registry.py @@ -0,0 +1,97 @@ +"""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" + +# 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 unless explicitly disabled.""" + if get_env_value(FLATTEN_ENABLED_ENV) is None: + return True + return is_env_flag_enabled(FLATTEN_ENABLED_ENV) + + +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 _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..df9633e --- /dev/null +++ b/src/harness_sdk/flatten_dict_span_processor.py @@ -0,0 +1,147 @@ +"""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, + is_raw_json_enabled, +) + +logger = get_custom_logger(__name__) + +MAX_DEPTH = 3 +MAX_LEAF_ATTRIBUTES = 32 + +_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): + """Walk ``mapping`` into ``flattened``; returns False once the cap is hit.""" + for key, value in mapping.items(): + if len(flattened) >= MAX_LEAF_ATTRIBUTES: + return False + flat_key = f"{prefix}.{key}" + if isinstance(value, Mapping): + if depth < MAX_DEPTH: + if not _collect(flat_key, value, depth + 1, flattened): + 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() + for root_key, value in pending.items(): + flattened = {} + if not _collect(root_key, value, 1, flattened): + logger.debug( + "Flatten: dict attribute %r on span %s exceeded %s leaf " + "attributes; remaining entries dropped", + root_key, + getattr(span, "name", None), + MAX_LEAF_ATTRIBUTES, + ) + 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..10472a9 --- /dev/null +++ b/test/flatten_dict_span_processor_test.py @@ -0,0 +1,251 @@ +"""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 ( + FLATTEN_ENABLED_ENV, + FLATTEN_RAW_JSON_ENV, + FlattenDictRegistry, + get_registry, + is_flatten_enabled, + is_raw_json_enabled, +) +from harness_sdk.flatten_dict_span_processor import ( + MAX_LEAF_ATTRIBUTES, + 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(MAX_LEAF_ATTRIBUTES + 10)} + + attributes = _ended_attributes(downstream, {"agent": oversized}) + + flattened = [key for key in attributes if key.startswith("agent.")] + assert len(flattened) == MAX_LEAF_ATTRIBUTES + + +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}", "true") + assert is_flatten_enabled() is True 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) == {} From 8e381b9963ad871ea251eb543964ba91fd46a620 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Wed, 5 Aug 2026 17:40:01 +0530 Subject: [PATCH 2/3] fix: keep dict flattening enabled by default (opt-out only) Only explicit false disables flattening; unset or any other env value leaves it on so customers get queryable flat keys without configuring HARNESS_SPAN_ATTRIBUTE_FLATTEN_ENABLED=true. Co-authored-by: Cursor --- README.md | 5 +++-- src/harness_sdk/flatten_dict_registry.py | 7 ++++--- test/flatten_dict_span_processor_test.py | 9 +++++++++ 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 93736a3..e676d75 100644 --- a/README.md +++ b/README.md @@ -194,8 +194,9 @@ before export. Flattening rules: At most 32 leaf attributes are emitted per dictionary; the rest are dropped with a debug log. The original key (`agent`) is not set unless `HARNESS_SPAN_ATTRIBUTE_FLATTEN_RAW_JSON=true`, which additionally stores the whole dict -as JSON there. Set `HARNESS_SPAN_ATTRIBUTE_FLATTEN_ENABLED=false` to turn flattening off, -in which case OTel rejects dictionary values as it did before. +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 diff --git a/src/harness_sdk/flatten_dict_registry.py b/src/harness_sdk/flatten_dict_registry.py index 50ab4b9..10d6802 100644 --- a/src/harness_sdk/flatten_dict_registry.py +++ b/src/harness_sdk/flatten_dict_registry.py @@ -27,10 +27,11 @@ def is_flatten_enabled(): - """Dict flattening is on unless explicitly disabled.""" - if get_env_value(FLATTEN_ENABLED_ENV) is None: + """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 is_env_flag_enabled(FLATTEN_ENABLED_ENV) + return value.strip().lower() != "false" def is_raw_json_enabled(): diff --git a/test/flatten_dict_span_processor_test.py b/test/flatten_dict_span_processor_test.py index 10472a9..0d58a7d 100644 --- a/test/flatten_dict_span_processor_test.py +++ b/test/flatten_dict_span_processor_test.py @@ -247,5 +247,14 @@ 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 From 3c64a9791da1d9b1d444203ab3ce5c71a427d1a4 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Thu, 6 Aug 2026 13:57:27 +0530 Subject: [PATCH 3/3] feat: make flatten depth and leaf limits configurable via env Add HARNESS_SPAN_ATTRIBUTE_FLATTEN_MAX_DEPTH and MAX_LEAVES (defaults 3 and 32) so callers can tune safety rails without disabling flattening. Co-authored-by: Cursor --- AGENTS.md | 4 +- README.md | 7 +-- src/harness_sdk/flatten_dict_registry.py | 24 ++++++++++ .../flatten_dict_span_processor.py | 19 ++++---- test/flatten_dict_span_processor_test.py | 44 ++++++++++++++++--- 5 files changed, 79 insertions(+), 19 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index edeac56..f033ded 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,7 +38,9 @@ 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) | +| `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 e676d75..382c5e7 100644 --- a/README.md +++ b/README.md @@ -188,11 +188,12 @@ before export. Flattening rules: | 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 3 levels | JSON string at the depth-3 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 32 leaf attributes are emitted per dictionary; the rest are dropped with a debug -log. The original key (`agent`) is not set unless +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 diff --git a/src/harness_sdk/flatten_dict_registry.py b/src/harness_sdk/flatten_dict_registry.py index 10d6802..8364172 100644 --- a/src/harness_sdk/flatten_dict_registry.py +++ b/src/harness_sdk/flatten_dict_registry.py @@ -20,6 +20,11 @@ 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. @@ -39,6 +44,25 @@ def is_raw_json_enabled(): 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: diff --git a/src/harness_sdk/flatten_dict_span_processor.py b/src/harness_sdk/flatten_dict_span_processor.py index df9633e..23c6ed8 100644 --- a/src/harness_sdk/flatten_dict_span_processor.py +++ b/src/harness_sdk/flatten_dict_span_processor.py @@ -20,14 +20,13 @@ 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__) -MAX_DEPTH = 3 -MAX_LEAF_ATTRIBUTES = 32 - _SCALAR_TYPES = (bool, int, float, str) @@ -75,15 +74,15 @@ def _leaf_value(value): return str(value) -def _collect(prefix, mapping, depth, flattened): +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_LEAF_ATTRIBUTES: + 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): + 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) @@ -122,15 +121,17 @@ def _flatten(span, pending): 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): + 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_LEAF_ATTRIBUTES, + max_leaves, ) for flat_key, leaf in flattened.items(): # An explicitly set attribute always wins over a flattened one. diff --git a/test/flatten_dict_span_processor_test.py b/test/flatten_dict_span_processor_test.py index 0d58a7d..e4ec0a8 100644 --- a/test/flatten_dict_span_processor_test.py +++ b/test/flatten_dict_span_processor_test.py @@ -6,17 +6,19 @@ 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 ( - MAX_LEAF_ATTRIBUTES, - FlattenDictSpanProcessor, -) +from harness_sdk.flatten_dict_span_processor import FlattenDictSpanProcessor class RecordingProcessor: @@ -158,12 +160,12 @@ def test_explicit_attribute_wins_over_flattened_key(downstream): def test_leaf_cap_stops_flattening(downstream): - oversized = {f"key{index}": index for index in range(MAX_LEAF_ATTRIBUTES + 10)} + 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) == MAX_LEAF_ATTRIBUTES + assert len(flattened) == DEFAULT_MAX_LEAVES def test_raw_json_flag_also_sets_original_key(downstream, monkeypatch): @@ -258,3 +260,33 @@ 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