squash various bot discovered bugs

braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustBedrockInterceptor.java

@@ -1,5 +1,6 @@
 package dev.braintrust.instrumentation.awsbedrock.v2_30_0;
 
+import dev.braintrust.instrumentation.ConverseStreamAccumulator;
 import dev.braintrust.instrumentation.InstrumentationSemConv;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Span;
@@ -8,7 +9,6 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.StringWriter;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
@@ -30,9 +30,6 @@
 import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest;
-import software.amazon.awssdk.thirdparty.jackson.core.JsonFactory;
-import software.amazon.awssdk.thirdparty.jackson.core.JsonParser;
-import software.amazon.awssdk.thirdparty.jackson.core.JsonToken;
 import software.amazon.eventstream.Message;
 import software.amazon.eventstream.MessageDecoder;
 
@@ -49,8 +46,6 @@ class BraintrustBedrockInterceptor implements ExecutionInterceptor {
     private static final ExecutionAttribute<String> MODEL_ID_ATTRIBUTE =
             new ExecutionAttribute<>("braintrust.modelId");
 
-    private static final JsonFactory JSON_FACTORY = new JsonFactory();
-
     private final Tracer tracer;
 
     BraintrustBedrockInterceptor(OpenTelemetry openTelemetry) {
@@ -234,21 +229,18 @@ private static String extractModelIdFromPath(String path) {
     }
 
     /**
-     * Tees the reactive byte stream into a {@link MessageDecoder}. On completion, decodes each
-     * event-stream frame using the AWS SDK's shaded Jackson streaming parser, accumulates the
-     * response content, and hands a synthetic Converse-shaped JSON string to semconv.
+     * Tees the reactive byte stream into a {@link MessageDecoder}, forwarding each decoded
+     * event-stream frame to a {@link ConverseStreamAccumulator}. On completion the accumulator's
+     * reconstructed Converse-shaped body is handed to semconv, so a streaming span carries the same
+     * content-block shapes — text, tool use, reasoning — as a synchronous one.
      */
     private static class TeeingSubscriber implements Subscriber<ByteBuffer> {
         private final Subscriber<? super ByteBuffer> downstream;
         private final Span span;
         private final Tracer tracer;
         private final MessageDecoder decoder = new MessageDecoder();
+        private final ConverseStreamAccumulator accumulator = new ConverseStreamAccumulator();
 
-        // Accumulated incrementally in onNext — no message list retained.
-        private final StringBuilder text = new StringBuilder();
-        private String stopReason = null;
-        private int inputTokens = 0;
-        private int outputTokens = 0;
         private long startNanos;
         private Long timeToFirstTokenNanos = null;
 
@@ -271,28 +263,18 @@ public void onNext(ByteBuffer buf) {
             try {
                 decoder.feed(copy);
                 for (Message msg : decoder.getDecodedMessages()) {
-                    var h = msg.getHeaders().get(":event-type");
-                    if (h == null) continue;
-                    String eventType = h.getString();
-                    byte[] payload = msg.getPayload();
-                    switch (eventType) {
-                        case "contentBlockDelta" -> {
-                            String t = parseDeltaText(payload);
-                            if (t != null) {
-                                text.append(t);
-                                if (timeToFirstTokenNanos == null) {
-                                    timeToFirstTokenNanos = System.nanoTime() - startNanos;
-                                }
-                            }
-                        }
-                        case "messageStop" -> stopReason = parseStopReason(payload);
-                        case "metadata" -> {
-                            int[] tokens = parseTokenUsage(payload);
-                            inputTokens = tokens[0];
-                            outputTokens = tokens[1];
-                        }
-                        default -> {}
+                    var header = msg.getHeaders().get(":event-type");
+                    if (header == null) continue;
+                    String eventType = header.getString();
+                    // First content frame marks time-to-first-token, whether the model opened with
+                    // text, a tool call, or a reasoning block.
+                    if (timeToFirstTokenNanos == null
+                            && ("contentBlockDelta".equals(eventType)
+                                    || "contentBlockStart".equals(eventType))) {
+                        timeToFirstTokenNanos = System.nanoTime() - startNanos;
                     }
+                    accumulator.accept(
+                            eventType, new String(msg.getPayload(), StandardCharsets.UTF_8));
                 }
             } catch (Exception e) {
                 log.debug("Failed to feed event-stream decoder", e);
@@ -312,107 +294,13 @@ public void onComplete() {
                         tracer,
                         span,
                         InstrumentationSemConv.PROVIDER_NAME_BEDROCK,
-                        buildConverseJson(text.toString(), stopReason, inputTokens, outputTokens),
+                        accumulator.build(),
                         timeToFirstTokenNanos);
             } catch (Exception e) {
                 log.debug("Failed to tag span from streaming response", e);
             } finally {
                 downstream.onComplete();
             }
         }
-
-        /**
-         * Parses {@code delta.text} from a {@code contentBlockDelta} payload: {@code
-         * {"contentBlockIndex":0,"delta":{"text":"...","type":"text_delta"}}}
-         */
-        private static String parseDeltaText(byte[] payload) throws Exception {
-            try (JsonParser p = JSON_FACTORY.createParser(payload)) {
-                boolean inDelta = false;
-                while (p.nextToken() != null) {
-                    if (p.currentToken() == JsonToken.FIELD_NAME) {
-                        if ("delta".equals(p.currentName())) {
-                            inDelta = true;
-                        } else if (inDelta && "text".equals(p.currentName())) {
-                            p.nextToken();
-                            return p.getText();
-                        }
-                    } else if (p.currentToken() == JsonToken.END_OBJECT) {
-                        inDelta = false;
-                    }
-                }
-            }
-            return null;
-        }
-
-        /**
-         * Parses {@code stopReason} from a {@code messageStop} payload: {@code
-         * {"stopReason":"end_turn"}}
-         */
-        private static String parseStopReason(byte[] payload) throws Exception {
-            try (JsonParser p = JSON_FACTORY.createParser(payload)) {
-                while (p.nextToken() != null) {
-                    if (p.currentToken() == JsonToken.FIELD_NAME
-                            && "stopReason".equals(p.currentName())) {
-                        p.nextToken();
-                        return p.getText();
-                    }
-                }
-            }
-            return null;
-        }
-
-        /**
-         * Parses {@code [inputTokens, outputTokens]} from a {@code metadata} payload: {@code
-         * {"usage":{"inputTokens":N,"outputTokens":M},"metrics":{...}}}
-         */
-        private static int[] parseTokenUsage(byte[] payload) throws Exception {
-            int inputTokens = 0;
-            int outputTokens = 0;
-            try (JsonParser p = JSON_FACTORY.createParser(payload)) {
-                while (p.nextToken() != null) {
-                    if (p.currentToken() == JsonToken.FIELD_NAME) {
-                        if ("inputTokens".equals(p.currentName())) {
-                            p.nextToken();
-                            inputTokens = p.getIntValue();
-                        } else if ("outputTokens".equals(p.currentName())) {
-                            p.nextToken();
-                            outputTokens = p.getIntValue();
-                        }
-                    }
-                }
-            }
-            return new int[] {inputTokens, outputTokens};
-        }
-
-        /**
-         * Builds a synthetic Converse-shaped JSON string matching what {@code tagBedrockResponse}
-         * expects, using the shaded Jackson generator for correct escaping.
-         */
-        private static String buildConverseJson(
-                String text, String stopReason, int inputTokens, int outputTokens)
-                throws Exception {
-            StringWriter sw = new StringWriter();
-            try (var gen = JSON_FACTORY.createGenerator(sw)) {
-                gen.writeStartObject();
-                gen.writeObjectFieldStart("output");
-                gen.writeObjectFieldStart("message");
-                gen.writeStringField("role", "assistant");
-                gen.writeArrayFieldStart("content");
-                gen.writeStartObject();
-                gen.writeStringField("text", text);
-                gen.writeEndObject();
-                gen.writeEndArray();
-                gen.writeEndObject(); // message
-                gen.writeEndObject(); // output
-                gen.writeStringField("stopReason", stopReason != null ? stopReason : "end_turn");
-                gen.writeObjectFieldStart("usage");
-                gen.writeNumberField("inputTokens", inputTokens);
-                gen.writeNumberField("outputTokens", outputTokens);
-                gen.writeNumberField("totalTokens", inputTokens + outputTokens);
-                gen.writeEndObject(); // usage
-                gen.writeEndObject();
-            }
-            return sw.toString();
-        }
     }
 }

braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/auto/AWSBedrockInstrumentationModule.java

@@ -39,7 +39,8 @@ public List<String> getHelperClassNames() {
                 MANUAL_INSTRUMENTATION_PACKAGE + "BraintrustBedrockInterceptor",
                 MANUAL_INSTRUMENTATION_PACKAGE + "BraintrustBedrockInterceptor$TeeingSubscriber",
                 "dev.braintrust.json.BraintrustJsonMapper",
-                "dev.braintrust.instrumentation.InstrumentationSemConv");
+                "dev.braintrust.instrumentation.InstrumentationSemConv",
+                "dev.braintrust.instrumentation.ConverseStreamAccumulator");
     }
 
     @Override

