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,89 @@ 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) {
+      logger.warn("NO ACTIVE SPAN CONTEXT");
+      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()) {
+      logger.warn("NO ACTIVE SPAN CONTEXT FOR TRACE ID {}", traceId);
+      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);
+    logger.warn("ADDED TRACE ID: {} SPAN CONTEXTS: {}", 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

@@ -55,7 +55,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);
 
       span.setTag(LLMObsTags.MODEL_NAME, modelName);
       span.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider);
@@ -66,28 +71,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) {