Fix DD_APM_TRACING_ENABLED to work with LLMObs

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

@@ -30,6 +30,11 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
       return;
     }
 
+    if (!config.isTraceEnabled()) {
+      LOGGER.debug("LLM Observability is disabled: tracing is disabled");
+      return;
+    }
+
     sco.createRemaining(config);
 
     String mlApp = config.getLlmObsMlApp();

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

@@ -5,7 +5,9 @@
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTraceApiInfo;
 import datadog.trace.api.DDTraceId;
+import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.WellKnownTags;
+import datadog.trace.api.internal.TraceSegment;
 import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsContext;
 import datadog.trace.api.llmobs.LLMObsSpan;
@@ -110,6 +112,12 @@ public DDLLMObsSpan(
     }
     span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID);
     scope = LLMObsContext.attach(span.context());
+
+    // Mark this span as originating from LLM Observability product
+    TraceSegment segment = AgentTracer.get().getTraceSegment();
+    if (segment != null) {
+      segment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.LLMOBS);
+    }
   }
 
   @Override

dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/LLMObsSystemTest.groovy

@@ -0,0 +1,53 @@
+package datadog.trace.llmobs
+
+import datadog.communication.ddagent.SharedCommunicationObjects
+import datadog.trace.test.util.DDSpecification
+import okhttp3.HttpUrl
+
+class LLMObsSystemTest extends DDSpecification {
+
+  void 'start disabled when llmobs is disabled'() {
+    setup:
+    injectSysConfig('llmobs.enabled', 'false')
+    rebuildConfig()
+    final inst = Mock(java.lang.instrument.Instrumentation)
+    final sco = Mock(SharedCommunicationObjects)
+
+    when:
+    LLMObsSystem.start(inst, sco)
+
+    then:
+    0 * sco._
+  }
+
+  void 'start disabled when trace is disabled'() {
+    setup:
+    injectSysConfig('llmobs.enabled', 'true')
+    injectSysConfig('trace.enabled', 'false')
+    rebuildConfig()
+    final inst = Mock(java.lang.instrument.Instrumentation)
+    final sco = Mock(SharedCommunicationObjects)
+
+    when:
+    LLMObsSystem.start(inst, sco)
+
+    then:
+    0 * sco._
+  }
+
+  void 'start enabled when apm tracing disabled but llmobs enabled'() {
+    setup:
+    injectSysConfig('llmobs.enabled', 'true')
+    injectSysConfig('apm.tracing.enabled', 'false')
+    rebuildConfig()
+    final inst = Mock(java.lang.instrument.Instrumentation)
+    final sco = Mock(SharedCommunicationObjects)
+    sco.agentUrl = HttpUrl.parse('http://localhost:8126')
+
+    when:
+    LLMObsSystem.start(inst, sco)
+
+    then:
+    1 * sco.createRemaining(_)
+  }
+}

dd-smoke-tests/apm-tracing-disabled/src/main/java/datadog/smoketest/apmtracingdisabled/Controller.java

@@ -1,5 +1,7 @@
 package datadog.smoketest.apmtracingdisabled;
 
+import datadog.trace.api.llmobs.LLMObs;
+import datadog.trace.api.llmobs.LLMObsSpan;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import io.opentracing.Span;
 import io.opentracing.util.GlobalTracer;
@@ -73,6 +75,17 @@ public void write(
     }
   }
 
