Add BTrace extensions and MCP server: runtime contracts, AI/ML observability, LLM-driven diagnostics

.github/workflows/continuous.yml

@@ -26,7 +26,7 @@ jobs:
       - name: Set up Java
         uses: actions/setup-java@v5
         with:
-          java-version: 11
+          java-version: 21
           distribution: temurin
       - name: Checkout
         uses: actions/checkout@v6

btrace-client/src/main/java/io/btrace/client/Client.java

@@ -801,7 +801,7 @@ public void attach(String pid, String agentPath, String sysCp, String bootCp) th
     }
   }
 
-  void connectAndListProbes(String host, CommandListener listener) throws IOException {
+  public void connectAndListProbes(String host, CommandListener listener) throws IOException {
     if (sock != null) {
       throw new IllegalStateException();
     }
@@ -1111,11 +1111,11 @@ public synchronized void close() throws IOException {
     reset();
   }
 
-  boolean isDisconnected() {
+  public boolean isDisconnected() {
     return disconnected;
   }
 
-  void disconnect() throws IOException {
+  public void disconnect() throws IOException {
     disconnected = true;
     if (log.isDebugEnabled()) {
       log.debug("sending DISCONNECT request to agent");

btrace-client/src/main/java/io/btrace/extcli/Installer.java

@@ -108,7 +108,9 @@ static void install(String target, List<String> repos, String id, boolean dryRun
         throw new IOException("Failed to download extension from provided repositories.");
     } else {
       throw new IllegalArgumentException(
-          "Unrecognized input: provide a zip path, URL, or group:artifact:version");
+          "Unrecognised target: '"
+              + target
+              + "'. Expected a URL, .zip path, or groupId:artifactId:version coordinate.");
     }
 
     // Validate zip contains -api.jar and -impl.jar, and install

btrace-client/src/test/java/io/btrace/extcli/ExtensionListerTest.java

@@ -36,13 +36,11 @@ class ExtensionListerTest {
   private final ByteArrayOutputStream errContent = new ByteArrayOutputStream();
   private final PrintStream originalOut = System.out;
   private final PrintStream originalErr = System.err;
-  private String originalBtraceHome;
 
   @BeforeEach
   void setUpStreams() {
     System.setOut(new PrintStream(outContent));
     System.setErr(new PrintStream(errContent));
-    originalBtraceHome = System.getenv("BTRACE_HOME");
   }
 
   @AfterEach
@@ -68,17 +66,6 @@ void listFromBtraceHome() throws IOException {
     assertNotNull(output);
   }
 
-  @Test
-  void listWithJsonFormat() throws IOException {
-    ExtensionLister.list(true);
-
-    String output = outContent.toString();
-    // Should output valid JSON (starts with [ and ends with ])
-    assertTrue(
-        output.trim().startsWith("[") && output.trim().endsWith("]"),
-        "JSON output should be an array");
-  }
-
   @Test
   void listHandlesEmptyDirectories() throws IOException {
     // Create empty extensions directory

btrace-compiler/src/main/java/io/btrace/compiler/VerifierVisitor.java

@@ -537,14 +537,22 @@ public Void visitVariable(VariableTree vt, Void p) {
   /**
    * Returns true if the given service class name is declared by any extension.
    *
-   * <p>Compile-time validation notes: - This check operates without loading classes. It inspects
-   * extension metadata (BTrace-Extension-Services in MANIFEST.MF and legacy
-   * META-INF/btrace-extension.properties) to verify that an @Injected service type is declared by
-   * some extension. - It complements the bytecode-time check in instr (BTraceProbeNode) and the
-   * runtime reflection-based validation in the agent (Client#validateDeclaredServices). The latter
-   * ensures correctness under the actual runtime classloader/JPMS environment.
+   * <p>Compile-time validation notes: checks first via the annotation processing type model (which
+   * sees -cp JARs), then falls back to scanning jar manifests on the JVM classpath. The agent
+   * performs a definitive runtime check (Client#validateDeclaredServices).
    */
   private boolean isDeclaredExtensionService(String serviceClassName) {
+    // Primary: use the annotation processing Elements API — sees the compilation classpath
+    TypeElement te = verifier.getElementUtils().getTypeElement(serviceClassName);
+    if (te != null) {
+      for (javax.lang.model.element.AnnotationMirror am : te.getAnnotationMirrors()) {
+        String annName = am.getAnnotationType().asElement().toString();
+        if ("io.btrace.core.extensions.ServiceDescriptor".equals(annName)) {
+          return true;
+        }
+      }
+    }
+    // Fallback: scan JAR manifests visible to the JVM classloader
     String resourceName = serviceClassName.replace('.', '/') + ".class";
     ClassLoader cl = Thread.currentThread().getContextClassLoader();
     if (cl == null) {
@@ -566,7 +574,6 @@ private boolean isDeclaredExtensionService(String serviceClassName) {
           }
         }
       }
-      // Fallback: scan system classpath jars
       String cp = System.getProperty("java.class.path", "");
       String[] parts = cp.split(java.io.File.pathSeparator);
       for (String p : parts) {

btrace-dist/build.gradle

@@ -638,7 +638,7 @@ task buildDockerContext(type: Copy, dependsOn: [btraceJar, copyDtraceLib, proces
 
     from "${distTarget}"
     into layout.buildDirectory.dir("docker-context/btrace")
-    fileMode = 0644
+    filePermissions { unix(0644) }
 
     doLast {
         // Copy entrypoint script
@@ -807,7 +807,7 @@ task buildDockerImages {
         description "Generates Javadoc API documentation for the btrace-${name}."
 
         title = "btrace-${name}"
-        destinationDir = file("${project.docsDir}/${name}/javadoc")
+        destinationDir = layout.buildDirectory.dir("docs/${name}/javadoc").get().asFile
         classpath = files(compileJava.destinationDirectory) + configurations.artifact.asFileTree
         options.addStringOption('Xdoclint:all,-missing', '-quiet')
         failOnError false

btrace-dist/src/main/resources/samples/ContractCheck.java

@@ -0,0 +1,103 @@
+import io.btrace.contracts.ContractService;
+import io.btrace.core.annotations.BTrace;
+import io.btrace.core.annotations.Duration;
+import io.btrace.core.annotations.Injected;
+import io.btrace.core.annotations.Kind;
+import io.btrace.core.annotations.Location;
+import io.btrace.core.annotations.OnEvent;
+import io.btrace.core.annotations.OnMethod;
+import io.btrace.core.annotations.OnTimer;
+import io.btrace.core.annotations.ProbeClassName;
+import io.btrace.core.annotations.ProbeMethodName;
+import io.btrace.core.annotations.Return;
+
+import static io.btrace.core.BTraceUtils.*;
+
+/**
+ * Runtime behavioral contracts: latency budgets, call-rate limits, null-safety, and tagged
+ * code path profiling — without modifying the target code.
+ *
+ * <p>Requires the btrace-contracts extension in $BTRACE_HOME/extensions/.
+ *
+ * <p>Attach to a running JVM:
+ * <pre>
+ * btrace &lt;pid&gt; ContractCheck.java
+ * </pre>
+ *
+ * <p>Trigger an on-demand summary:
+ * <pre>
+ * btrace &lt;pid&gt; --event summary
+ * </pre>
+ */
+@BTrace
+public class ContractCheck {
+
+  @Injected
+  private static ContractService contracts;
+
+  // ==================== Latency budgets ====================
+
+  /**
+   * Enforce a 500ms latency budget on all methods in the target package.
+   * Adjust the clazz pattern to match your project structure.
+   */
+  @OnMethod(
+      clazz = "/com\\.myapp\\..*/",
+      method = "/.*/",
+      location = @Location(Kind.RETURN))
+  public static void checkLatency(
+      @ProbeClassName String cls,
+      @ProbeMethodName String method,
+      @Duration long dur) {
+    contracts.checkLatency(strcat(cls, strcat(".", method)), dur, 500_000_000L);
+  }
+
+  // ==================== Null safety ====================
+
+  @OnMethod(
+      clazz = "/com\\.myapp\\..*/",
+      method = "/.*/",
+      location = @Location(Kind.RETURN))
+  public static void checkNullReturn(
+      @ProbeClassName String cls,
+      @ProbeMethodName String method,
+      @Return Object ret) {
+    contracts.checkNotNull(strcat(cls, strcat(".", method)), ret);
+  }
+
+  // ==================== Tagged path comparison ====================
+
+  // Example: compare two implementations of the same operation.
+  // Tag them differently so getSummary() renders them side by side.
+
+  @OnMethod(
+      clazz = "com.myapp.CachedQueryService",
+      method = "query",
+      location = @Location(Kind.RETURN))
+  public static void onCachedQuery(@Duration long dur) {
+    contracts.trackCodePath("QueryService.query", dur, "cached");
+  }
+
+  @OnMethod(
+      clazz = "com.myapp.DirectQueryService",
+      method = "query",
+      location = @Location(Kind.RETURN))
+  public static void onDirectQuery(@Duration long dur) {
+    contracts.trackCodePath("QueryService.query", dur, "direct");
+  }
+
+  // ==================== Reporting ====================
+
+  @OnTimer(10000)
+  public static void periodicCheck() {
+    if (contracts.hasViolations()) {
+      println("=== CONTRACT VIOLATIONS DETECTED ===");
+      println(contracts.getSummary());
+    }
+  }
+
+  @OnEvent("summary")
+  public static void onDemandSummary() {
+    println(contracts.getSummary());
+  }
+}

btrace-dist/src/main/resources/samples/GpuBridge.java

@@ -0,0 +1,85 @@
+import io.btrace.core.annotations.BTrace;
+import io.btrace.core.annotations.Duration;
+import io.btrace.core.annotations.Injected;
+import io.btrace.core.annotations.Kind;
+import io.btrace.core.annotations.Location;
+import io.btrace.core.annotations.OnEvent;
+import io.btrace.core.annotations.OnMethod;
+import io.btrace.core.annotations.OnTimer;
+import io.btrace.gpu.GpuBridgeService;
+
+import static io.btrace.core.BTraceUtils.*;
+
+/**
+ * Traces GPU model inference via ONNX Runtime and DJL (Deep Java Library).
+ * Tracks inference latency, batch sizes, and model load times.
+ *
+ * <p>Attach to a JVM running ONNX or DJL inference:
+ * <pre>
+ * btrace &lt;pid&gt; GpuBridge.java
+ * </pre>
+ */
+@BTrace
+public class GpuBridge {
+
+  @Injected
+  private static GpuBridgeService gpu;
+
+  // ==================== ONNX Runtime ====================
+
+  @OnMethod(
+      clazz = "ai.onnxruntime.OrtSession",
+      method = "run",
+      location = @Location(Kind.RETURN))
+  public static void onOnnxInference(@Duration long dur) {
+    gpu.recordInference("onnx", "session", dur);
+  }
+
+  @OnMethod(
+      clazz = "ai.onnxruntime.OrtSession",
+      method = "<init>",
+      location = @Location(Kind.RETURN))
+  public static void onOnnxModelLoad(@Duration long dur) {
+    gpu.recordModelLoad("onnx", "session", dur);
+  }
+
+  // ==================== DJL (Deep Java Library) ====================
+
+  @OnMethod(
+      clazz = "/ai\\.djl\\.inference\\.Predictor/",
+      method = "predict",
+      location = @Location(Kind.RETURN))
+  public static void onDjlPredict(@Duration long dur) {
+    gpu.recordInference("djl", "predictor", dur);
+  }
+
+  @OnMethod(
+      clazz = "/ai\\.djl\\.repository\\.zoo\\.ModelZoo/",
+      method = "loadModel",
+      location = @Location(Kind.RETURN))
+  public static void onDjlModelLoad(@Duration long dur) {
+    gpu.recordModelLoad("djl", "model-zoo", dur);
+  }
+
+  // ==================== TensorFlow Java ====================
+
+  @OnMethod(
+      clazz = "/org\\.tensorflow\\.Session/",
+      method = "run",
+      location = @Location(Kind.RETURN))
+  public static void onTensorFlowRun(@Duration long dur) {
+    gpu.recordInference("tensorflow", "session", dur);
+  }
+
+  // ==================== Periodic summary ====================
+
+  @OnTimer(30000)
+  public static void periodicSummary() {
+    println(gpu.getSummary());
+  }
+
+  @OnEvent("summary")
+  public static void onDemandSummary() {
+    println(gpu.getSummary());
+  }
+}