use trace ID to map to individual stacks of active contexts

Author: gary-huang

communication/build.gradle

@@ -4,7 +4,6 @@ dependencies {
   implementation libs.slf4j
 
   api project(':remote-config:remote-config-api')
-  implementation project(':components:json')
   implementation project(':remote-config:remote-config-core')
   implementation project(':internal-api')
   implementation project(':utils:container-utils')

communication/src/main/java/datadog/communication/serialization/json/JSONWritableFormatter.java

@@ -4,6 +4,17 @@
 import datadog.communication.BackendApiFactory;
 import datadog.communication.ddagent.SharedCommunicationObjects;
 import datadog.trace.api.Config;
+import datadog.trace.api.DDTraceId;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.llmobs.domain.SpanContextInfo;
+import java.util.Deque;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import javax.annotation.Nonnull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -14,9 +25,86 @@ public class LLMObsServices {
   final Config config;
   final BackendApi backendApi;
 
+  Map<DDTraceId, Deque<SpanContextInfo>> activeSpanContextByTID = new ConcurrentHashMap<>();
+
   LLMObsServices(Config config, SharedCommunicationObjects sco) {
     this.config = config;
     this.backendApi =
         new BackendApiFactory(config, sco).createBackendApi(BackendApiFactory.Intake.LLMOBS_API);
   }
+
+  @Nonnull
+  public SpanContextInfo getActiveSpanContext() {
+    // Valid case: possibly start root llm obs span/trace while there is NOT an active apm trace
+    AgentScope activeScope = AgentTracer.activeScope();
+    if (activeScope == null) {
+      return new SpanContextInfo();
+    }
+
+    // Unexpected case: null active scope span, log to avoid crashes
+    if (activeScope.span() == null) {
+      logger.warn("active span scope found but no null span");
+      return new SpanContextInfo();
+    }
+
+    // Unexpected case: null trace ID, log to avoid crashes
+    DDTraceId traceId = activeScope.span().getTraceId();
+    if (traceId == null) {
+      logger.warn("active scope found but unexpectedly null trace ID");
+      return new SpanContextInfo();
+    }
+
+    Deque<SpanContextInfo> activeSpanCtxForTID = activeSpanContextByTID.get(traceId);
+    // Valid case: possibly start root llm obs span/trace while there's an active apm trace
+    if (activeSpanCtxForTID == null || activeSpanCtxForTID.isEmpty()) {
+      return new SpanContextInfo();
+    }
+
+    // Valid case: possibly start child llm obs span for a given trace ID
+    return activeSpanCtxForTID.peek();
+  }
+
+  public void setActiveSpanContext(SpanContextInfo spanContext) {
+    AgentSpanContext activeCtx = spanContext.getActiveContext();
+    if (activeCtx == null) {
+      logger.warn("unexpected null active context");
+      return;
+    }
+
+    DDTraceId traceId = activeCtx.getTraceId();
+    if (traceId == null) {
+      logger.warn("unexpected null trace ID");
+      return;
+    }
+
+    Deque<SpanContextInfo> contexts = activeSpanContextByTID.get(activeCtx.getTraceId());
+    if (contexts == null) {
+      contexts = new ConcurrentLinkedDeque<>();
+    }
+    contexts.push(spanContext);
+    this.activeSpanContextByTID.put(traceId, contexts);
+  }
+
+  public void removeActiveSpanContext(DDTraceId traceId) {
+    if (!activeSpanContextByTID.containsKey(traceId)) {
+      logger.debug("active span contexts not found for trace {}", traceId);
+      return;
+    }
+    Deque<SpanContextInfo> contexts = activeSpanContextByTID.get(traceId);
+    if (contexts == null) {
+      return;
+    }
+    if (!contexts.isEmpty()) {
+      try {
+        contexts.pop();
+        if (contexts.isEmpty()) {
+          // the trace MAY still be active, however, the next set should re-create the hierarchy as
+          // needed
+          activeSpanContextByTID.remove(traceId);
+        }
+      } catch (NoSuchElementException noSuchElementException) {
+        logger.debug("failed to pop context stack for trace {}", traceId);
+      }
+    }
+  }
 }

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsServices.java

@@ -57,7 +57,12 @@ public LLMObsSpan startLLMSpan(
 
       DDLLMObsSpan span =
           new DDLLMObsSpan(
-              Tags.LLMOBS_LLM_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
+              Tags.LLMOBS_LLM_SPAN_KIND,
+              spanName,
+              getMLApp(mlApp),
+              sessionID,
+              serviceName,
+              llmObsServices);
 
       if (modelName == null || modelName.isEmpty()) {
         modelName = CUSTOM_MODEL_VAL;
@@ -75,28 +80,48 @@ public LLMObsSpan startLLMSpan(
     public LLMObsSpan startAgentSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionID) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_AGENT_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
+          Tags.LLMOBS_AGENT_SPAN_KIND,
+          spanName,
+          getMLApp(mlApp),
+          sessionID,
+          serviceName,
+          llmObsServices);
     }
 
     @Override
     public LLMObsSpan startToolSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionID) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_TOOL_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
