Add span origin provenance

braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java

@@ -3,10 +3,14 @@
 import dev.braintrust.Braintrust;
 import dev.braintrust.api.BraintrustOpenApiClient;
 import java.net.URI;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
+import javax.annotation.Nullable;
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.TrustManagerFactory;
 import javax.net.ssl.X509TrustManager;
@@ -39,6 +43,8 @@ public final class BraintrustConfig extends BaseConfig {
     private final Duration requestTimeout =
             Duration.ofSeconds(getConfig("BRAINTRUST_REQUEST_TIMEOUT", 30));
     private final Boolean filterAISpans = getConfig("BRAINTRUST_FILTER_AI_SPANS", false);
+    private final Optional<SpanOriginEnvironment> spanOriginEnvironment =
+            Optional.ofNullable(detectSpanOriginEnvironment());
 
     /** compress otel data before exporting to braintrust */
     private final Boolean compressOtelPayload = getConfig("BRAINTRUST_COMPRESS_OTEL_PAYLOAD", true);
@@ -206,6 +212,29 @@ public Builder filterAISpans(boolean value) {
             return this;
         }
 
+        public Builder environment(String type) {
+            return environment(type, null);
+        }
+
+        public Builder environment(String type, String name) {
+            envOverrides.put("BRAINTRUST_ENVIRONMENT_TYPE", type);
+            if (name != null) {
+                envOverrides.put("BRAINTRUST_ENVIRONMENT_NAME", name);
+            } else {
+                envOverrides.put("BRAINTRUST_ENVIRONMENT_NAME", NULL_OVERRIDE);
+            }
+            return this;
+        }
+
+        public Builder spanOriginEnvironment(SpanOriginEnvironment value) {
+            if (value != null) {
+                return environment(value.type(), value.name());
+            }
+            envOverrides.put("BRAINTRUST_ENVIRONMENT_TYPE", NULL_OVERRIDE);
+            envOverrides.put("BRAINTRUST_ENVIRONMENT_NAME", NULL_OVERRIDE);
+            return this;
+        }
+
         public Builder compressOtelPayload(boolean value) {
             envOverrides.put("BRAINTRUST_COMPRESS_OTEL_PAYLOAD", String.valueOf(value));
             return this;
@@ -235,4 +264,135 @@ public BraintrustConfig build() {
             return new BraintrustConfig(envOverrides, sslContext, x509TrustManager);
         }
     }
+
+    private SpanOriginEnvironment detectSpanOriginEnvironment() {
+        var environmentType = getEnvironmentConfigValue("BRAINTRUST_ENVIRONMENT_TYPE");
+        var environmentName = getEnvironmentConfigValue("BRAINTRUST_ENVIRONMENT_NAME");
+        if (environmentType != null || environmentName != null) {
+            return new SpanOriginEnvironment(environmentType, environmentName);
+        }
+        if (envOverrides.containsKey("BRAINTRUST_ENVIRONMENT_TYPE")
+                || envOverrides.containsKey("BRAINTRUST_ENVIRONMENT_NAME")) {
+            return null;
+        }
+
+        var serverName = detectServerEnvironmentName();
+        if (serverName != null) {
+            return new SpanOriginEnvironment("server", serverName);
+        }
+
+        var ciName =
+                detectFirstPresent(
+                        Map.ofEntries(
+                                Map.entry("GITHUB_ACTIONS", "github_actions"),
+                                Map.entry("GITLAB_CI", "gitlab_ci"),
+                                Map.entry("CIRCLECI", "circleci"),
+                                Map.entry("BUILDKITE", "buildkite"),
+                                Map.entry("JENKINS_URL", "jenkins"),
+                                Map.entry("JENKINS_HOME", "jenkins"),
+                                Map.entry("TF_BUILD", "azure_pipelines"),
+                                Map.entry("TEAMCITY_VERSION", "teamcity"),
+                                Map.entry("TRAVIS", "travis"),
+                                Map.entry("BITBUCKET_BUILD_NUMBER", "bitbucket")));
+        if (ciName != null) {
+            return new SpanOriginEnvironment("ci", ciName);
+        }
+        if (isPresent("CI")) {
+            return new SpanOriginEnvironment("ci", "ci");
+        }
+
+        return null;
+    }
+
+    private @Nullable String getEnvironmentConfigValue(String key) {
+        var value = getEnvValue(key);
+        if (envOverrides.containsKey(key)) {
+            return firstNonBlank(value);
+        }
+        return firstNonBlank(value, getBraintrustEnvFileValue(key));
+    }
+
+    private @Nullable String detectServerEnvironmentName() {
+        var serverName =
+                detectFirstPresent(
+                        Map.ofEntries(
+                                Map.entry("VERCEL", "vercel"), Map.entry("NETLIFY", "netlify")));
+        if (serverName != null) {
+            return serverName;
+        }
+        if (isPresent("ECS_CONTAINER_METADATA_URI") || isPresent("ECS_CONTAINER_METADATA_URI_V4")) {
+            return "ecs";
+        }
+        var awsExecutionEnv = firstNonBlank(getEnvValue("AWS_EXECUTION_ENV"));
+        if (awsExecutionEnv != null) {
+            if (awsExecutionEnv.startsWith("AWS_ECS_")) {
+                return "ecs";
+            }
+            if (awsExecutionEnv.startsWith("AWS_Lambda_")) {
+                return "aws_lambda";
+            }
+        }
+        if (isPresent("AWS_LAMBDA_FUNCTION_NAME")) {
+            return "aws_lambda";
+        }
+        return detectFirstPresent(
+                Map.ofEntries(
+                        Map.entry("K_SERVICE", "cloud_run"),
+                        Map.entry("FUNCTION_TARGET", "gcp_functions"),
+                        Map.entry("KUBERNETES_SERVICE_HOST", "kubernetes"),
+                        Map.entry("DYNO", "heroku"),
+                        Map.entry("FLY_APP_NAME", "fly"),
+                        Map.entry("RAILWAY_ENVIRONMENT", "railway"),
+                        Map.entry("RENDER_SERVICE_NAME", "render")));
+    }
+
+    private @Nullable String detectFirstPresent(Map<String, String> vars) {
+        return vars.entrySet().stream()
+                .filter(entry -> isPresent(entry.getKey()))
+                .map(Map.Entry::getValue)
+                .findFirst()
+                .orElse(null);
+    }
+
+    private boolean isPresent(String key) {
+        return firstNonBlank(getEnvValue(key)) != null;
+    }
+
+    private static @Nullable String firstNonBlank(String... values) {
+        return Arrays.stream(values)
+                .filter(value -> value != null && !value.isBlank())
+                .map(String::trim)
+                .findFirst()
+                .orElse(null);
+    }
+
+    private static @Nullable String getBraintrustEnvFileValue(String key) {
+        for (var dir = Path.of("").toAbsolutePath(); dir != null; dir = dir.getParent()) {
+            var path = dir.resolve(".env.braintrust");
+            if (!Files.isRegularFile(path)) {
+                continue;
+            }
+            try {
+                for (var line : Files.readAllLines(path)) {
+                    var trimmed = line.trim();
+                    if (trimmed.isEmpty() || trimmed.startsWith("#")) {
+                        continue;
+                    }
+                    var equals = trimmed.indexOf('=');
+                    if (equals <= 0 || !trimmed.substring(0, equals).trim().equals(key)) {
+                        continue;
+                    }
+                    var value = trimmed.substring(equals + 1).trim();
+                    if ((value.startsWith("\"") && value.endsWith("\""))
+                            || (value.startsWith("'") && value.endsWith("'"))) {
+                        value = value.substring(1, value.length() - 1);
+                    }
+                    return value;
+                }
+            } catch (Exception ignored) {
+                return null;
+            }
+        }
+        return null;
+    }
 }

