LLM Obs SDK use context API for parent children span linkage

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java

@@ -1,10 +1,12 @@
 package datadog.trace.llmobs.domain;
 
+import datadog.context.ContextScope;
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import java.util.Collections;
@@ -29,13 +31,15 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private static final String OUTPUT = LLMOBS_TAG_PREFIX + "output";
   private static final String SPAN_KIND = LLMOBS_TAG_PREFIX + Tags.SPAN_KIND;
   private static final String METADATA = LLMOBS_TAG_PREFIX + LLMObsTags.METADATA;
+  private static final String PARENT_ID_TAG_INTERNAL = "parent_id";
 
   private static final String LLM_OBS_INSTRUMENTATION_NAME = "llmobs";
 
   private static final Logger LOGGER = LoggerFactory.getLogger(DDLLMObsSpan.class);
 
   private final AgentSpan span;
   private final String spanKind;
+  private final ContextScope scope;
 
   private boolean finished = false;
 
@@ -63,6 +67,24 @@ public DDLLMObsSpan(
     if (sessionId != null && !sessionId.isEmpty()) {
       this.span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
     }
+
+    AgentSpanContext parent = LLMObsState.getLLMObsParentContext();
+    String parentSpanID = LLMObsState.ROOT_SPAN_ID;
+    if (null != parent) {
+      if (parent.getTraceId() != this.span.getTraceId()) {
+        LOGGER.error(
+            "trace ID mismatch, retrieved parent from context trace_id={}, span_id={}, started span trace_id={}, span_id={}",
+            parent.getTraceId(),
+            parent.getSpanId(),
+            this.span.getTraceId(),
+            this.span.getSpanId());
+      } else {
+        parentSpanID = String.valueOf(parent.getSpanId());
+      }
+    }
+    this.span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID);
+    this.scope = LLMObsState.attach();
+    LLMObsState.setLLMObsParentContext(this.span.context());
   }
 
   @Override
@@ -271,6 +293,7 @@ public void finish() {
       return;
     }
     this.span.finish();
+    this.scope.close();
     this.finished = true;
   }
 }

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsState.java

@@ -0,0 +1,37 @@
+package datadog.trace.llmobs.domain;
+
+import datadog.context.Context;
+import datadog.context.ContextKey;
+import datadog.context.ContextScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+
+public class LLMObsState {
+  public static final String ROOT_SPAN_ID = "undefined";
+
+  private static final ContextKey<LLMObsState> CONTEXT_KEY = ContextKey.named("llmobs_span");
+
+  private AgentSpanContext parentSpanID;
+
+  public static ContextScope attach() {
+    return Context.current().with(CONTEXT_KEY, new LLMObsState()).attach();
+  }
+
+  private static LLMObsState fromContext() {
+    return Context.current().get(CONTEXT_KEY);
+  }
+
+  public static AgentSpanContext getLLMObsParentContext() {
+    LLMObsState state = fromContext();
+    if (state != null) {
+      return state.parentSpanID;
+    }
+    return null;
+  }
+
+  public static void setLLMObsParentContext(AgentSpanContext llmObsParentContext) {
+    LLMObsState state = fromContext();
+    if (state != null) {
+      state.parentSpanID = llmObsParentContext;
+    }
+  }
+}

dd-trace-api/build.gradle

@@ -51,6 +51,7 @@ excludedClassesCoverage += [
   'datadog.trace.context.NoopTraceScope',
   'datadog.trace.payloadtags.PayloadTagsData',
   'datadog.trace.payloadtags.PayloadTagsData.PathAndValue',
+  'datadog.trace.api.llmobs.LLMObsTags',
 ]
 
 description = 'dd-trace-api'

dd-trace-api/src/main/java/datadog/trace/api/DDSpanTypes.java

