springai 2 instrumentation

- enhance vanilla openai / anthropic instrumentation
- add module for springai 2
- add multi version support to btx

Author: realark

braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropic.java

@@ -1,14 +1,13 @@
 package dev.braintrust.instrumentation.anthropic.v2_2_0;
 
 import com.anthropic.client.AnthropicClient;
+import com.anthropic.client.AnthropicClientAsync;
 import com.anthropic.core.ClientOptions;
 import com.anthropic.core.http.HttpClient;
 import io.opentelemetry.api.OpenTelemetry;
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
-import java.util.function.BiConsumer;
 import java.util.function.Consumer;
-import kotlin.Lazy;
 import lombok.extern.slf4j.Slf4j;
 
 /** Braintrust Anthropic client instrumentation. */
@@ -17,48 +16,79 @@ public final class BraintrustAnthropic {
 
     /** Instrument Anthropic client with Braintrust traces. */
     public static AnthropicClient wrap(OpenTelemetry openTelemetry, AnthropicClient client) {
+        if (!instrument(openTelemetry, client)) {
+            return client;
+        }
+        return ContextCapturingProxy.wrap(client, AnthropicClient.class);
+    }
+
+    /** Instrument an async Anthropic client with Braintrust traces. */
+    public static AnthropicClientAsync wrap(
+            OpenTelemetry openTelemetry, AnthropicClientAsync client) {
+        if (!instrument(openTelemetry, client)) {
+            return client;
+        }
+        return ContextCapturingProxy.wrap(client, AnthropicClientAsync.class);
+    }
+
+    /**
+     * Swaps the client's {@code ClientOptions.httpClient} for a {@link TracingHttpClient} in place;
+     * wrapping is idempotent.
+     *
+     * @return whether the HTTP layer is instrumented. When false (custom client implementations or
+     *     changed SDK internals), the caller must NOT install {@link ContextCapturingProxy}: the
+     *     proxy's internal context header is only stripped by {@link TracingHttpClient}, so
+     *     installing it without one would leak trace/span IDs to the provider.
+     */
+    private static boolean instrument(OpenTelemetry openTelemetry, Object client) {
+        if (ContextCapturingProxy.isContextCapturingProxy(client)) {
+            // already instrumented
+            return true;
+        }
         try {
             instrumentHttpClient(openTelemetry, client);
-            return client;
+            return true;
         } catch (Exception e) {
-            log.error("failed to apply anthropic instrumentation", e);
-            return client;
+            log.error(
+                    "failed to apply anthropic instrumentation to {} — leaving client untouched",
+                    client.getClass().getName(),
+                    e);
+            return false;
         }
     }
 
-    private static void instrumentHttpClient(OpenTelemetry openTelemetry, AnthropicClient client) {
+    private static void instrumentHttpClient(OpenTelemetry openTelemetry, Object client) {
+        int[] instrumented = {0};
         forAllFields(
                 client,
                 fieldName -> {
                     try {
-                        var field = getField(client, fieldName);
-                        if (field instanceof ClientOptions clientOptions) {
-                            instrumentClientOptions(
-                                    openTelemetry, clientOptions, "originalHttpClient");
-                            instrumentClientOptions(openTelemetry, clientOptions, "httpClient");
-                        } else if (field instanceof Lazy<?> lazyField) {
-                            var resolved = lazyField.getValue();
-                            forAllFieldsOfType(
-                                    resolved,
-                                    ClientOptions.class,
-                                    (clientOptions, subfieldName) ->
-                                            instrumentClientOptions(
-                                                    openTelemetry, clientOptions, subfieldName));
-                        } else {
-                            forAllFieldsOfType(
-                                    field,
-                                    ClientOptions.class,
-                                    (clientOptions, subfieldName) ->
-                                            instrumentClientOptions(
-                                                    openTelemetry, clientOptions, subfieldName));
+                        if (getField(client, fieldName) instanceof ClientOptions clientOptions) {
+                            instrumentClientOptions(openTelemetry, clientOptions);
+                            instrumented[0]++;
                         }
                     } catch (ReflectiveOperationException e) {
                         throw new RuntimeException(e);
                     }
                 });
+        if (instrumented[0] == 0) {
+            // Finding nothing is as much a failure as a reflection error: the request path
+            // would bypass TracingHttpClient entirely.
+            throw new IllegalStateException(
+                    "no ClientOptions field found on "
+                            + client.getClass().getName()
+                            + " — unrecognized client shape");
+        }
     }
 
+    /** Swaps both HTTP client fields on a {@link ClientOptions} for tracing wrappers. */
     private static void instrumentClientOptions(
+            OpenTelemetry openTelemetry, ClientOptions clientOptions) {
+        swapHttpClient(openTelemetry, clientOptions, "originalHttpClient");
+        swapHttpClient(openTelemetry, clientOptions, "httpClient");
+    }
+
+    private static void swapHttpClient(
             OpenTelemetry openTelemetry, ClientOptions clientOptions, String fieldName) {
         try {
             HttpClient httpClient = getField(clientOptions, fieldName);
@@ -84,21 +114,6 @@ private static void forAllFields(Object object, Consumer<String> consumer) {
         }
     }
 
-    private static <T> void forAllFieldsOfType(
-            Object object, Class<T> targetClazz, BiConsumer<T, String> consumer) {
-        forAllFields(
-                object,
-                fieldName -> {
-                    try {
-                        if (targetClazz.isAssignableFrom(object.getClass())) {
-                            consumer.accept(getField(object, fieldName), fieldName);
-                        }
-                    } catch (ReflectiveOperationException e) {
-                        throw new RuntimeException(e);
-                    }
-                });
-    }
-
     @SuppressWarnings("unchecked")
     private static <T> T getField(Object obj, String fieldName)
             throws ReflectiveOperationException {

braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/ContextCapturingProxy.java

@@ -0,0 +1,179 @@
+package dev.braintrust.instrumentation.anthropic.v2_2_0;
+
+import com.anthropic.core.Params;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Captures the caller's OTel context at the service-call boundary of a wrapped anthropic-java
+ * client.
+ *
+ * <p>anthropic-java dispatches async requests through CompletableFuture continuations on the common
+ * pool, so by the time {@link TracingHttpClient#executeAsync} runs, the caller's thread-local
+ * context (e.g. an active application span) is gone. The service-method invocation itself, however,
+ * happens on the caller's thread — and the only data that travels from there into the HTTP layer is
+ * the request. So this proxy rewrites any {@link Params} argument to carry the current span as an
+ * internal {@code traceparent}-format header, which {@link TracingHttpClient} extracts (and strips
+ * from the outgoing request) to parent the LLM span.
+ *
+ * <p>Purely reflective and API-shape based, so it works for every service and endpoint: methods
+ * returning other {@code com.anthropic} interfaces (service accessors like {@code messages()},
+ * {@code async()}) return proxied instances, so context capture follows the whole call graph.
+ */
+@Slf4j
+final class ContextCapturingProxy implements InvocationHandler {
+
+    /** Internal correlation header; never sent — TracingHttpClient removes it. */
+    static final String CONTEXT_HEADER = "x-braintrust-instrumentation-context";
+
+    /** Per-params-class reflection handles: [toBuilder, putAdditionalHeader, build]. */
+    private static final Map<Class<?>, Method[]> PARAMS_METHODS = new ConcurrentHashMap<>();
+
+    private static final Method[] UNSUPPORTED = new Method[0];
+
+    private final Object delegate;
+
+    private ContextCapturingProxy(Object delegate) {
+        this.delegate = delegate;
+    }
+
+    /** Wraps {@code delegate} in a context-capturing proxy of {@code iface}. Idempotent. */
+    @SuppressWarnings("unchecked")
+    static <T> T wrap(T delegate, Class<T> iface) {
+        if (delegate == null || isContextCapturingProxy(delegate)) {
+            return delegate;
+        }
+        return (T)
+                Proxy.newProxyInstance(
+                        iface.getClassLoader(),
+                        new Class<?>[] {iface},
+                        new ContextCapturingProxy(delegate));
+    }
+
+    static boolean isContextCapturingProxy(Object o) {
+        return o != null
+                && Proxy.isProxyClass(o.getClass())
+                && Proxy.getInvocationHandler(o) instanceof ContextCapturingProxy;
+    }
+
+    @Override
+    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+        // equals/hashCode/toString are the only Object methods routed to an InvocationHandler.
+        // They must not be forwarded like service methods: delegate.equals(proxy) is false, so a
+        // forwarded equals would make proxy.equals(proxy) false — violating reflexivity and
+        // breaking map/set usage of wrapped clients.
+        if (method.getDeclaringClass() == Object.class) {
+            return switch (method.getName()) {
+                case "equals" -> proxy == args[0] || delegateEquals(args[0]);
+                case "hashCode" -> delegate.hashCode();
+                default -> "ContextCapturingProxy(" + delegate + ")"; // toString
+            };
+        }
+        Object[] invokeArgs = injectContextHeader(args);
+        Object result;
+        try {
+            result = method.invoke(delegate, invokeArgs);
+        } catch (InvocationTargetException e) {
+            throw e.getCause();
+        }
+        // Follow the service graph: accessors like messages(), async() return com.anthropic
+        // interfaces whose methods must also capture context.
+        Class<?> returnType = method.getReturnType();
+        if (result != null
+                && returnType.isInterface()
+                && returnType.getName().startsWith("com.anthropic.")) {
+            return Proxy.newProxyInstance(
+                    returnType.getClassLoader(),
+                    new Class<?>[] {returnType},
+                    new ContextCapturingProxy(result));
+        }
+        return result;
+    }
+
+    /** Two proxies are equal iff their delegates are (consistent with delegate-based hashCode). */
+    private boolean delegateEquals(Object other) {
+        return isContextCapturingProxy(other)
+                && delegate.equals(
+                        ((ContextCapturingProxy) Proxy.getInvocationHandler(other)).delegate);
+    }
+
+    /** Rewrites any {@link Params} argument to carry the current span as an internal header. */
+    private Object[] injectContextHeader(Object[] args) {
+        if (args == null) {
+            return null;
+        }
+        String traceparent = currentTraceparent();
+        if (traceparent == null) {
+            return args;
+        }
+        Object[] result = args;
+        for (int i = 0; i < args.length; i++) {
+            if (args[i] instanceof Params params) {
+                Object rewritten = withContextHeader(params, traceparent);
+                if (rewritten != null) {
+                    if (result == args) {
+                        result = args.clone();
+                    }
+                    result[i] = rewritten;
+                }
+            }
+        }
+        return result;
+    }
+
+    private static String currentTraceparent() {
+        SpanContext spanContext = Span.current().getSpanContext();
+        if (!spanContext.isValid()) {
+            return null;
+        }
+        return "00-"
+                + spanContext.getTraceId()
+                + "-"
+                + spanContext.getSpanId()
+                + "-"
+                + spanContext.getTraceFlags().asHex();
+    }
+
+    /**
+     * {@code params.toBuilder().putAdditionalHeader(CONTEXT_HEADER, traceparent).build()}, done
+     * reflectively so it works for every generated params type. Returns {@code null} (leaving the
+     * original untouched) when the shape doesn't match.
+     */
+    private static Object withContextHeader(Params params, String traceparent) {
+        Method[] methods =
+                PARAMS_METHODS.computeIfAbsent(
+                        params.getClass(), ContextCapturingProxy::resolveParamsMethods);
+        if (methods == UNSUPPORTED) {
+            return null;
+        }
+        try {
+            Object builder = methods[0].invoke(params);
+            methods[1].invoke(builder, CONTEXT_HEADER, traceparent);
+            return methods[2].invoke(builder);
+        } catch (Exception e) {
+            log.debug("failed to inject context header into {}", params.getClass().getName(), e);
+            return null;
+        }
+    }
+
+    private static Method[] resolveParamsMethods(Class<?> paramsClass) {
+        try {
+            Method toBuilder = paramsClass.getMethod("toBuilder");
+            Class<?> builderClass = toBuilder.getReturnType();
+            Method putHeader =
+                    builderClass.getMethod("putAdditionalHeader", String.class, String.class);
+            Method build = builderClass.getMethod("build");
+            return new Method[] {toBuilder, putHeader, build};
+        } catch (NoSuchMethodException e) {
+            log.debug("params type {} has no header builder shape", paramsClass.getName());
+            return UNSUPPORTED;
+        }
+    }
+}