Attach via {@code -javaagent:braintrust-java-agent.jar} to automatically instrument AI client + * libraries (OpenAI, Anthropic, Google GenAI, LangChain4j) for Braintrust tracing. + * + *
This class is intentionally dependency-free. It lives on the system classloader (since that's + * where {@code -javaagent} JARs are loaded) and its only job is to create a {@link + * BraintrustClassLoader} that loads the real agent implementation from hidden {@code .classdata} + * entries inside the agent JAR. + */ +public class BraintrustAgent { + + private static final String INSTALLER_CLASS = "dev.braintrust.agent.internal.AgentInstaller"; + private static final String INSTALLER_METHOD = "install"; + + private static volatile boolean installed = false; + + /** + * Entry point when the agent is loaded at JVM startup via {@code -javaagent}. + */ + public static void premain(String agentArgs, Instrumentation inst) { + install(agentArgs, inst); + } + + /** + * Entry point when the agent is attached to a running JVM via the Attach API. + */ + public static void agentmain(String agentArgs, Instrumentation inst) { + install(agentArgs, inst); + } + + private static synchronized void install(String agentArgs, Instrumentation inst) { + if (installed) { + log("Agent already installed, skipping."); + return; + } + installed = true; + + log("Braintrust Java Agent starting..."); + + try { + // Locate the agent JAR from our own code source + URL agentJarURL = + BraintrustAgent.class.getProtectionDomain().getCodeSource().getLocation(); + log("Agent JAR: " + agentJarURL); + + // Create the isolated classloader that reads .classdata entries from inst/ + ClassLoader parent = BraintrustAgent.class.getClassLoader(); + BraintrustClassLoader agentClassLoader = new BraintrustClassLoader(agentJarURL, parent); + + // Load and invoke the real agent installer through the isolated classloader + Class installerClass = agentClassLoader.loadClass(INSTALLER_CLASS); + installerClass + .getMethod(INSTALLER_METHOD, String.class, Instrumentation.class) + .invoke(null, agentArgs, inst); + + log("Braintrust Java Agent installed."); + } catch (Throwable t) { + log("ERROR: Failed to install Braintrust Java Agent: " + t.getMessage()); + t.printStackTrace(System.err); + } + } + + private static void log(String msg) { + System.out.println("[braintrust] " + msg); + } +} diff --git a/braintrust-java-agent/agent-bootstrap/src/main/java/dev/braintrust/agent/BraintrustClassLoader.java b/braintrust-java-agent/agent-bootstrap/src/main/java/dev/braintrust/agent/BraintrustClassLoader.java new file mode 100644 index 0000000..8539716 --- /dev/null +++ b/braintrust-java-agent/agent-bootstrap/src/main/java/dev/braintrust/agent/BraintrustClassLoader.java @@ -0,0 +1,109 @@ +package dev.braintrust.agent; + +import java.io.IOException; +import java.io.InputStream; +import java.net.URL; +import java.security.CodeSource; +import java.security.SecureClassLoader; +import java.security.cert.Certificate; +import java.util.jar.JarEntry; +import java.util.jar.JarFile; + +/** + * A classloader that loads agent-internal classes from {@code .classdata} entries inside the agent + * JAR. + * + *
Classes stored under the {@code inst/} prefix with a {@code .classdata} extension are + * invisible to the JVM's default classloading mechanism. This classloader knows how to find them, + * providing full classloader isolation between the agent's internals and the application's + * classpath. + * + *
The delegation model is: parent-first (standard), falling back to reading {@code .classdata} + * entries from the agent JAR. This means bootstrap classes (like this class itself and {@link + * BraintrustAgent}) are loaded by the parent (system classloader), while agent internals (ByteBuddy, + * OTel SDK, instrumentation code) are loaded here and invisible to the application. + */ +public class BraintrustClassLoader extends SecureClassLoader { + + private static final String ENTRY_PREFIX = "inst/"; + private static final String CLASS_DATA_SUFFIX = ".classdata"; + + private final JarFile agentJarFile; + private final CodeSource agentCodeSource; + private final String agentResourcePrefix; + + static { + registerAsParallelCapable(); + } + + /** + * Creates a new BraintrustClassLoader. + * + * @param agentJarURL the URL of the agent JAR file (from the -javaagent path) + * @param parent the parent classloader (typically the system/platform classloader) + */ + public BraintrustClassLoader(URL agentJarURL, ClassLoader parent) throws Exception { + super(parent); + this.agentJarFile = new JarFile(new java.io.File(agentJarURL.toURI()), false); + this.agentCodeSource = new CodeSource(agentJarURL, (Certificate[]) null); + this.agentResourcePrefix = "jar:file:" + agentJarFile.getName() + "!/"; + } + + @Override + protected Class findClass(String name) throws ClassNotFoundException { + // Convert "dev.braintrust.agent.internal.AgentInstaller" + // -> "inst/dev/braintrust/agent/internal/AgentInstaller.classdata" + String entryName = ENTRY_PREFIX + name.replace('.', '/') + CLASS_DATA_SUFFIX; + JarEntry entry = agentJarFile.getJarEntry(entryName); + if (entry == null) { + throw new ClassNotFoundException(name); + } + + byte[] classBytes = readEntry(entry, name); + return defineClass(name, classBytes, 0, classBytes.length, agentCodeSource); + } + + @Override + protected URL findResource(String name) { + // For .class resource lookups, map to .classdata + String entryName; + if (name.endsWith(".class")) { + entryName = ENTRY_PREFIX + name.substring(0, name.length() - ".class".length()) + + CLASS_DATA_SUFFIX; + } else { + entryName = ENTRY_PREFIX + name; + } + + JarEntry entry = agentJarFile.getJarEntry(entryName); + if (entry != null) { + try { + return new URL(agentResourcePrefix + entryName); + } catch (java.net.MalformedURLException e) { + // fall through + } + } + return null; + } + + private byte[] readEntry(JarEntry entry, String className) throws ClassNotFoundException { + int size = (int) entry.getSize(); + byte[] buf = new byte[size]; + try (InputStream in = agentJarFile.getInputStream(entry)) { + int offset = 0; + while (offset < size) { + int bytesRead = in.read(buf, offset, size - offset); + if (bytesRead < 0) { + break; + } + offset += bytesRead; + } + if (offset != size) { + throw new ClassNotFoundException( + className + " (incomplete read: " + offset + "/" + size + " bytes)"); + } + return buf; + } catch (IOException e) { + throw new ClassNotFoundException(className, e); + } + } +} diff --git a/braintrust-java-agent/agent-internal/build.gradle b/braintrust-java-agent/agent-internal/build.gradle new file mode 100644 index 0000000..0962236 --- /dev/null +++ b/braintrust-java-agent/agent-internal/build.gradle @@ -0,0 +1,41 @@ +plugins { + id 'java' + id 'com.gradleup.shadow' version '8.3.6' +} + +java { + toolchain { + languageVersion = JavaLanguageVersion.of(17) + vendor = JvmVendorSpec.ADOPTIUM + } +} + +tasks.withType(JavaCompile).configureEach { + options.release = 17 +} + +repositories { + mavenCentral() +} + +dependencies { + // ByteBuddy for bytecode manipulation + implementation 'net.bytebuddy:byte-buddy:1.17.5' +} + +// Build a shadow JAR with all dependencies relocated to avoid conflicts. +// This JAR is NOT the final agent — it's an intermediate artifact whose +// contents get repacked as .classdata entries in the final agent JAR. +shadowJar { + archiveClassifier.set('all') + + // Relocate all external deps under our namespace + relocate 'net.bytebuddy', 'dev.braintrust.agent.shaded.bytebuddy' + + mergeServiceFiles() + + // Strip out empty directory stubs left behind by relocation + exclude 'META-INF/versions/*/net/**' + // Strip module-info that could confuse JPMS + exclude 'module-info.class' +} diff --git a/braintrust-java-agent/agent-internal/src/main/java/dev/braintrust/agent/internal/AgentInstaller.java b/braintrust-java-agent/agent-internal/src/main/java/dev/braintrust/agent/internal/AgentInstaller.java new file mode 100644 index 0000000..3dd8290 --- /dev/null +++ b/braintrust-java-agent/agent-internal/src/main/java/dev/braintrust/agent/internal/AgentInstaller.java @@ -0,0 +1,32 @@ +package dev.braintrust.agent.internal; + +import java.lang.instrument.Instrumentation; + +/** + * The real agent installation logic, loaded by {@code BraintrustClassLoader} in an isolated + * classloader. + * + *
This class and all its dependencies (ByteBuddy, OTel SDK, etc.) are invisible to the + * application's classpath. They exist as {@code .classdata} entries inside the agent JAR and are + * only accessible through the agent's custom classloader. + */ +public class AgentInstaller { + + /** + * Called reflectively from {@code BraintrustAgent.premain()} via the isolated classloader. + * + * @param agentArgs the agent arguments from the {@code -javaagent} flag + * @param inst the JVM instrumentation instance + */ + public static void install(String agentArgs, Instrumentation inst) { + System.out.println("[braintrust] AgentInstaller.install() called"); + System.out.println("[braintrust] AgentInstaller classloader: " + + AgentInstaller.class.getClassLoader().getClass().getName()); + System.out.println("[braintrust] Agent args: " + agentArgs); + System.out.println("[braintrust] Instrumentation: retransform=" + + inst.isRetransformClassesSupported()); + + // TODO: This is where ByteBuddy AgentBuilder setup will go. + // For now, just prove the classloader isolation works. + } +} diff --git a/braintrust-java-agent/build.gradle b/braintrust-java-agent/build.gradle new file mode 100644 index 0000000..1aa7831 --- /dev/null +++ b/braintrust-java-agent/build.gradle @@ -0,0 +1,56 @@ +plugins { + id 'base' +} + +configurations { + bootstrap + internal +} + +dependencies { + bootstrap project(':braintrust-java-agent:agent-bootstrap') + internal project(path: ':braintrust-java-agent:agent-internal', configuration: 'shadow') +} + +/** + * Assembles the final agent JAR with classloader isolation. + * + * Layout: + * dev/braintrust/agent/BraintrustAgent.class <- bootstrap (normal .class) + * dev/braintrust/agent/BraintrustClassLoader.class <- bootstrap (normal .class) + * inst/dev/braintrust/agent/internal/AgentInstaller.classdata <- hidden from normal CL + * inst/dev/braintrust/agent/shaded/bytebuddy/...classdata <- hidden from normal CL + */ +task agentJar(type: Jar) { + dependsOn configurations.bootstrap, configurations.internal + + destinationDirectory.set(layout.buildDirectory.dir('libs')) + archiveBaseName.set('braintrust-java-agent') + archiveClassifier.set('') + + manifest { + attributes( + 'Premain-Class': 'dev.braintrust.agent.BraintrustAgent', + 'Agent-Class': 'dev.braintrust.agent.BraintrustAgent', + 'Can-Redefine-Classes': 'true', + 'Can-Retransform-Classes': 'true', + 'Implementation-Title': 'Braintrust Java Agent', + 'Implementation-Version': project.version, + 'Implementation-Vendor': 'Braintrust', + ) + } + + // 1) Include bootstrap classes as normal .class files + from(provider { configurations.bootstrap.collect { zipTree(it) } }) { + exclude 'META-INF/**' + } + + // 2) Include internal classes as .classdata files under inst/ + from(provider { configurations.internal.collect { zipTree(it) } }) { + into 'inst' + rename '(.*)\\.class$', '$1.classdata' + exclude 'META-INF/**' + } +} + +assemble.dependsOn agentJar diff --git a/settings.gradle b/settings.gradle index cbaa458..810754f 100644 --- a/settings.gradle +++ b/settings.gradle @@ -5,3 +5,6 @@ plugins { rootProject.name = 'braintrust-sdk-java' include 'examples' +include 'braintrust-java-agent' +include 'braintrust-java-agent:agent-bootstrap' +include 'braintrust-java-agent:agent-internal'