From f4b6093ea283014a8dfdbfbe947ff2dd10c38b32 Mon Sep 17 00:00:00 2001
From: Jaroslav Bachorik <jaroslav.bachorik@datadoghq.com>
Date: Tue, 2 Jun 2026 14:09:27 +0200
Subject: [PATCH] feat: register llm.agent.phase slot and add phase set/clear
 API

---
 .../datadog/profiling/ddprof/DatadogProfiler.java | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfiler.java b/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfiler.java
index 9aedde9e49f..9f1ce49dfab 100644
--- a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfiler.java
+++ b/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfiler.java
@@ -17,6 +17,7 @@
 import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isCpuProfilerEnabled;
 import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isLiveHeapSizeTrackingEnabled;
 import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isMemoryLeakProfilingEnabled;
+import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isLlmPhaseAttributeEnabled;
 import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isResourceNameContextAttributeEnabled;
 import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isSpanNameContextAttributeEnabled;
 import static com.datadog.profiling.ddprof.DatadogProfilerConfig.isTrackingGenerations;
@@ -69,6 +70,7 @@ public final class DatadogProfiler {
 
   private static final String OPERATION = "_dd.trace.operation";
   private static final String RESOURCE = "_dd.trace.resource";
+  private static final String LLM_PHASE = "llm.agent.phase";
 
   private static final int MAX_NUM_ENDPOINTS = 8192;
 
@@ -105,6 +107,7 @@ public static DatadogProfiler newInstance(ConfigProvider configProvider) {
   private final Set<ProfilingMode> profilingModes = EnumSet.noneOf(ProfilingMode.class);
 
   private final ContextSetter contextSetter;
+  private final int llmPhaseOffset;
 
   private final List<String> orderedContextAttributes;
 
@@ -150,7 +153,11 @@ private DatadogProfiler(ConfigProvider configProvider) {
     if (isResourceNameContextAttributeEnabled(configProvider)) {
       orderedContextAttributes.add(RESOURCE);
     }
+    if (isLlmPhaseAttributeEnabled(configProvider)) {
+      orderedContextAttributes.add(LLM_PHASE);
+    }
     this.contextSetter = new ContextSetter(profiler, orderedContextAttributes);
+    this.llmPhaseOffset = contextSetter.offsetOf(LLM_PHASE);
     this.queueTimeThresholdMillis =
         configProvider.getLong(
             PROFILING_QUEUEING_TIME_THRESHOLD_MILLIS,
@@ -408,6 +415,14 @@ public boolean clearContextValue(int offset) {
     return false;
   }
 
+  public boolean setAgentPhase(String phaseToken) {
+    return contextSetter.setContextValue(llmPhaseOffset, phaseToken);
+  }
+
+  public boolean clearAgentPhase() {
+    return contextSetter.clearContextValue(llmPhaseOffset);
+  }
+
   private void debugLogging(long localRootSpanId) {
     if (detailedDebugLogging && log.isDebugEnabled()) {
       log.debug("localRootSpanId={}", localRootSpanId, new Throwable());