+          Tags.LLMOBS_TOOL_SPAN_KIND,
+          spanName,
+          getMLApp(mlApp),
+          sessionID,
+          serviceName,
+          llmObsServices);
     }
 
     @Override
     public LLMObsSpan startTaskSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionID) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_TASK_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
+          Tags.LLMOBS_TASK_SPAN_KIND,
+          spanName,
+          getMLApp(mlApp),
+          sessionID,
+          serviceName,
+          llmObsServices);
     }
 
     @Override
     public LLMObsSpan startWorkflowSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionID) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
+          Tags.LLMOBS_WORKFLOW_SPAN_KIND,
+          spanName,
+          getMLApp(mlApp),
+          sessionID,
+          serviceName,
+          llmObsServices);
     }
 
     private String getMLApp(String mlApp) {

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java

@@ -1,12 +1,15 @@
 package datadog.trace.llmobs.domain;
 
 import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTraceId;
 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 datadog.trace.llmobs.LLMObsServices;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -32,36 +35,51 @@ public class DDLLMObsSpan implements LLMObsSpan {
 
   private static final String LLM_OBS_INSTRUMENTATION_NAME = "llmobs";
 
+  private static final String PARENT_ID_TAG_INTERNAL = "parent_id";
+
   private static final Logger LOGGER = LoggerFactory.getLogger(DDLLMObsSpan.class);
 
   private final AgentSpan span;
   private final String spanKind;
 
   private boolean finished = false;
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(DDLLMObsSpan.class);
+  private final LLMObsServices llmObsServices;
 
   public DDLLMObsSpan(
       @Nonnull String kind,
       String spanName,
       @Nonnull String mlApp,
       String sessionID,
-      @Nonnull String serviceName) {
+      @Nonnull String serviceName,
+      @Nonnull LLMObsServices llmObsServices) {
 
     if (null == spanName || spanName.isEmpty()) {
       spanName = kind;
     }
 
+    this.llmObsServices = llmObsServices;
+
     AgentTracer.SpanBuilder spanBuilder =
         AgentTracer.get()
             .buildSpan(LLM_OBS_INSTRUMENTATION_NAME, spanName)
             .withServiceName(serviceName)
             .withSpanType(DDSpanTypes.LLMOBS);
 
+    SpanContextInfo activeSpanCtxInfo = this.llmObsServices.getActiveSpanContext();
+    AgentSpanContext activeCtx = activeSpanCtxInfo.getActiveContext();
+    if (!activeSpanCtxInfo.isRoot() && null != activeCtx) {
+      spanBuilder.asChildOf(activeCtx);
+    }
     this.span = spanBuilder.start();
+    this.llmObsServices.setActiveSpanContext(
+        new SpanContextInfo(span.context(), String.valueOf(span.context().getSpanId())));
+
     this.span.setTag(SPAN_KIND, kind);
     this.spanKind = kind;
     this.span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp);
+    this.span.setTag(
+        LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, activeSpanCtxInfo.getParentSpanID());
     if (sessionID != null && !sessionID.isEmpty()) {
       this.span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionID);
     }
@@ -87,7 +105,6 @@ public void annotateIO(List<LLMObs.LLMMessage> inputData, List<LLMObs.LLMMessage
     if (finished) {
       return;
     }
-    LOGGER.warn("ANNOTATE IN {} OUT {}", inputData, outputData);
     if (inputData != null && !inputData.isEmpty()) {
       this.span.setTag(INPUT, inputData);
     }
@@ -273,7 +290,9 @@ public void finish() {
     if (finished) {
       return;
     }
+    DDTraceId traceId = this.span.getTraceId();
     this.span.finish();
     this.finished = true;
+    this.llmObsServices.removeActiveSpanContext(traceId);
   }
 }