add Custom Trace Attributes to extend observability (#708)

* add custom trace attributes

* refactor: prefix custom trace attributes and update schema handlers tests configs

* refactor: rename custom_attribute_prefixes to span_attribute_header_prefixes in configuration and related handlers

* docs: add section on custom span attributes

* refactor: update tracing configuration to use span attributes and adjust related handlers

* docs: custom span attributes section to include static attributes and clarify configuration

* add custom trace attributes

* refactor: prefix custom trace attributes and update schema handlers tests configs

* refactor: rename custom_attribute_prefixes to span_attribute_header_prefixes in configuration and related handlers

* docs: add section on custom span attributes

* refactor: update tracing configuration to use span attributes and adjust related handlers

* docs: custom span attributes section to include static attributes and clarify configuration

* refactor: remove TraceCollector usage and enhance logging with structured attributes

* refactor: custom trace attribute extraction to improve clarity

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

crates/brightstaff/src/tracing/custom_attributes.rs

@@ -0,0 +1,156 @@
+use std::collections::HashMap;
+
+use common::configuration::SpanAttributes;
+use common::traces::SpanBuilder;
+use hyper::header::HeaderMap;
+
+pub fn collect_custom_trace_attributes(
+    headers: &HeaderMap,
+    span_attributes: Option<&SpanAttributes>,
+) -> HashMap<String, String> {
+    let mut attributes = HashMap::new();
+    let Some(span_attributes) = span_attributes else {
+        return attributes;
+    };
+
+    if let Some(static_attributes) = span_attributes.static_attributes.as_ref() {
+        for (key, value) in static_attributes {
+            attributes.insert(key.clone(), value.clone());
+        }
+    }
+
+    let Some(header_prefixes) = span_attributes.header_prefixes.as_deref() else {
+        return attributes;
+    };
+    if header_prefixes.is_empty() {
+        return attributes;
+    }
+
+    for (name, value) in headers.iter() {
+        let header_name = name.as_str();
+        let matched_prefix = header_prefixes
+            .iter()
+            .find(|prefix| header_name.starts_with(prefix.as_str()))
+            .map(String::as_str);
+        let Some(prefix) = matched_prefix else {
+            continue;
+        };
+
+        let Some(raw_value) = value.to_str().ok().map(str::trim) else {
+            continue;
+        };
+
+        let suffix = header_name.strip_prefix(prefix).unwrap_or("");
+        let suffix_key = suffix.trim_start_matches('-').replace('-', ".");
+        if suffix_key.is_empty() {
+            continue;
+        }
+
+        attributes.insert(suffix_key, raw_value.to_string());
+    }
+
+    attributes
+}
+
+pub fn append_span_attributes(
+    mut span_builder: SpanBuilder,
+    attributes: &HashMap<String, String>,
+) -> SpanBuilder {
+    for (key, value) in attributes {
+        span_builder = span_builder.with_attribute(key, value);
+    }
+    span_builder
+}
+
+#[cfg(test)]
+mod tests {
+    use super::collect_custom_trace_attributes;
+    use common::configuration::SpanAttributes;
+    use hyper::header::{HeaderMap, HeaderValue};
+    use std::collections::HashMap;
+
+    #[test]
+    fn extracts_headers_by_prefix() {
+        let mut headers = HeaderMap::new();
+        headers.insert("x-katanemo-tenant-id", HeaderValue::from_static("ten_456"));
+        headers.insert("x-katanemo-user-id", HeaderValue::from_static("usr_789"));
+        headers.insert("x-katanemo-admin-level", HeaderValue::from_static("3"));
+        headers.insert("x-other-id", HeaderValue::from_static("ignored"));
+
+        let attrs = collect_custom_trace_attributes(
+            &headers,
+            Some(&SpanAttributes {
+                header_prefixes: Some(vec!["x-katanemo-".to_string()]),
+                static_attributes: None,
+            }),
+        );
+
+        assert_eq!(attrs.get("tenant.id"), Some(&"ten_456".to_string()));
+        assert_eq!(attrs.get("user.id"), Some(&"usr_789".to_string()));
+        assert_eq!(attrs.get("admin.level"), Some(&"3".to_string()));
+        assert!(!attrs.contains_key("other.id"));
+    }
+
+    #[test]
+    fn returns_empty_when_prefixes_missing_or_empty() {
+        let mut headers = HeaderMap::new();
+        headers.insert("x-katanemo-tenant-id", HeaderValue::from_static("ten_456"));
+
+        let attrs_none = collect_custom_trace_attributes(
+            &headers,
+            Some(&SpanAttributes {
+                header_prefixes: None,
+                static_attributes: None,
+            }),
+        );
+        assert!(attrs_none.is_empty());
+
+        let attrs_empty = collect_custom_trace_attributes(
+            &headers,
+            Some(&SpanAttributes {
+                header_prefixes: Some(Vec::new()),
+                static_attributes: None,
+            }),
+        );
+        assert!(attrs_empty.is_empty());
+    }
+
+    #[test]
+    fn supports_multiple_prefixes() {
+        let mut headers = HeaderMap::new();
+        headers.insert("x-katanemo-tenant-id", HeaderValue::from_static("ten_456"));
+        headers.insert("x-tenant-user-id", HeaderValue::from_static("usr_789"));
+
+        let attrs = collect_custom_trace_attributes(
+            &headers,
+            Some(&SpanAttributes {
+                header_prefixes: Some(vec!["x-katanemo-".to_string(), "x-tenant-".to_string()]),
+                static_attributes: None,
+            }),
+        );
+
+        assert_eq!(attrs.get("tenant.id"), Some(&"ten_456".to_string()));
+        assert_eq!(attrs.get("user.id"), Some(&"usr_789".to_string()));
+    }
+
+    #[test]
+    fn header_attributes_override_static_attributes() {
+        let mut headers = HeaderMap::new();
+        headers.insert("x-katanemo-tenant-id", HeaderValue::from_static("ten_456"));
+
+        let mut static_attributes = HashMap::new();
+        static_attributes.insert("tenant.id".to_string(), "ten_static".to_string());
+        static_attributes.insert("environment".to_string(), "prod".to_string());
+
+        let attrs = collect_custom_trace_attributes(
+            &headers,
+            Some(&SpanAttributes {
+                header_prefixes: Some(vec!["x-katanemo-".to_string()]),
+                static_attributes: Some(static_attributes),
+            }),
+        );
+
+        assert_eq!(attrs.get("tenant.id"), Some(&"ten_456".to_string()));
+        assert_eq!(attrs.get("environment"), Some(&"prod".to_string()));
+    }
+}

crates/brightstaff/src/tracing/mod.rs

@@ -1,9 +1,11 @@
 mod constants;
+mod custom_attributes;
 mod service_name_exporter;
 
 pub use constants::{
     error, http, llm, operation_component, routing, signals, OperationNameBuilder,
 };
+pub use custom_attributes::{append_span_attributes, collect_custom_trace_attributes};
 pub use service_name_exporter::{ServiceNameOverrideExporter, SERVICE_NAME_OVERRIDE_KEY};
 
 use opentelemetry::trace::get_active_span;