Mark the outermost GenAI span of traces with a non-GenAI root

data-prepper-plugins/otel-trace-raw-processor/README.md

@@ -23,8 +23,16 @@ This processor automatically enriches GenAI agent traces. The enrichment is alwa
 
 **What it does:**
 1. **Vendor attribute normalization** — Maps vendor-specific attributes (OpenInference, OpenLLMetry) to [OTel GenAI Semantic Conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/). Original attributes are preserved. See [`genai-attribute-mappings.yaml`](src/main/resources/genai-attribute-mappings.yaml) for the full mapping table.
-2. **Root span enrichment** — Propagates `gen_ai.system`, `gen_ai.provider.name`, `gen_ai.agent.name`, `gen_ai.request.model`, and `gen_ai.operation.name` from child spans to root (first-child-wins, skip-if-present). Aggregates `gen_ai.usage.input_tokens` and `gen_ai.usage.output_tokens` across children (sum).
-3. **Flattened sub-key stripping** — Removes flattened sub-keys (e.g. `llm.input_messages.0.message.content`) that conflict with parent string values, preventing OpenSearch mapping failures.
+2. **Root span enrichment** — Propagates `gen_ai.system`, `gen_ai.provider.name`, `gen_ai.agent.name`, and `gen_ai.request.model` from child spans to root (skip-if-present). Values come from the outermost GenAI span first (earliest start time), then other GenAI spans, then the remaining children. Aggregates `gen_ai.usage.input_tokens` and `gen_ai.usage.output_tokens` across children, counting only spans that have no ancestor reporting usage: agent spans such as `invoke_agent` often report the total of their own LLM calls, which would otherwise be counted twice.
+3. **Agent root marker** — Sets `derived.gen_ai.is_root: true` on each outermost GenAI span of a trace: a span with `gen_ai.operation.name` and no GenAI ancestor. That is the root itself when it is a GenAI span, or the agent span below a non-GenAI root (for example an agent served behind an HTTP handler). It also aggregates the token usage of that span's descendants onto it (with the same rule), unless it already has `gen_ai.usage.input_tokens`.
+   * Several top-level agents under one root are each marked.
+   * An agent that calls another agent over HTTP or RPC in the same trace (`invoke_agent` > HTTP client > HTTP server > `invoke_agent`) gets one mark, on the calling agent, because the called agent has a GenAI ancestor. If the called service's spans are flushed after `trace_flush_interval` without the caller's spans, the called agent is marked too, and consumers that list marked spans show it as a separate entry.
+   * Spans that arrive after their trace's root has been processed are enriched on their own. The processor remembers, for `trace_group_cache_ttl`, which span ids of a trace are already covered by a GenAI span or a marker, so such a span is not mistaken for the outermost one. A span that arrives later than that may be marked as well.
+   * If a trace's children are flushed before a GenAI root arrives, both the root and the outermost child can end up marked. Consumers should prefer the marked span with no parent.
+   * A marker set at the source with any value, including `false`, is kept, and the processor then marks nothing else in that part of the trace. Set it to `false` to opt a trace out.
+   * Spans flushed after `trace_flush_interval` without their root are still marked, because the marker only needs the spans below it.
+   * If any span in the trace already carries `derived.gen_ai.is_root` (set by an SDK or by hand), no spans are marked by the processor. To choose the root yourself, set `derived.gen_ai.is_root` to `true` on that span when you create it.
+4. **Flattened sub-key stripping** — Removes flattened sub-keys (e.g. `llm.input_messages.0.message.content`) that conflict with parent string values, preventing OpenSearch mapping failures.
 
 **Important:** The enrichment matches attributes by their original OTel key names (e.g. `http.method`). This requires the source to use `output_format: otel` so that attribute keys are not transformed. The default `opensearch` format rewrites keys (e.g. `http.method` → `span.attributes.http@method`), which prevents the enrichment from working. See the [otel-trace-source README](../otel-trace-source/README.md#output-format) for details.
 

data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/GenAiCoverage.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.processor.oteltrace;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Remembers, per trace, the span ids already covered by a GenAI span or an agent root marker (the
+ * span itself or one of its ancestors). Spans that arrive after their trace's root has been
+ * processed are enriched without their ancestors; this lets the agent root marker tell whether
+ * such a span is really outermost.
+ */
+final class GenAiCoverage {
+
+    /** Keeps nothing: every parent outside the batch counts as uncovered. */
+    static final GenAiCoverage NONE = new GenAiCoverage(null);
+
+    private final Cache<String, Set<String>> coveredByTrace;
+
+    private GenAiCoverage(final Cache<String, Set<String>> coveredByTrace) {
+        this.coveredByTrace = coveredByTrace;
+    }
+
+    static GenAiCoverage create(final long maxSpanIds, final Duration timeToLive) {
+        return new GenAiCoverage(Caffeine.newBuilder()
+                .maximumWeight(maxSpanIds)
+                .<String, Set<String>>weigher((traceId, spanIds) -> spanIds.size())
+                .expireAfterWrite(timeToLive)
+                .build());
+    }
+
+    Set<String> covered(final String traceId) {
+        if (coveredByTrace == null || traceId == null) {
+            return Collections.emptySet();
+        }
+        final Set<String> spanIds = coveredByTrace.getIfPresent(traceId);
+        return spanIds == null ? Collections.emptySet() : spanIds;
+    }
+
+    void record(final String traceId, final Collection<String> spanIds) {
+        if (coveredByTrace == null || traceId == null || spanIds.isEmpty()) {
+            return;
+        }
+        coveredByTrace.asMap().merge(traceId, new HashSet<>(spanIds), (existing, added) -> {
+            final Set<String> merged = new HashSet<>(existing);
+            merged.addAll(added);
+            return merged;
+        });
+    }
+
+    long estimatedSize() {
+        return coveredByTrace == null ? 0 : coveredByTrace.estimatedSize();
+    }
+}