@@ -36,6 +36,7 @@ public class DDSpanTypes {
   public static final String PROTOBUF = "protobuf";
 
   public static final String MULE = "mule";
+
   public static final String VALKEY = "valkey";
   public static final String WEBSOCKET = "websocket";
 

dd-trace-core/build.gradle

@@ -37,6 +37,8 @@ excludedClassesCoverage += [
   // Interface with an empty defender method
   'datadog.trace.core.propagation.HttpCodec.Extractor',
   'datadog.trace.core.flare.*',
+  'datadog.trace.llmobs.writer.ddintake.LLMObsSpanMapper',
+  'datadog.trace.llmobs.writer.ddintake.LLMObsSpanMapper.PayloadV1',
   // FIXME(DSM): test coverage needed
   'datadog.trace.core.datastreams.DataStreamContextInjector',
   'datadog.trace.common.sampling.TraceSamplingRules.RuleAdapter',

dd-trace-core/src/main/java/datadog/trace/common/writer/WriterFactory.java

@@ -68,7 +68,8 @@ public static Writer createWriter(
     if (!DD_AGENT_WRITER_TYPE.equals(configuredType)
         && !DD_INTAKE_WRITER_TYPE.equals(configuredType)) {
       log.warn(
-          "Writer type not configured correctly: Type {} not recognized. Ignoring", configuredType);
+          "Writer type not configured correctly: Type {} not recognized. Ignoring ",
+          configuredType);
       configuredType = datadog.trace.api.ConfigDefaults.DEFAULT_AGENT_WRITER_TYPE;
     }
 
@@ -84,14 +85,22 @@ public static Writer createWriter(
 
     // The AgentWriter doesn't support the CI Visibility protocol. If CI Visibility is
     // enabled, check if we can use the IntakeWriter instead.
-    if (DD_AGENT_WRITER_TYPE.equals(configuredType) && config.isCiVisibilityEnabled()) {
+    if (DD_AGENT_WRITER_TYPE.equals(configuredType) && (config.isCiVisibilityEnabled())) {
       if (featuresDiscovery.supportsEvpProxy() || config.isCiVisibilityAgentlessEnabled()) {
         configuredType = DD_INTAKE_WRITER_TYPE;
       } else {
         log.info(
             "CI Visibility functionality is limited. Please upgrade to Agent v6.40+ or v7.40+ or enable Agentless mode.");
       }
     }
+    if (DD_AGENT_WRITER_TYPE.equals(configuredType) && (config.isLlmObsEnabled())) {
+      if (featuresDiscovery.supportsEvpProxy() || config.isLlmObsAgentlessEnabled()) {
+        configuredType = DD_INTAKE_WRITER_TYPE;
+      } else {
+        log.info("LLM Observability functionality is limited.");
+        // TODO: add supported agent version to this log line for llm obs
+      }
+    }
 
     RemoteWriter remoteWriter;
     if (DD_INTAKE_WRITER_TYPE.equals(configuredType)) {
@@ -117,7 +126,11 @@ public static Writer createWriter(
             createDDIntakeRemoteApi(config, commObjects, featuresDiscovery, TrackType.CITESTCOV);
         builder.addTrack(TrackType.CITESTCOV, coverageApi);
       }
-
+      if (config.isLlmObsEnabled()) {
+        final RemoteApi llmobsApi =
+            createDDIntakeRemoteApi(config, commObjects, featuresDiscovery, TrackType.LLMOBS);
+        builder.addTrack(TrackType.LLMOBS, llmobsApi);
+      }
       remoteWriter = builder.build();
 
     } else { // configuredType == DDAgentWriter
@@ -173,20 +186,34 @@ private static RemoteApi createDDIntakeRemoteApi(
       SharedCommunicationObjects commObjects,
       DDAgentFeaturesDiscovery featuresDiscovery,
       TrackType trackType) {
-    if (featuresDiscovery.supportsEvpProxy() && !config.isCiVisibilityAgentlessEnabled()) {
+    boolean evpProxySupported = featuresDiscovery.supportsEvpProxy();
+    boolean useProxyApi =
+        (evpProxySupported && TrackType.LLMOBS == trackType && !config.isLlmObsAgentlessEnabled())
+            || (evpProxySupported
+                && (TrackType.CITESTCOV == trackType || TrackType.CITESTCYCLE == trackType)
+                && !config.isCiVisibilityAgentlessEnabled());
+
+    if (useProxyApi) {
       return DDEvpProxyApi.builder()
           .httpClient(commObjects.okHttpClient)
           .agentUrl(commObjects.agentUrl)
           .evpProxyEndpoint(featuresDiscovery.getEvpProxyEndpoint())
           .trackType(trackType)
           .compressionEnabled(featuresDiscovery.supportsContentEncodingHeadersWithEvpProxy())
           .build();
-
     } else {
       HttpUrl hostUrl = null;
+      String llmObsAgentlessUrl = config.getLlMObsAgentlessUrl();
+
       if (config.getCiVisibilityAgentlessUrl() != null) {
         hostUrl = HttpUrl.get(config.getCiVisibilityAgentlessUrl());
         log.info("Using host URL '{}' to report CI Visibility traces in Agentless mode.", hostUrl);
+      } else if (config.isLlmObsEnabled()
+          && config.isLlmObsAgentlessEnabled()
+          && llmObsAgentlessUrl != null
+          && !llmObsAgentlessUrl.isEmpty()) {
+        hostUrl = HttpUrl.get(llmObsAgentlessUrl);
+        log.info("Using host URL '{}' to report LLM Obs traces in Agentless mode.", hostUrl);
       }
       return DDIntakeApi.builder()
           .hostUrl(hostUrl)

dd-trace-core/src/main/java/datadog/trace/common/writer/ddintake/DDIntakeMapperDiscovery.java

@@ -6,6 +6,7 @@
 import datadog.trace.civisibility.writer.ddintake.CiTestCycleMapperV1;
 import datadog.trace.common.writer.RemoteMapper;
 import datadog.trace.common.writer.RemoteMapperDiscovery;
+import datadog.trace.llmobs.writer.ddintake.LLMObsSpanMapper;
 
 /**
  * Mapper discovery logic when a DDIntake is used. The mapper is discovered based on a backend
@@ -40,6 +41,8 @@ public void discover() {
       mapper = new CiTestCycleMapperV1(wellKnownTags, compressionEnabled);
     } else if (TrackType.CITESTCOV.equals(trackType)) {
       mapper = new CiTestCovMapperV2(compressionEnabled);
+    } else if (TrackType.LLMOBS.equals(trackType)) {
+      mapper = new LLMObsSpanMapper();
     } else {
       mapper = RemoteMapper.NO_OP;
     }