Propagate inherited fields properly (#54)

ankrgyl · web-flow · commit a66e78312e09 · 2026-05-15T11:38:22.000-07:00
diff --git a/src/dataset.rs b/src/dataset.rs
@@ -741,6 +741,7 @@ impl<S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer +
                 } else {
                     None
                 },
+                extra: std::collections::HashMap::new(),
             };
 
             let parent_info = ParentSpanInfo::Dataset {
diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs
@@ -198,6 +198,7 @@ impl<
             tags: feedback.tags,
             context: None,
             span_attributes: None,
+            extra: HashMap::new(),
         };
 
         let parent_info = ParentSpanInfo::Experiment {
@@ -384,6 +385,7 @@ impl<
                 purpose: None,
                 extra: HashMap::new(),
             }),
+            extra: HashMap::new(),
         };
 
         let parent_info = ParentSpanInfo::Experiment {
diff --git a/src/log_queue/merge.rs b/src/log_queue/merge.rs
@@ -185,6 +185,15 @@ pub(crate) fn merge_row_into(target: &mut Logs3Row, source: Logs3Row) {
         }
     }
 
+    for (key, source_value) in source.extra {
+        match target.extra.get_mut(&key) {
+            Some(target_value) => deep_merge(target_value, &source_value),
+            None => {
+                target.extra.insert(key, source_value);
+            }
+        }
+    }
+
     // Combine merge_paths from both entries.
     // Only retain merge paths if both entries are merge-type entries.
     let final_is_merge = target.is_merge.unwrap_or(false) && source.is_merge.unwrap_or(false);
diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs
@@ -2,7 +2,6 @@ use anyhow::Context;
 use chrono::Utc;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-use std::collections::HashMap;
 
 use super::batching::batch_and_serialize_rows;
 use super::config::LogQueueConfig;
