feat: add span customizers (SDK-316)

src/Braintrust.Sdk/Braintrust.Sdk.csproj

@@ -55,6 +55,8 @@
   </Target>
 
   <ItemGroup>
+    <PackageReference Include="Google.Protobuf" Version="3.29.3" />
+    <PackageReference Include="Grpc.Tools" Version="2.67.0" PrivateAssets="all" />
     <PackageReference Include="OpenTelemetry" Version="1.15.3" />
     <PackageReference Include="OpenTelemetry.Api" Version="1.15.3" />
     <PackageReference Include="OpenTelemetry.Exporter.OpenTelemetryProtocol" Version="1.15.3" />
@@ -65,4 +67,12 @@
     <PackageReference Include="System.Text.Json" Version="9.0.19" />
   </ItemGroup>
 
+  <!-- Official OTLP message schemas, pinned to opentelemetry-proto v1.5.0:
+       https://github.com/open-telemetry/opentelemetry-proto/tree/v1.5.0
+       Messages are internal implementation details; the existing OTel exporter owns transport. -->
+  <ItemGroup>
+    <Protobuf Include="Trace/Protos/opentelemetry/proto/**/*.proto"
+              ProtoRoot="Trace/Protos" GrpcServices="None" Access="Internal" />
+  </ItemGroup>
+
 </Project>

src/Braintrust.Sdk/Config/BraintrustConfig.cs

@@ -1,3 +1,5 @@
+using Braintrust.Sdk.Trace;
+
 namespace Braintrust.Sdk.Config;
 
 public sealed record SpanOriginEnvironment(string Type, string? Name = null);
@@ -29,12 +31,33 @@ public sealed class BraintrustConfig : BaseConfig
     public TimeSpan RequestTimeout { get; }
     public SpanOriginEnvironment? Environment { get; }
 
+    /// <summary>
+    /// Ordered registration snapshot. Hooks run synchronously on each detached outgoing span.
+    /// </summary>
+    public IReadOnlyList<ISpanCustomizer> SpanCustomizers { get; }
+
     public static BraintrustConfig FromEnvironment()
     {
         return Of();
     }
 
+    public static BraintrustConfig FromEnvironment(IEnumerable<ISpanCustomizer> spanCustomizers)
+    {
+        return Of(spanCustomizers);
+    }
+
     public static BraintrustConfig Of(params (string Key, string? Value)[] envOverrides)
+    {
+        return Of(Array.Empty<ISpanCustomizer>(), envOverrides);
+    }
+
+    /// <summary>
+    /// Creates a configuration with an immutable copy of the ordered customizer registrations.
+    /// Customizer instances themselves are not cloned and must be safe for export threads.
+    /// </summary>
+    public static BraintrustConfig Of(
+        IEnumerable<ISpanCustomizer> spanCustomizers,
+        params (string Key, string? Value)[] envOverrides)
     {
         var overridesMap = new Dictionary<string, string?>();
 
@@ -43,11 +66,21 @@ public static BraintrustConfig Of(params (string Key, string? Value)[] envOverri
             overridesMap[key] = value;
         }
 
-        return new BraintrustConfig(overridesMap);
+        return new BraintrustConfig(overridesMap, spanCustomizers);
     }
 
