fixup: onEnding

Author: realark

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

@@ -4,10 +4,8 @@
 import java.util.List;
 
 /**
- * A filter that decides whether a span should be enriched and exported. Samplers are evaluated in
- * {@link BraintrustSpanProcessor#onEnding} before enrichment and again in {@link
- * BraintrustSpanProcessor#onEnd} before export, so implementations must be deterministic and
- * side-effect-free.
+ * A filter that decides whether a finalized span should be exported. Samplers are evaluated in
+ * {@link BraintrustSpanProcessor#onEnd} after enrichment and all ending hooks have completed.
  */
 interface BraintrustSampler {
     /** Returns {@code true} if the span should be exported, {@code false} to discard it. */

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

@@ -159,12 +159,6 @@ public boolean isStartRequired() {
 
     @Override
     public void onEnding(ReadWriteSpan span) {
-        for (var sampler : samplers) {
-            if (!sampler.sample(span)) {
-                return;
-            }
-        }
-
         try {
             var inputJson = span.getAttribute(INPUT_JSON);
             var newInputJson = attachmentProcessor.processAndUpload(inputJson);

braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java

@@ -232,14 +232,15 @@ public boolean isStartRequired() {
 
                     @Override
                     public void onEnding(ReadWriteSpan span) {
-                        assertEquals("accepted", span.getName());
                         var context =
                                 BraintrustJsonMapper.fromJson(
                                         span.getAttribute(BraintrustSpanProcessor.CONTEXT_JSON),
                                         Map.class);
-                        assertEquals(
-                                "late-attribute",
-                                ((Map<?, ?>) context.get("metadata")).get("source"));
+                        if (span.getName().equals("accepted")) {
+                            assertEquals(
+                                    "late-attribute",
+                                    ((Map<?, ?>) context.get("metadata")).get("source"));
+                        }
                         assertEquals(
                                 "braintrust.sdk.java",
                                 ((Map<?, ?>) context.get("span_origin")).get("name"));
@@ -282,6 +283,13 @@ public boolean isEndRequired() {
             var rejected = tracer.spanBuilder("rejected").startSpan();
             rejected.setAttribute("http.method", "GET");
             rejected.end();
+            assertNotNull(
+                    ((ReadableSpan) rejected).getAttribute(BraintrustSpanProcessor.CONTEXT_JSON),
+                    "Rejected spans are enriched before filtering");
+            assertEquals(
+                    Boolean.TRUE,
+                    ((ReadableSpan) rejected).getAttribute(AttributeKey.booleanKey("smoke.ending")),
+                    "Ending hooks run even for spans rejected at export");
 
             var spans = exporter.getFinishedSpanItems();
             assertEquals(1, spans.size());