test(inst/langchain4j): add mock ChatModel CI test and Ollama manual demo

dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java

@@ -82,6 +82,14 @@ public void clearContext() {
     DDPROF.clearContextValue(RESOURCE_NAME_INDEX);
   }
 
+  public static void setLlmPhase(String phaseToken) {
+    DDPROF.setAgentPhase(phaseToken);
+  }
+
+  public static void clearLlmPhase() {
+    DDPROF.clearAgentPhase();
+  }
+
   @Override
   public ProfilingContextAttribute createContextAttribute(String attribute) {
     return new DatadogProfilerContextSetter(attribute, DDPROF);

dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/build.gradle

@@ -13,7 +13,32 @@ muzzle {
 addTestSuiteForDir('latestDepTest', 'test')
 
 dependencies {
+  compileOnly project(':dd-java-agent:agent-profiling:profiling-ddprof')
   compileOnly group: 'dev.langchain4j', name: 'langchain4j-core', version: minVer
   testImplementation group: 'dev.langchain4j', name: 'langchain4j', version: minVer
+  testImplementation group: 'dev.langchain4j', name: 'langchain4j-ollama', version: '1.2.0'
   latestDepTestImplementation group: 'dev.langchain4j', name: 'langchain4j', version: '+'
+  latestDepTestCompileOnly group: 'dev.langchain4j', name: 'langchain4j-ollama', version: '+'
+}
+
+tasks.register('runOllamaDemo', JavaExec) {
+  dependsOn testClasses, ':dd-java-agent:shadowJar'
+  classpath = sourceSets.test.runtimeClasspath
+  mainClass = 'datadog.trace.instrumentation.langchain4j.demo.OllamaLlmPipelineDemo'
+  // Pick the highest semver agent jar in build/libs, ignoring classifier jars.
+  def agentJar = fileTree("$rootDir/dd-java-agent/build/libs")
+      .include("dd-java-agent-*.jar")
+      .exclude("*-sources.jar", "*-javadoc.jar")
+      .files
+      .sort { f ->
+          def m = f.name =~ /dd-java-agent-(\d+)\.(\d+)\.(\d+)/
+          m ? (m[0][1] as int) * 1_000_000 + (m[0][2] as int) * 1_000 + (m[0][3] as int) : 0
+      }.last()
+  jvmArgs "-javaagent:${agentJar}",
+      "-Ddd.profiling.enabled=true",
+      "-Ddd.profiling.context.attributes.llm.phase.enabled=true",
+      "-Ddd.trace.enabled=false",
+      "-Ddd.profiling.start-force-first=true",
+      "-Ddd.profiling.upload.period=10",
+      "-Ddd.profiling.debug.dump_path=/tmp/llm-demo-profiles"
 }

dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/AiServicesInstrumentation.java

@@ -3,9 +3,8 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
+import com.datadog.profiling.ddprof.DatadogProfilingIntegration;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.profiling.Profiling;
-import datadog.trace.api.profiling.ProfilingScope;
 import net.bytebuddy.asm.Advice;
 
 public class AiServicesInstrumentation
@@ -25,16 +24,13 @@ public void methodAdvice(MethodTransformer transformer) {
 
   public static final class InvokeAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static ProfilingScope enter() {
-      ProfilingScope scope = Profiling.get().newScope();
-      scope.setContextValue("llm.agent.phase", "context_build");
-      return scope;
+    public static void enter() {
+      DatadogProfilingIntegration.setLlmPhase("context_build");
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void exit(@Advice.Enter final ProfilingScope scope) {
-      scope.clearContextValue("llm.agent.phase");
-      scope.close();
+    public static void exit() {
+      DatadogProfilingIntegration.clearLlmPhase();
     }
   }
 }

dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ChatModelInstrumentation.java

@@ -5,9 +5,8 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
+import com.datadog.profiling.ddprof.DatadogProfilingIntegration;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.profiling.Profiling;
-import datadog.trace.api.profiling.ProfilingScope;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -36,16 +35,13 @@ public void methodAdvice(MethodTransformer transformer) {
 
   public static final class ChatAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static ProfilingScope enter() {
-      ProfilingScope scope = Profiling.get().newScope();
-      scope.setContextValue("llm.agent.phase", "awaiting_inference");
-      return scope;
+    public static void enter() {
+      DatadogProfilingIntegration.setLlmPhase("awaiting_inference");
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void exit(@Advice.Enter final ProfilingScope scope) {
-      scope.clearContextValue("llm.agent.phase");
-      scope.close();
+    public static void exit() {
+      DatadogProfilingIntegration.clearLlmPhase();
     }
   }
 }

dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ToolExecutorInstrumentation.java

@@ -5,9 +5,8 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
+import com.datadog.profiling.ddprof.DatadogProfilingIntegration;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.profiling.Profiling;
-import datadog.trace.api.profiling.ProfilingScope;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -36,16 +35,13 @@ public void methodAdvice(MethodTransformer transformer) {
 
   public static final class ExecuteAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static ProfilingScope enter() {
-      ProfilingScope scope = Profiling.get().newScope();
-      scope.setContextValue("llm.agent.phase", "tool_execution");
-      return scope;
+    public static void enter() {
+      DatadogProfilingIntegration.setLlmPhase("tool_execution");
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void exit(@Advice.Enter final ProfilingScope scope) {
-      scope.clearContextValue("llm.agent.phase");
-      scope.close();
+    public static void exit() {
+      DatadogProfilingIntegration.clearLlmPhase();
     }
   }
 }

dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/MockLlmPipelineTest.java

@@ -0,0 +1,62 @@
+package datadog.trace.instrumentation.langchain4j.demo;
+
+import dev.langchain4j.agent.tool.Tool;
+import dev.langchain4j.agent.tool.ToolExecutionRequest;
+import dev.langchain4j.data.message.AiMessage;
+import dev.langchain4j.memory.chat.MessageWindowChatMemory;
+import dev.langchain4j.model.chat.ChatModel;
+import dev.langchain4j.model.chat.request.ChatRequest;
+import dev.langchain4j.model.chat.response.ChatResponse;
+import dev.langchain4j.service.AiServices;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+public class MockLlmPipelineTest {
+
+    interface WeatherAssistant {
+        String chat(String message);
+    }
+
+    static class MockWeatherTool {
+        @Tool("Get current weather for a location")
+        public String getWeather(String location) {
+            return "Sunny, 22°C in " + location;
+        }
+    }
+
+    static class TwoTurnMockModel implements ChatModel {
+        private final AtomicInteger calls = new AtomicInteger();
+
+        @Override
+        public ChatResponse chat(ChatRequest request) {
+            try { Thread.sleep(30); } catch (InterruptedException ignored) {}
+            if (calls.getAndIncrement() == 0) {
+                return ChatResponse.builder()
+                    .aiMessage(AiMessage.from(
+                        ToolExecutionRequest.builder()
+                            .name("getWeather")
+                            .arguments("{\"location\": \"Amsterdam\"}")
+                            .build()))
+                    .build();
+            }
+            return ChatResponse.builder()
+                .aiMessage(AiMessage.from("The weather in Amsterdam is Sunny, 22°C."))
+                .build();
+        }
+    }
+
+    @Test
+    public void pipelineExercisesAllThreeInstrumentationPoints() {
+        WeatherAssistant assistant = AiServices.builder(WeatherAssistant.class)
+            .chatModel(new TwoTurnMockModel())
+            .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
+            .tools(new MockWeatherTool())
+            .build();
+
+        String response = assistant.chat("What is the weather in Amsterdam?");
+        assertNotNull(response, "Expected a non-null response from the mock pipeline");
+    }
+}

dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/OllamaLlmPipelineDemo.java

@@ -0,0 +1,58 @@
+package datadog.trace.instrumentation.langchain4j.demo;
+
+import dev.langchain4j.memory.chat.MessageWindowChatMemory;
+import dev.langchain4j.model.ollama.OllamaChatModel;
+import dev.langchain4j.service.AiServices;
+
+import java.time.Duration;
+
+/**
+ * Generates a JFR recording with llm.agent.phase tags against a local Ollama server.
+ *
+ * Run:
+ *   java -javaagent:dd-java-agent.jar
+ *     -Ddd.profiling.enabled=true
+ *     -Ddd.profiling.context.attributes.llm.phase.enabled=true
+ *     -cp <classpath>
+ *     datadog.trace.instrumentation.langchain4j.demo.OllamaLlmPipelineDemo
+ *
+ * Prerequisites: ollama serve && ollama pull llama3
+ */
+public class OllamaLlmPipelineDemo {
+
+    interface Assistant {
+        String chat(String message);
+    }
+
+    public static void main(String[] args) {
+        OllamaChatModel model = OllamaChatModel.builder()
+            .baseUrl("http://localhost:11434")
+            .modelName("llama3")
+            .timeout(Duration.ofMinutes(2))
+            .build();
+
+        Assistant assistant = AiServices.builder(Assistant.class)
+            .chatModel(model)
+            .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
+            .build();
+
+        String[] questions = {
+            "Explain Java garbage collection in one sentence.",
+            "What is a Java virtual thread?",
+            "Name one advantage of the G1 garbage collector."
+        };
+        for (String q : questions) {
+            System.out.println("Q: " + q);
+            System.out.println("A: " + assistant.chat(q));
+        }
+
+        // Hold the JVM alive so the profiler flushes at least two recording cycles
+        // (dd.profiling.upload.period=10 → flush at ~10 s and ~20 s).
+        System.out.println("Waiting 25 s for profiling data to flush...");
+        try {
+            Thread.sleep(25_000);
+        } catch (InterruptedException ignored) {
+        }
+        System.out.println("Done.");
+    }
+}