-    private BraintrustConfig(IDictionary<string, string?> envOverrides) : base(envOverrides)
+    private BraintrustConfig(
+        IDictionary<string, string?> envOverrides,
+        IEnumerable<ISpanCustomizer> spanCustomizers) : base(envOverrides)
     {
+        ArgumentNullException.ThrowIfNull(spanCustomizers);
+        var customizers = spanCustomizers.ToArray();
+        if (customizers.Any(customizer => customizer is null))
+        {
+            throw new ArgumentException("Span customizers cannot contain null.", nameof(spanCustomizers));
+        }
+        SpanCustomizers = Array.AsReadOnly(customizers);
+
         try
         {
             _braintrustEnvSearchRoot = Directory.GetCurrentDirectory();

src/Braintrust.Sdk/Trace/BraintrustSpanCustomizerHandler.cs

@@ -0,0 +1,159 @@
+using Braintrust.Sdk.Trace.Protos.Collector.Trace.V1;
+using Braintrust.Sdk.Trace.Protos.Common.V1;
+using Braintrust.Sdk.Trace.Protos.Resource.V1;
+using Braintrust.Sdk.Trace.Protos.Trace.V1;
+using OtlpSpan = Braintrust.Sdk.Trace.Protos.Trace.V1.Span;
+using Google.Protobuf;
+
+namespace Braintrust.Sdk.Trace;
+
+/// <summary>
+/// Adapts the upstream Activity-only exporter to detached per-span OTLP JSON data.
+/// Hooks run before final serialization and before auth or network transmission.
+/// </summary>
+internal sealed class BraintrustSpanCustomizerHandler : DelegatingHandler
+{
+    private readonly IReadOnlyList<ISpanCustomizer> _customizers;
+
+    internal BraintrustSpanCustomizerHandler(IReadOnlyList<ISpanCustomizer> customizers)
+    {
+        _customizers = customizers;
+    }
+
+    protected override HttpResponseMessage Send(HttpRequestMessage request, CancellationToken cancellationToken)
+    {
+        CustomizeAsync(request, cancellationToken).GetAwaiter().GetResult();
+        return base.Send(request, cancellationToken);
+    }
+
+    protected override async Task<HttpResponseMessage> SendAsync(
+        HttpRequestMessage request, CancellationToken cancellationToken)
+    {
+        await CustomizeAsync(request, cancellationToken).ConfigureAwait(false);
+        return await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
+    }
+
+    private async Task CustomizeAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+    {
+        try
+        {
+            var original = request.Content ?? throw new InvalidOperationException("Missing OTLP request content.");
+            if (original.Headers.ContentType?.MediaType != "application/x-protobuf" ||
+                original.Headers.ContentEncoding.Count != 0)
+            {
+                throw new InvalidOperationException("Span customization requires uncompressed OTLP protobuf content.");
+            }
+
+            var bytes = await original.ReadAsByteArrayAsync(cancellationToken).ConfigureAwait(false);
+            var batch = ExportTraceServiceRequest.Parser.ParseFrom(bytes);
+            batch = CustomizeBatch(batch);
+
+            var replacement = new ByteArrayContent(batch.ToByteArray());
+            try
+            {
+                foreach (var header in original.Headers)
+                {
+                    if (!header.Key.Equals("Content-Length", StringComparison.OrdinalIgnoreCase))
+                    {
+                        replacement.Headers.TryAddWithoutValidation(header.Key, header.Value);
+                    }
+                }
+                request.Content = replacement;
+            }
+            catch
+            {
+                replacement.Dispose();
+                throw;
+            }
+            original.Dispose();
+        }
+        catch (Exception ex)
+        {
+            // Never log the snapshot or exception message, which may contain secrets.
+            Console.Error.WriteLine($"[Braintrust] Span customization failed; no spans sent ({ex.GetType().Name}).");
+            // The OTel exporter catches this and reports ExportResult.Failure. Do not
+            // use HttpRequestException: customization failures are not network retries.
+            throw new InvalidOperationException("Braintrust span customization failed; the batch was not sent.", ex);
+        }
+    }
+
+    private ExportTraceServiceRequest CustomizeBatch(ExportTraceServiceRequest batch)
+    {
+        OtlpJsonPayload.ValidateKnownFields(batch);
+        var result = new ExportTraceServiceRequest();
+        var groups = new Dictionary<(Resource? Resource, string SchemaUrl),
+            (ResourceSpans Resource, Dictionary<(InstrumentationScope? Scope, string SchemaUrl), ScopeSpans> Scopes)>();
+
+        foreach (var resource in batch.ResourceSpans)
+        {
+            if (resource.ScopeSpans.Count == 0)
+            {
+                Add(resource);
+            }
+            foreach (var scope in resource.ScopeSpans)
+            {
+                if (scope.Spans.Count == 0)
+                {
+                    Add(new ResourceSpans
+                    {
+                        Resource = resource.Resource,
+                        SchemaUrl = resource.SchemaUrl,
+                        ScopeSpans = { scope }
+                    });
+                }
+                foreach (var span in scope.Spans)
+                {
+                    var data = OtlpJsonPayload.ToSpanJson(resource, scope, span);
+                    for (var i = 0; i < _customizers.Count; i++)
+                    {
+                        data = _customizers[i].OnSpanExport(data)
+                            ?? throw new InvalidOperationException("A span customizer returned null.");
+                        var customized = OtlpJsonPayload.FromSpanJson(data);
+                        ValidateIdentity(customized.ScopeSpans[0].Spans[0], span);
+                        if (i == _customizers.Count - 1)
+                        {
+                            // Snapshot the final result before invoking hooks on another span.
+                            // Retained JsonObjects cannot mutate a previously customized span.
+                            Add(customized);
+                        }
+                    }
+                }
+            }
+        }
+        return result;
+
+        void Add(ResourceSpans data)
+        {
+            // Group only after customization. Metadata objects are private protobuf
+            // snapshots, never mutated after becoming dictionary keys.
+            var key = (data.Resource, data.SchemaUrl);
+            if (!groups.TryGetValue(key, out var group))
+            {
+                group = (new ResourceSpans { Resource = data.Resource, SchemaUrl = data.SchemaUrl }, new());
+                groups.Add(key, group);
+                result.ResourceSpans.Add(group.Resource);
+            }
+            foreach (var incomingScope in data.ScopeSpans)
+            {
+                var scopeKey = (incomingScope.Scope, incomingScope.SchemaUrl);
+                if (!group.Scopes.TryGetValue(scopeKey, out var scope))
+                {
+                    scope = new ScopeSpans { Scope = incomingScope.Scope, SchemaUrl = incomingScope.SchemaUrl };
+                    group.Scopes.Add(scopeKey, scope);
+                    group.Resource.ScopeSpans.Add(scope);
+                }
+                scope.Spans.Add(incomingScope.Spans);
+            }
+        }
+    }
+
+    private static void ValidateIdentity(OtlpSpan span, OtlpSpan original)
+    {
+        if (!span.TraceId.Equals(original.TraceId) ||
+            !span.SpanId.Equals(original.SpanId) ||
+            !span.ParentSpanId.Equals(original.ParentSpanId))
+        {
+            throw new InvalidOperationException("A span customizer changed the trace, span or parent span ID.");
+        }
+    }
+}

src/Braintrust.Sdk/Trace/BraintrustTracing.cs

@@ -83,17 +83,27 @@ public static void Enable(BraintrustConfig config, TracerProviderBuilder tracerP
                 otlpOptions.Protocol = OtlpExportProtocol.HttpProtobuf;
                 otlpOptions.Endpoint = new Uri($"{config.ApiUrl}{config.TracesPath}");
                 otlpOptions.TimeoutMilliseconds = (int)config.RequestTimeout.TotalMilliseconds;
-                otlpOptions.HttpClientFactory = () => new HttpClient(new BraintrustOtlpAuthHandler(config)
-                {
-                    InnerHandler = new HttpClientHandler()
-                })
-                {
-                    Timeout = config.RequestTimeout
-                };
+                otlpOptions.HttpClientFactory = () => CreateHttpClient(config);
             })
             .SetSampler(new AlwaysOnSampler());
     }
 