@@ -379,6 +378,7 @@ impl LogQueueCore {
             tags,
             context,
             span_attributes,
+            extra,
         } = payload;
 
         if let Some(span_components) = span_components {
@@ -415,7 +415,7 @@ impl LogQueueCore {
                 tags,
                 context,
                 span_attributes,
-                extra: HashMap::new(),
+                extra,
                 created: Utc::now(),
                 xact_id: None,
                 object_delete: None,
@@ -539,7 +539,7 @@ impl LogQueueCore {
             tags,
             context,
             span_attributes,
-            extra: HashMap::new(),
+            extra,
             created: Utc::now(),
             xact_id: None,
             object_delete: None,
@@ -821,7 +821,7 @@ impl LogQueueCore {
                 tags: payload.tags,
                 context: payload.context,
                 span_attributes: payload.span_attributes,
-                extra: HashMap::new(),
+                extra: payload.extra,
                 created: Utc::now(),
                 xact_id: None,
                 object_delete: None,
@@ -1005,6 +1005,7 @@ mod tests {
                 tags: None,
                 context: None,
                 span_attributes: None,
+                extra: std::collections::HashMap::new(),
             },
             parent_info: Some(ParentSpanInfo::Experiment {
                 object_id: "exp-test".to_string(),
diff --git a/src/logger.rs b/src/logger.rs
@@ -23,7 +23,7 @@ use crate::http::build_http_client;
 use crate::log_queue::{LogQueue, LogQueueConfig};
 use crate::span::SpanSubmitter;
 use crate::span_components::SpanComponents;
-use crate::types::{ParentSpanInfo, SpanAttributes, SpanObjectType, SpanPayload};
+use crate::types::{ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType, SpanPayload};
 
 // 30s covers both quick API calls (login, project registration) and slower batch log uploads.
 // The TypeScript SDK applies no explicit timeout.
@@ -632,15 +632,7 @@ impl BraintrustClient {
         }
 
         let span_id = components.span_id.clone().unwrap_or_else(|| row_id.clone());
-
-        let payload = SpanPayload {
-            row_id,
-            span_id,
-            is_merge: true,
-            span_components: Some(components),
-            org_id,
-            org_name,
-            project_name: None,
+        let mut event_data = SpanEventData {
             input: event.input,
             output: event.output,
             expected: event.expected,
@@ -656,6 +648,31 @@ impl BraintrustClient {
                 purpose: None,
                 extra: HashMap::new(),
             }),
+            extra: HashMap::new(),
+        };
+        if let Some(propagated_event) = components.propagated_event.as_ref() {
+            event_data.apply_propagated_event(propagated_event);
+        }
+
+        let payload = SpanPayload {
+            row_id,
+            span_id,
+            is_merge: true,
+            span_components: Some(components),
+            org_id,
+            org_name,
+            project_name: None,
+            input: event_data.input,
+            output: event_data.output,
+            expected: event_data.expected,
+            error: event_data.error,
+            scores: event_data.scores,
+            metadata: event_data.metadata,
+            metrics: event_data.metrics,
+            tags: event_data.tags,
+            context: event_data.context,
+            span_attributes: event_data.span_attributes,
+            extra: event_data.extra,
         };
 
         self.submit_payload(token, payload, None);
diff --git a/src/span.rs b/src/span.rs
@@ -9,7 +9,9 @@ use uuid::Uuid;
 
 use crate::error::Result;
 use crate::span_components::SpanComponents;
-use crate::types::{ParentSpanInfo, SpanAttributes, SpanObjectType, SpanPayload, SpanType};
+use crate::types::{
+    ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType, SpanPayload, SpanType,
+};
 
 /// Error type for SpanLog builder validation.
 #[derive(Debug, Clone, PartialEq)]
@@ -578,14 +580,7 @@ struct SpanData {
 }
 
 impl From<SpanData> for SpanPayload {
-    fn from(mut data: SpanData) -> Self {
-        // Apply propagated_event if present - merge its fields into the span data
-        if let Some(propagated_event) = data.propagated_event.take() {
-            apply_propagated_event_to_span_data(&mut data, &propagated_event);
-            // Restore the propagated_event after applying it
-            data.propagated_event = Some(propagated_event);
-        }
-
+    fn from(data: SpanData) -> Self {
         let span_attributes = SpanAttributes {
             name: data.name,
             span_type: Some(data.span_type),
@@ -598,14 +593,7 @@ impl From<SpanData> for SpanPayload {
             || span_attributes.span_type.is_some()
             || span_attributes.purpose.is_some();
 
-        Self {
-            row_id: data.row_id,
-            span_id: data.span_id,
-            is_merge: data.has_flushed, // First flush = false (replace), subsequent = true (merge)
-            span_components: None,
-            org_id: data.org_id,
-            org_name: data.org_name,
-            project_name: data.project_name,
+        let mut event_data = SpanEventData {
             input: data.input,
             output: data.output,
             expected: data.expected,
@@ -616,76 +604,32 @@ impl From<SpanData> for SpanPayload {
             tags: (!data.tags.is_empty()).then_some(data.tags),
             context: data.context,
             span_attributes: has_attributes.then_some(span_attributes),
+            extra: HashMap::new(),
+        };
+
+        if let Some(propagated_event) = data.propagated_event {
+            event_data.apply_propagated_event(&propagated_event);
         }
-    }
-}
 
-/// Apply propagated_event data to SpanData by merging it into the span's fields.
-/// This allows parent span metadata to flow down to child spans.
-fn apply_propagated_event_to_span_data(data: &mut SpanData, propagated_event: &Map<String, Value>) {
-    for (key, value) in propagated_event {
-        match key.as_str() {
-            "input" if data.input.is_none() => {
-                data.input = Some(value.clone());
-            }
-            "output" if data.output.is_none() => {
-                data.output = Some(value.clone());
-            }
-            "expected" if data.expected.is_none() => {
-                data.expected = Some(value.clone());
-            }
-            "error" if data.error.is_none() => {
-                if let Some(s) = value.as_str() {
-                    data.error = Some(Value::String(s.to_string()));
-                }
-            }
-            "scores" => {
-                if let Some(obj) = value.as_object() {
-                    for (score_key, score_val) in obj {
-                        if let Some(score) = score_val.as_f64() {
-                            data.scores.entry(score_key.clone()).or_insert(score);
-                        }
-                    }
-                }
-            }
-            "metadata" => {
-                if let Some(obj) = value.as_object() {
-                    for (meta_key, meta_val) in obj {
-                        data.metadata
-                            .entry(meta_key.clone())
-                            .or_insert_with(|| meta_val.clone());
-                    }
-                }
-            }
-            "metrics" => {
-                if let Some(obj) = value.as_object() {
-                    for (metric_key, metric_val) in obj {
-                        if let Some(metric) = metric_val.as_f64() {
-                            data.metrics.entry(metric_key.clone()).or_insert(metric);
-                        }
-                    }
-                }
-            }
-            "tags" => {
-                if let Some(arr) = value.as_array() {
-                    for tag_val in arr {
-                        if let Some(tag) = tag_val.as_str() {
-                            if !data.tags.contains(&tag.to_string()) {
-                                data.tags.push(tag.to_string());
-                            }
-                        }
-                    }
-                }
-            }
-            "context" if data.context.is_none() => {
-                data.context = Some(value.clone());
-            }
-            _ => {
-                // Unknown fields go into metadata
-                data.metadata
-                    .entry(key.clone())
-                    .or_insert_with(|| value.clone());
-            }
+        Self {
+            row_id: data.row_id,
+            span_id: data.span_id,
+            is_merge: data.has_flushed, // First flush = false (replace), subsequent = true (merge)
+            span_components: None,
+            org_id: data.org_id,
+            org_name: data.org_name,
+            project_name: data.project_name,
+            input: event_data.input,
+            output: event_data.output,
+            expected: event_data.expected,
+            error: event_data.error,
+            scores: event_data.scores,
+            metadata: event_data.metadata,
+            metrics: event_data.metrics,
+            tags: event_data.tags,
+            context: event_data.context,
+            span_attributes: event_data.span_attributes,
+            extra: event_data.extra,
         }
     }
 }
@@ -874,6 +818,10 @@ mod tests {
         let mut parent_propagated = Map::new();
         parent_propagated.insert("parent_key".to_string(), json!("parent_value"));
         parent_propagated.insert("metadata".to_string(), json!({"inherited": "data"}));
+        parent_propagated.insert(
+            "span_attributes".to_string(),
+            json!({"purpose": "scorer", "skip_realtime": true}),
+        );
 
         let parent_info = ParentSpanInfo::FullSpan {
             object_type: SpanObjectType::ProjectLogs,
@@ -910,9 +858,19 @@ mod tests {
             "Inherited metadata should be present"
         );
         assert_eq!(
-            metadata.get("parent_key").and_then(|v| v.as_str()),
+            captured
+                .payload
+                .extra
+                .get("parent_key")
+                .and_then(|v| v.as_str()),
             Some("parent_value"),
-            "Parent key should be in metadata"
+            "Unknown propagated fields should remain top-level extras"
+        );
+        let span_attributes = captured.payload.span_attributes.as_ref().unwrap();
+        assert_eq!(span_attributes.purpose.as_deref(), Some("scorer"));
+        assert_eq!(
+            span_attributes.extra.get("skip_realtime"),
+            Some(&json!(true))
         );
     }
 
diff --git a/src/types.rs b/src/types.rs
diff --git a/tests/span_lifecycle.rs b/tests/span_lifecycle.rs