+  @GetMapping("/llmobs/test")
+  public String llmobsTest() {
+    // Create LLMObs span using public API
+    LLMObsSpan llmSpan =
+        LLMObs.startLLMSpan("llmobs-test-operation", "gpt-4", "openai", null, null);
+    llmSpan.annotateIO("test input", "test output");
+    llmSpan.finish();
+
+    return "LLMObs test completed";
+  }
+
   private String forceKeepSpan() {
     final Span span = GlobalTracer.get().activeSpan();
     if (span != null) {

dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsApmDisabledSmokeTest.groovy

@@ -0,0 +1,92 @@
+package datadog.smoketest.apmtracingdisabled
+
+import datadog.trace.api.sampling.PrioritySampling
+import okhttp3.Request
+
+class LlmObsApmDisabledSmokeTest extends AbstractApmTracingDisabledSmokeTest {
+
+  static final String LLMOBS_SERVICE_NAME = "llmobs-apm-disabled-test"
+
+  static final String[] LLMOBS_APM_DISABLED_PROPERTIES = [
+    "-Ddd.apm.tracing.enabled=false",
+    "-Ddd.llmobs.enabled=true",
+    "-Ddd.llmobs.ml-app=test-app",
+    "-Ddd.service.name=${LLMOBS_SERVICE_NAME}",
+  ]
+
+  @Override
+  ProcessBuilder createProcessBuilder() {
+    return createProcess(LLMOBS_APM_DISABLED_PROPERTIES)
+  }
+
+  void 'When APM disabled and LLMObs enabled, LLMObs spans should be kept and APM spans should be dropped'() {
+    setup:
+    final llmobsUrl = "http://localhost:${httpPort}/rest-api/llmobs/test"
+    final llmobsRequest = new Request.Builder().url(llmobsUrl).get().build()
+
+    final apmUrl = "http://localhost:${httpPort}/rest-api/greetings"
+    final apmRequest = new Request.Builder().url(apmUrl).get().build()
+
+    when: "Create LLMObs span"
+    final llmobsResponse = client.newCall(llmobsRequest).execute()
+
+    then: "LLMObs request should succeed"
+    llmobsResponse.successful
+
+    when: "Create regular APM span"
+    final apmResponse = client.newCall(apmRequest).execute()
+
+    then: "APM request should succeed"
+    apmResponse.successful
+
+    and: "Wait for traces"
+    waitForTraceCount(2)
+
+    and: "LLMObs trace should be kept (SAMPLER_KEEP)"
+    def llmobsTrace = traces.find { trace ->
+      trace.spans.find { span ->
+        span.meta["http.url"] == llmobsUrl
+      }
+    }
+    assert llmobsTrace != null
+    // The LLMObs child span should have LLMObs tags
+    def llmobsChildSpan = llmobsTrace.spans.find { span ->
+      span.meta["_ml_obs_tag.model_name"] == "gpt-4"
+    }
+    assert llmobsChildSpan != null : "LLMObs child span with model_name=gpt-4 should exist"
+
+    and: "Regular APM trace should be dropped (SAMPLER_DROP)"
+    def apmTrace = traces.find { trace ->
+      trace.spans.find { span ->
+        span.meta["http.url"] == apmUrl
+      }
+    }
+    assert apmTrace != null
+    checkRootSpanPrioritySampling(apmTrace, PrioritySampling.SAMPLER_DROP)
+
+    and: "No NPE or errors in logs"
+    !isLogPresent { it.contains("NullPointerException") }
+    !isLogPresent { it.contains("ERROR") }
+  }
+
+  void 'LLMObs spans should have PROPAGATED_TRACE_SOURCE tag set'() {
+    setup:
+    final llmobsUrl = "http://localhost:${httpPort}/rest-api/llmobs/test"
+    final llmobsRequest = new Request.Builder().url(llmobsUrl).get().build()
+
+    when:
+    final response = client.newCall(llmobsRequest).execute()
+
+    then:
+    response.successful
+    waitForTraceCount(1)
+
+    and: "LLMObs span should be created successfully"
+    def trace = traces[0]
+    assert trace != null
+    def llmobsSpan = trace.spans.find { span ->
+      span.meta["_ml_obs_tag.model_name"] == "gpt-4"
+    }
+    assert llmobsSpan != null : "LLMObs span with model_name should exist"
+  }
+}

dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsTraceDisabledSmokeTest.groovy

@@ -0,0 +1,34 @@
+package datadog.smoketest.apmtracingdisabled
+
+import okhttp3.Request
+
+class LlmObsTraceDisabledSmokeTest extends AbstractApmTracingDisabledSmokeTest {
+
+  static final String[] LLMOBS_TRACE_DISABLED_PROPERTIES = [
+    "-Ddd.trace.enabled=false",
+    "-Ddd.llmobs.enabled=true",
+    "-Ddd.llmobs.ml-app=test-app",
+    "-Ddd.service.name=llmobs-trace-disabled-test",
+  ]
+
+  @Override
+  ProcessBuilder createProcessBuilder() {
+    return createProcess(LLMOBS_TRACE_DISABLED_PROPERTIES)
+  }
+
+  void 'DD_TRACE_ENABLED=false with DD_LLMOBS_ENABLED=true should disable LLMObs gracefully'() {
+    setup:
+    final llmobsUrl = "http://localhost:${httpPort}/rest-api/llmobs/test"
+    final llmobsRequest = new Request.Builder().url(llmobsUrl).get().build()
+
+    when: "Call LLMObs endpoint"
+    final response = client.newCall(llmobsRequest).execute()
+
+    then: "Request should succeed"
+    response.successful
+    response.code() == 200
+
+    and: "LLMObs disabled message in logs"
+    isLogPresent { it.contains("LLM Observability is disabled: tracing is disabled") }
+  }
+}

dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsAndAsmStandaloneSampler.java

@@ -0,0 +1,73 @@
+package datadog.trace.common.sampling;
+
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP;
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP;
+
+import datadog.trace.api.ProductTraceSource;
+import datadog.trace.api.sampling.SamplingMechanism;
+import datadog.trace.core.CoreSpan;
+import datadog.trace.core.DDSpan;
+import java.time.Clock;
+import java.util.concurrent.atomic.AtomicLong;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This sampler is used when APM tracing is disabled but both LLM Observability and ASM are enabled.
+ * It keeps all LLMObs and ASM traces, and allows 1 APM trace per minute for billing/service catalog
+ * purposes.
+ */
+public class LlmObsAndAsmStandaloneSampler implements Sampler, PrioritySampler {
+
+  private static final Logger log = LoggerFactory.getLogger(LlmObsAndAsmStandaloneSampler.class);
+  private static final int RATE_IN_MILLISECONDS = 60000; // 1 minute
+
+  private final AtomicLong lastSampleTime;
+  private final Clock clock;
+
+  public LlmObsAndAsmStandaloneSampler(final Clock clock) {
+    this.clock = clock;
+    this.lastSampleTime = new AtomicLong(clock.millis() - RATE_IN_MILLISECONDS);
+  }
+
+  @Override
+  public <T extends CoreSpan<T>> boolean sample(final T span) {
+    // Priority sampling sends all traces to the core agent, including traces marked dropped.
+    // This allows the core agent to collect stats on all traces.
+    return true;
+  }
+
+  @Override
+  public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
+    T rootSpan = span.getLocalRootSpan();
+    if (rootSpan instanceof DDSpan) {
+      DDSpan ddRootSpan = (DDSpan) rootSpan;
+      int traceSource = ddRootSpan.context().getPropagationTags().getTraceSource();
+      if (ProductTraceSource.isProductMarked(traceSource, ProductTraceSource.LLMOBS)) {
+        log.debug("Set SAMPLER_KEEP for LLMObs span {}", span.getSpanId());
+        span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.DEFAULT);
+        return;
+      }
+      if (ProductTraceSource.isProductMarked(traceSource, ProductTraceSource.ASM)) {
+        log.debug("Set SAMPLER_KEEP for ASM span {}", span.getSpanId());
+        span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.APPSEC);
+        return;
+      }
+    }
+    // For APM-only traces, allow 1 per minute for billing/catalog purposes
+    if (shouldSample()) {
+      log.debug("Set SAMPLER_KEEP for APM span {}", span.getSpanId());
+      span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.APPSEC);
+    } else {
+      log.debug("Set SAMPLER_DROP for APM span {}", span.getSpanId());
+      span.setSamplingPriority(SAMPLER_DROP, SamplingMechanism.APPSEC);
+    }
+  }
+
+  private boolean shouldSample() {
+    long now = clock.millis();
+    return lastSampleTime.updateAndGet(
+            lastTime -> now - lastTime >= RATE_IN_MILLISECONDS ? now : lastTime)
+        == now;
+  }
+}

dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsStandaloneSampler.java

@@ -0,0 +1,47 @@
+package datadog.trace.common.sampling;
+
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP;
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP;
+
+import datadog.trace.api.ProductTraceSource;
+import datadog.trace.api.sampling.SamplingMechanism;
+import datadog.trace.core.CoreSpan;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This sampler is used when APM tracing is disabled but LLM Observability is enabled. Unlike ASM
+ * standalone mode which only needs 1 trace per minute for billing/catalog purposes, LLM
+ * Observability needs to capture all LLM interactions to track costs, latency, and quality metrics.
+ * Therefore, this sampler keeps all LLMOBS traces and drops all APM-only traces.
+ */
+public class LlmObsStandaloneSampler implements Sampler, PrioritySampler {
+
+  private static final Logger log = LoggerFactory.getLogger(LlmObsStandaloneSampler.class);
+
+  @Override
+  public <T extends CoreSpan<T>> boolean sample(final T span) {
+    // Priority sampling sends all traces to the core agent, including traces marked dropped.
+    // This allows the core agent to collect stats on all traces.
+    return true;
+  }
+
+  @Override
+  public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
+    // Only keep traces that have the LLMOBS product flag
+    // Drop regular APM traces when APM tracing is disabled
+    T rootSpan = span.getLocalRootSpan();
+    if (rootSpan instanceof datadog.trace.core.DDSpan) {
+      datadog.trace.core.DDSpan ddRootSpan = (datadog.trace.core.DDSpan) rootSpan;
+      int traceSource = ddRootSpan.context().getPropagationTags().getTraceSource();
+      if (ProductTraceSource.isProductMarked(traceSource, ProductTraceSource.LLMOBS)) {
+        log.debug("Set SAMPLER_KEEP for LLMObs span {}", span.getSpanId());
+        span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.DEFAULT);
+        return;
+      }
+    }
+    // Drop APM-only traces when APM tracing is disabled
+    log.debug("Set SAMPLER_DROP for APM-only span {}", span.getSpanId());
+    span.setSamplingPriority(SAMPLER_DROP, SamplingMechanism.DEFAULT);
+  }
+}

