diff --git a/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java b/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java index 86b9bb67..5a34ff31 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/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 = + 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 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; + } } diff --git a/braintrust-sdk/src/main/java/dev/braintrust/config/SpanOriginEnvironment.java b/braintrust-sdk/src/main/java/dev/braintrust/config/SpanOriginEnvironment.java new file mode 100644 index 00000000..77a5901b --- /dev/null +++ b/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) {} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java index aa0efb04..450bb407 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java +++ b/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) { diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java index 813a12c8..c243bb36 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java +++ b/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 PARENT = AttributeKey.stringKey(BraintrustTracing.PARENT_KEY); + static final AttributeKey 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 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 spanOrigin = + existingOrigin instanceof Map map + ? new HashMap<>((Map) 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(); + 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 getAttribute(AttributeKey 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); } diff --git a/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java b/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java index 0ebe306f..d395ef5c 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java @@ -91,4 +91,58 @@ void testDefaultSslContextWhenNotProvided() throws Exception { // Should be the system defaults assertEquals(SSLContext.getDefault(), config.sslContext()); } + + @Test + void explicitNullSpanOriginEnvironmentDisablesDetection() { + var config = + BraintrustConfig.of( + "BRAINTRUST_ENVIRONMENT_TYPE", + BaseConfig.NULL_OVERRIDE, + "BRAINTRUST_ENVIRONMENT_NAME", + BaseConfig.NULL_OVERRIDE, + "CI", + "true", + "AWS_EXECUTION_ENV", + "AWS_ECS_FARGATE"); + + assertTrue(config.spanOriginEnvironment().isEmpty()); + } + + @Test + void explicitEnvironmentNameWithoutTypeIsPreserved() { + var config = + BraintrustConfig.of( + "BRAINTRUST_ENVIRONMENT_TYPE", + BaseConfig.NULL_OVERRIDE, + "BRAINTRUST_ENVIRONMENT_NAME", + "staging", + "CI", + "true"); + + var environment = config.spanOriginEnvironment().orElseThrow(); + assertNull(environment.type()); + assertEquals("staging", environment.name()); + } + + @Test + void awsExecutionEnvClassifiesEcsBeforeLambda() { + var config = + BraintrustConfig.of( + "GITHUB_ACTIONS", "true", "AWS_EXECUTION_ENV", "AWS_ECS_FARGATE"); + + var environment = config.spanOriginEnvironment().orElseThrow(); + assertEquals("server", environment.type()); + assertEquals("ecs", environment.name()); + } + + @Test + void awsExecutionEnvClassifiesLambdaWhenLambdaSpecific() { + var config = + BraintrustConfig.of( + "GITHUB_ACTIONS", "true", "AWS_EXECUTION_ENV", "AWS_Lambda_java17"); + + var environment = config.spanOriginEnvironment().orElseThrow(); + assertEquals("server", environment.type()); + assertEquals("aws_lambda", environment.name()); + } } diff --git a/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java b/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java index 6b3bda06..a3437458 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java @@ -5,10 +5,12 @@ import dev.braintrust.TestHarness; import dev.braintrust.UnitTestSpanExporter; import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.json.BraintrustJsonMapper; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.util.Map; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -112,6 +114,41 @@ void filterAISpansDropsNonAISpans() { assertEquals("llm-call", spans.get(0).getName()); } + @Test + void spanOriginMergesWithContextJsonSetAfterSpanStart() throws Exception { + var config = + BraintrustConfig.builder() + .apiKey("test-key") + .apiUrl("http://localhost:1234") + .build(); + + var spanExporter = new UnitTestSpanExporter(); + var processor = + new BraintrustSpanProcessor(config, SimpleSpanProcessor.create(spanExporter)); + var tracerProvider = SdkTracerProvider.builder().addSpanProcessor(processor).build(); + var tracer = tracerProvider.get("test"); + + var span = tracer.spanBuilder("late-context").startSpan(); + span.setAttribute( + "braintrust.context_json", "{\"metadata\":{\"source\":\"late-attribute\"}}"); + span.end(); + + tracerProvider.forceFlush().join(5, TimeUnit.SECONDS); + var spans = spanExporter.getFinishedSpanItems(); + assertEquals(1, spans.size()); + + var contextJson = + spans.get(0).getAttributes().get(AttributeKey.stringKey("braintrust.context_json")); + assertNotNull(contextJson); + var context = BraintrustJsonMapper.get().readValue(contextJson, Map.class); + assertEquals("late-attribute", ((Map) context.get("metadata")).get("source")); + var spanOrigin = (Map) context.get("span_origin"); + assertEquals("braintrust.sdk.java", spanOrigin.get("name")); + assertEquals( + BraintrustTracing.INSTRUMENTATION_NAME, + ((Map) spanOrigin.get("instrumentation")).get("name")); + } + private void doSimpleOtelTrace(Tracer tracer) { // use tracer to create a simple trace with a root span and a child span var span = tracer.spanBuilder("unit-test-root").startSpan();