braintrust-sdk/src/main/java/dev/braintrust/config/SpanOriginEnvironment.java

@@ -0,0 +1,5 @@
+package dev.braintrust.config;
+
+import javax.annotation.Nullable;
+
+public record SpanOriginEnvironment(@Nullable String type, @Nullable String name) {}

braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java

@@ -21,7 +21,8 @@ public boolean sample(ReadableSpan span) {
             for (var key : span.getAttributes().asMap().keySet()) {
                 var keyName = key.getKey();
                 // Skip internal attributes injected by the span processor itself
-                if (keyName.equals(BraintrustTracing.PARENT_KEY)) {
+                if (keyName.equals(BraintrustTracing.PARENT_KEY)
+                        || keyName.equals(BraintrustSpanProcessor.CONTEXT_JSON.getKey())) {
                     continue;
                 }
                 for (var prefix : AI_ATTR_PREFIXES) {

braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java

@@ -2,6 +2,8 @@
 
 import dev.braintrust.api.BraintrustOpenApiClient;
 import dev.braintrust.config.BraintrustConfig;
+import dev.braintrust.config.SpanOriginEnvironment;
+import dev.braintrust.json.BraintrustJsonMapper;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.context.Context;
@@ -12,7 +14,9 @@
 import io.opentelemetry.sdk.trace.SpanProcessor;
 import io.opentelemetry.sdk.trace.data.DelegatingSpanData;
 import io.opentelemetry.sdk.trace.data.SpanData;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
@@ -33,6 +37,8 @@ public class BraintrustSpanProcessor implements SpanProcessor {
 
     public static final AttributeKey<String> PARENT =
             AttributeKey.stringKey(BraintrustTracing.PARENT_KEY);
+    static final AttributeKey<String> CONTEXT_JSON =
+            AttributeKey.stringKey("braintrust.context_json");
 
     private final BraintrustConfig config;
     private final SpanProcessor delegate;
@@ -109,6 +115,42 @@ public void onStart(@Nonnull Context parentContext, ReadWriteSpan span) {
         delegate.onStart(parentContext, span);
     }
 
+    @SuppressWarnings("unchecked")
+    private static String mergedContextJson(
+            @Nullable String existingContextJson, @Nullable SpanOriginEnvironment environment) {
+        Map<String, Object> context = new HashMap<>();
+        if (existingContextJson != null && !existingContextJson.isBlank()) {
+            try {
+                context.putAll(
+                        BraintrustJsonMapper.get().readValue(existingContextJson, Map.class));
+            } catch (Exception ignored) {
+                // Invalid user-provided context should not prevent adding provenance.
+            }
+        }
+
+        var existingOrigin = context.get("span_origin");
+        Map<String, Object> spanOrigin =
+                existingOrigin instanceof Map<?, ?> map
+                        ? new HashMap<>((Map<String, Object>) map)
+                        : new HashMap<>();
+        spanOrigin.putIfAbsent("name", "braintrust.sdk.java");
+        spanOrigin.putIfAbsent("version", BraintrustTracing.INSTRUMENTATION_VERSION);
+        spanOrigin.putIfAbsent(
+                "instrumentation", Map.of("name", BraintrustTracing.INSTRUMENTATION_NAME));
+        if (!spanOrigin.containsKey("environment") && environment != null) {
+            var env = new HashMap<String, Object>();
+            if (environment.type() != null && !environment.type().isBlank()) {
+                env.put("type", environment.type());
+            }
+            if (environment.name() != null && !environment.name().isBlank()) {
+                env.put("name", environment.name());
+            }
+            spanOrigin.put("environment", env);
+        }
+        context.put("span_origin", spanOrigin);
+        return BraintrustJsonMapper.toJson(context);
+    }
+
     @Override
     public boolean isStartRequired() {
         return true;
@@ -135,18 +177,30 @@ public void onEnd(ReadableSpan span) {
 
         @Nullable String newInputJson = attachmentProcessor.processAndUpload(inputJson);
         @Nullable String newOutputJson = attachmentProcessor.processAndUpload(outputJson);
+        var spanOriginEnvironment = config.spanOriginEnvironment().orElse(null);
+        var newContextJson =
+                mergedContextJson(
+                        spanData.getAttributes().get(CONTEXT_JSON), spanOriginEnvironment);
 
         if (!Objects.equals(newInputJson, inputJson)
                 || !Objects.equals(newOutputJson, outputJson)) {
-            delegate.onEnd(new TransformedReadableSpan(span, newInputJson, newOutputJson));
+            delegate.onEnd(
+                    new TransformedReadableSpan(
+                            span,
+                            newInputJson,
+                            newOutputJson,
+                            newContextJson,
+                            spanOriginEnvironment));
         } else {
-            delegate.onEnd(span);
+            delegate.onEnd(
+                    new TransformedReadableSpan(
+                            span, inputJson, outputJson, newContextJson, spanOriginEnvironment));
         }
     }
 
     @Override
     public boolean isEndRequired() {
-        return config.autoConvertAIAttachments() || !samplers.isEmpty() || delegate.isEndRequired();
+        return true;
     }
 
     @Override
@@ -207,11 +261,21 @@ private static class TransformedReadableSpan implements ReadableSpan {
         private final ReadableSpan delegate;
         private final Attributes attributes;
 
-        TransformedReadableSpan(ReadableSpan delegate, String inputJson, String outputJson) {
+        TransformedReadableSpan(
+                ReadableSpan delegate,
+                String inputJson,
+                String outputJson,
+                String contextJson,
+                @Nullable SpanOriginEnvironment environment) {
             this.delegate = delegate;
             var builder = delegate.getAttributes().toBuilder();
-            builder.put(INPUT_JSON, inputJson);
-            builder.put(OUTPUT_JSON, outputJson);
+            if (inputJson != null) {
+                builder.put(INPUT_JSON, inputJson);
+            }
+            if (outputJson != null) {
+                builder.put(OUTPUT_JSON, outputJson);
+            }
+            builder.put(CONTEXT_JSON, contextJson);
             attributes = builder.build();
         }
 
@@ -229,6 +293,9 @@ public <T> T getAttribute(AttributeKey<T> key) {
             if (key.equals(OUTPUT_JSON)) {
                 return (T) attributes.get(OUTPUT_JSON);
             }
+            if (key.equals(CONTEXT_JSON)) {
+                return (T) attributes.get(CONTEXT_JSON);
+            }
             return delegate.getAttribute(key);
         }