dd-trace-core/src/main/java/datadog/trace/common/sampling/Sampler.java

@@ -36,9 +36,22 @@ final class Builder {
     public static Sampler forConfig(final Config config, final TraceConfig traceConfig) {
       Sampler sampler;
       if (config != null) {
-        if (!config.isApmTracingEnabled() && isAsmEnabled(config)) {
-          log.debug("APM is disabled. Only 1 trace per minute will be sent.");
-          return new AsmStandaloneSampler(Clock.systemUTC());
+        if (!config.isApmTracingEnabled()) {
+          if (config.isLlmObsEnabled() && isAsmEnabled(config)) {
+            log.debug(
+                "APM is disabled, but both LLMObs and ASM are enabled. All LLMObs and ASM traces will be kept, only 1 APM trace per minute will be sent.");
+            return new LlmObsAndAsmStandaloneSampler(Clock.systemUTC());
+          } else if (config.isLlmObsEnabled()) {
+            log.debug("APM is disabled, but LLMObs is enabled. All LLMObs traces will be kept.");
+            return new LlmObsStandaloneSampler();
+          } else if (isAsmEnabled(config)) {
+            log.debug(
+                "APM is disabled, but ASM is enabled. Only 1 APM trace per minute will be sent, all ASM traces will be kept.");
+            return new AsmStandaloneSampler(Clock.systemUTC());
+          }
+          // APM disabled and no other products enabled - drop all APM traces
+          log.debug("APM is disabled. All APM traces will be dropped.");
+          return new ForcePrioritySampler(PrioritySampling.SAMPLER_DROP, SamplingMechanism.DEFAULT);
         }
         final Map<String, String> serviceRules = config.getTraceSamplingServiceRules();
         final Map<String, String> operationRules = config.getTraceSamplingOperationRules();