Gary/use ctx api

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,20 @@ 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 +289,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,38 @@
+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-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java

@@ -75,6 +75,9 @@ public class LLMObsSpanMapper implements RemoteMapper {
   private static final byte[] LLM_TOOL_CALL_ARGUMENTS =
       "arguments".getBytes(StandardCharsets.UTF_8);
 
+  // TODO is there a better place for this?
+  private static final String PARENT_ID_TAG_INTERNAL_FULL = LLMOBS_TAG_PREFIX + "parent_id";
+
   private final LLMObsSpanMapper.MetaWriter metaWriter = new MetaWriter();
   private final int size;
 
@@ -113,8 +116,8 @@ public void map(List<? extends CoreSpan<?>> trace, Writable writable) {
 
       // 3
       writable.writeUTF8(PARENT_ID);
-      // TODO fix after parent ID tracking is in place
-      writable.writeString("undefined", null);
+      writable.writeString(span.getTag(PARENT_ID_TAG_INTERNAL_FULL), null);
+      span.removeTag(PARENT_ID_TAG_INTERNAL_FULL);
 
       // 4
       writable.writeUTF8(NAME);