braintrust-sdk/instrumentation/langchain_1_14_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_14_0/BraintrustLangchain.java

@@ -64,16 +64,27 @@ public static <T> T wrap(OpenTelemetry openTelemetry, AiServices<T> aiServices)
 
             if (context.toolService != null) {
                 // ////// CREATE A SPAN FOR EACH TOOL CALL
+                // toolExecutors() hands back the live map on the AiServiceContext, and that
+                // context outlives build() and is shared with every service already built from
+                // this builder. Wrapping unconditionally would therefore nest a second tracing
+                // executor on the next build — duplicating every tool span, for the earlier
+                // services too — so already-wrapped entries are left alone. Same reasoning as the
+                // `instanceof WrappedHttpClient` guards on the model paths below.
                 for (Map.Entry<String, ToolExecutor> entry :
                         context.toolService.toolExecutors().entrySet()) {
                     String toolName = entry.getKey();
                     ToolExecutor original = entry.getValue();
+                    if (original instanceof TracingToolExecutor) {
+                        log.debug("tool already instrumented. skipping: {}", toolName);
+                        continue;
+                    }
                     entry.setValue(new TracingToolExecutor(original, toolName, tracer));
                 }
 
                 // ////// LINK SPANS ACROSS CONCURRENT TOOL CALLS
                 var underlyingExecutor = context.toolService.executor();
-                if (underlyingExecutor != null) {
+                if (underlyingExecutor != null
+                        && !(underlyingExecutor instanceof OtelContextPassingExecutor)) {
                     aiServices.executeToolsConcurrently(
                             new OtelContextPassingExecutor(underlyingExecutor));
                 }

braintrust-sdk/instrumentation/langchain_1_14_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_14_0/WrappedHttpClient.java

@@ -15,6 +15,7 @@
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Scope;
 import java.net.URI;
@@ -133,11 +134,24 @@ static class WrappedServerSentEventListener implements ServerSentEventListener {
         private final String providerName;
         private final Tracer tracer;
         private final long startNanos = System.nanoTime();
-        private final AtomicLong timeToFirstTokenNanos = new AtomicLong();
+        // Time-to-first-token is measured from the first payload that carries generated output,
+        // not the first payload of any kind — a Responses stream opens with response.created
+        // before the model has produced anything. firstPayloadNanos is a fallback for streams
+        // whose shape is not recognized at all; sawRecognizedShape is what keeps that fallback
+        // from firing on a recognized stream that simply never produced output, where the honest
+        // answer is that there was no first token.
+        private final AtomicLong firstOutputNanos = new AtomicLong();
+        private final AtomicLong firstPayloadNanos = new AtomicLong();
+        private volatile boolean sawRecognizedShape;
         // Handles both endpoints this module instruments: chat-completions chunk streams and
         // Responses API (`/v1/responses`) event streams.
         private final SseStreamAccumulator accumulator =
                 new SseStreamAccumulator(BraintrustJsonMapper.get());
+        // A stream can report a failed generation in band, after the HTTP request itself has
+        // succeeded. LangChain4j delivers those failures to the caller's own response handler and
+        // then closes the transport normally, so onError below is never reached — retaining the
+        // failure here is what stops onClose from finalizing a failed call as a successful span.
+        @javax.annotation.Nullable private volatile String streamFailure;
 
         WrappedServerSentEventListener(
                 ServerSentEventListener delegate, Span span, String providerName, Tracer tracer) {
@@ -157,15 +171,15 @@ public void onOpen(SuccessfulHttpResponse response) {
         @Override
         public void onEvent(ServerSentEvent event, ServerSentEventContext context) {
             try (Scope ignored = span.makeCurrent()) {
-                accumulateChunk(event.data());
+                accumulateEvent(event);
                 delegate.onEvent(event, context);
             }
         }
 
         @Override
         public void onEvent(ServerSentEvent event) {
             try (Scope ignored = span.makeCurrent()) {
-                accumulateChunk(event.data());
+                accumulateEvent(event);
                 delegate.onEvent(event);
             }
         }
@@ -190,17 +204,62 @@ public void onClose() {
             }
         }
 
-        private void accumulateChunk(String data) {
+        private void accumulateEvent(ServerSentEvent event) {
+            String data = event.data();
+            if (streamFailure == null) {
+                streamFailure =
+                        SseStreamAccumulator.streamFailure(
+                                BraintrustJsonMapper.get(), event.event(), data);
+            }
             if (data == null || data.isEmpty() || "[DONE]".equals(data)) return;
-            if (timeToFirstTokenNanos.get() == 0L) {
-                timeToFirstTokenNanos.compareAndExchange(0L, System.nanoTime() - startNanos);
+            firstPayloadNanos.compareAndExchange(0L, System.nanoTime() - startNanos);
+            // Only classify until the first output is seen; afterwards this is a single volatile
+            // read per chunk.
+            if (firstOutputNanos.get() == 0L) {
+                var kind = SseStreamAccumulator.classify(BraintrustJsonMapper.get(), data);
+                if (kind != SseStreamAccumulator.PayloadKind.UNRECOGNIZED) {
+                    sawRecognizedShape = true;
+                }
+                if (kind == SseStreamAccumulator.PayloadKind.OUTPUT) {
+                    firstOutputNanos.compareAndExchange(0L, System.nanoTime() - startNanos);
+                }
             }
             accumulator.merge(data);
         }
 
+        /**
+         * Source for {@code time_to_first_token}, or {@code null} when the stream produced no first
+         * token to time.
+         *
+         * <p>The first generated output when there was one. Otherwise the first payload, but only
+         * for a stream whose shape was never recognized — there the timestamp is a slightly early
+         * approximation, which beats dropping a metric the spec requires for streaming spans. A
+         * recognized stream that produced no output (one that failed before generating, or
+         * completed empty) reports nothing: its first payload is lifecycle metadata, and publishing
+         * that as a token latency would silently corrupt latency aggregates.
+         */
+        @javax.annotation.Nullable
+        private Long timeToFirstTokenNanos() {
+            long output = firstOutputNanos.get();
+            if (output != 0L) {
+                return output;
+            }
+            if (sawRecognizedShape) {
+                return null;
+            }
+            long payload = firstPayloadNanos.get();
+            return payload != 0L ? payload : null;
+        }
+
         private void finalizeSpan() {
+            String failure = streamFailure;
+            if (failure != null) {
+                // Recorded before tagging: a failed stream's body is often partial, and losing the
+                // error status to a tagging problem is worse than losing the partial output.
+                span.setStatus(StatusCode.ERROR, failure);
+            }
             try {
-                Long ttft = timeToFirstTokenNanos.get();
+                Long ttft = timeToFirstTokenNanos();
                 String responseBody = accumulator.build();
                 InstrumentationSemConv.tagLLMSpanResponse(
                         tracer, span, providerName, responseBody, ttft);

braintrust-sdk/instrumentation/langchain_1_14_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_14_0/auto/LangchainInstrumentationModule.java

@@ -53,6 +53,7 @@ public List<String> getHelperClassNames() {
                 MANUAL_PACKAGE + "TracingToolExecutor",
                 MANUAL_PACKAGE + "OtelContextPassingExecutor",
                 "dev.braintrust.instrumentation.SseStreamAccumulator",
+                "dev.braintrust.instrumentation.SseStreamAccumulator$PayloadKind",
                 "dev.braintrust.instrumentation.SseResponseAccumulator",
                 "dev.braintrust.instrumentation.InstrumentationSemConv",
                 "dev.braintrust.json.BraintrustJsonMapper");