+    internal static HttpClient CreateHttpClient(BraintrustConfig config, HttpMessageHandler? transport = null)
+    {
+        HttpMessageHandler handler = new BraintrustOtlpAuthHandler(config)
+        {
+            InnerHandler = transport ?? new HttpClientHandler()
+        };
+        if (config.SpanCustomizers.Count != 0)
+        {
+            handler = new BraintrustSpanCustomizerHandler(config.SpanCustomizers)
+            {
+                InnerHandler = handler
+            };
+        }
+        return new HttpClient(handler) { Timeout = config.RequestTimeout };
+    }
+
     /// <summary>
     /// Flush all pending spans to Braintrust.
     /// Returns true if the flush completed within the timeout, false otherwise.

src/Braintrust.Sdk/Trace/ISpanCustomizer.cs

@@ -0,0 +1,21 @@
+using System.Text.Json.Nodes;
+
+namespace Braintrust.Sdk.Trace;
+
+/// <summary>
+/// Transforms detached outgoing span data, never the application's Activity.
+/// Hooks run synchronously in registration order for each completed span.
+/// </summary>
+public interface ISpanCustomizer
+{
+    /// <summary>
+    /// Returns the supplied span or a replacement, for example from <see cref="JsonNode.DeepClone"/>.
+    /// The object contains OTLP span fields, plus optional resource, scope, resourceSchemaUrl
+    /// and scopeSchemaUrl metadata. Metadata changes apply only to this span.
+    /// traceId, spanId and parentSpanId must remain unchanged. IDs are hexadecimal strings,
+    /// enums are integers and 64-bit integers are decimal strings. Null, exceptions, invalid
+    /// data or changed IDs fail the entire batch before transmission. Do not retain and
+    /// mutate the object after returning.
+    /// </summary>
+    JsonObject OnSpanExport(JsonObject span) => span;
+}