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

config/arch_config_schema.yaml

@@ -382,28 +382,10 @@ properties:
         type: integer
       trace_arch_internal:
         type: boolean
-    custom_attributes:
-      type: array
-      items:
-        type: object
-        properties:
-          key:
-            type: string
-          type:
-            type: string
-            enum:
-              - str
-              - bool
-              - float
-              - int
-          header:
-            type: string
-        additionalProperties: false
-        required:
-          - key
-          - type
-          - header
-      additionalProperties: false
+      custom_attribute_prefixes:
+        type: array
+        items:
+          type: string
   mode:
     type: string
     enum:

crates/brightstaff/src/handlers/agent_chat_completions.rs

@@ -187,7 +187,7 @@ async fn handle_agent_chat(
         tracing_config
             .as_ref()
             .as_ref()
-            .and_then(|tracing| tracing.custom_attributes.as_deref()),
+            .and_then(|tracing| tracing.custom_attribute_prefixes.as_deref()),
     );
 
     let chat_request_bytes = request.collect().await?.to_bytes();

crates/brightstaff/src/handlers/llm.rs

@@ -53,7 +53,7 @@ pub async fn llm_chat(
         tracing_config
             .as_ref()
             .as_ref()
-            .and_then(|tracing| tracing.custom_attributes.as_deref()),
+            .and_then(|tracing| tracing.custom_attribute_prefixes.as_deref()),
     );
     let request_id: String = match request_headers
         .get(REQUEST_ID_HEADER)
@@ -434,7 +434,7 @@ async fn build_llm_span(
     tool_names: Option<Vec<String>>,
     user_message_preview: Option<String>,
     temperature: Option<f32>,
-    llm_providers: &Arc<RwLock<Vec<LlmProvider>>>,
+    llm_providers: &Arc<RwLock<LlmProviders>>,
     custom_attrs: &HashMap<String, String>,
 ) -> common::traces::Span {
     use crate::tracing::{http, llm, OperationNameBuilder};

crates/brightstaff/src/tracing/custom_attributes.rs

@@ -1,45 +1,44 @@
 use std::collections::HashMap;
 
-use common::configuration::{CustomTraceAttribute, CustomTraceAttributeType};
-use hyper::header::{HeaderMap, HeaderName};
+use hyper::header::HeaderMap;
 
 pub fn extract_custom_trace_attributes(
     headers: &HeaderMap,
-    custom_attributes: Option<&[CustomTraceAttribute]>,
+    custom_attribute_prefixes: Option<&[String]>,
 ) -> HashMap<String, String> {
     let mut attributes = HashMap::new();
-    let Some(custom_attributes) = custom_attributes else {
+    let Some(custom_attribute_prefixes) = custom_attribute_prefixes else {
         return attributes;
     };
+    if custom_attribute_prefixes.is_empty() {
+        return attributes;
+    }
 
-    for attribute in custom_attributes {
-        // Normalize/validate the configured header name; skip invalid names.
-        let header_name = match HeaderName::from_bytes(attribute.header.as_bytes()) {
-            Ok(name) => name,
-            Err(_) => continue,
+    for (name, value) in headers.iter() {
+        let header_name = name.as_str();
+        let mut matched_prefix: Option<&str> = None;
+        for prefix in custom_attribute_prefixes {
+            if header_name.starts_with(prefix) {
+                matched_prefix = Some(prefix.as_str());
+                break;
+            }
+        }
+        let Some(prefix) = matched_prefix else {
+            continue;
         };
 
-        // Extract header value as UTF-8 text; skip missing or invalid values.
-        let raw_value = match headers
-            .get(header_name)
-            .and_then(|value| value.to_str().ok())
-        {
+        let raw_value = match value.to_str().ok() {
             Some(value) => value.trim(),
             None => continue,
         };
 
-        // Parse the header value according to the configured type.
-        let parsed_value = match attribute.value_type {
-            CustomTraceAttributeType::Str => Some(raw_value.to_string()),
-            CustomTraceAttributeType::Bool => raw_value.parse::<bool>().ok().map(|v| v.to_string()),
-            CustomTraceAttributeType::Float => raw_value.parse::<f64>().ok().map(|v| v.to_string()),
-            CustomTraceAttributeType::Int => raw_value.parse::<i64>().ok().map(|v| v.to_string()),
-        };
-
-        // Only include attributes that successfully parsed.
-        if let Some(value) = parsed_value {
-            attributes.insert(attribute.key.clone(), value);
+        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
@@ -48,72 +47,22 @@ pub fn extract_custom_trace_attributes(
 #[cfg(test)]
 mod tests {
     use super::extract_custom_trace_attributes;
-    use common::configuration::{CustomTraceAttribute, CustomTraceAttributeType};
     use hyper::header::{HeaderMap, HeaderValue};
 
     #[test]
-    fn extracts_and_parses_custom_headers() {
+    fn extracts_headers_by_prefix() {
         let mut headers = HeaderMap::new();
-        headers.insert("x-workspace-id", HeaderValue::from_static("ws_123"));
-        headers.insert("x-tenant-id", HeaderValue::from_static("ten_456"));
-        headers.insert("x-user-id", HeaderValue::from_static("usr_789"));
-        headers.insert("x-admin-level", HeaderValue::from_static("3"));
-        headers.insert("x-is-internal", HeaderValue::from_static("true"));
-        headers.insert("x-budget", HeaderValue::from_static("42.5"));
-        headers.insert("x-bad-int", HeaderValue::from_static("nope"));
+        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 custom_attributes = vec![
-            CustomTraceAttribute {
-                key: "workspace.id".to_string(),
-                value_type: CustomTraceAttributeType::Str,
-                header: "x-workspace-id".to_string(),
-            },
-            CustomTraceAttribute {
-                key: "tenant.id".to_string(),
-                value_type: CustomTraceAttributeType::Str,
-                header: "x-tenant-id".to_string(),
-            },
-            CustomTraceAttribute {
-                key: "user.id".to_string(),
-                value_type: CustomTraceAttributeType::Str,
-                header: "x-user-id".to_string(),
-            },
-            CustomTraceAttribute {
-                key: "admin.level".to_string(),
-                value_type: CustomTraceAttributeType::Int,
-                header: "x-admin-level".to_string(),
-            },
-            CustomTraceAttribute {
-                key: "is.internal".to_string(),
-                value_type: CustomTraceAttributeType::Bool,
-                header: "x-is-internal".to_string(),
-            },
-            CustomTraceAttribute {
-                key: "budget.value".to_string(),
-                value_type: CustomTraceAttributeType::Float,
-                header: "x-budget".to_string(),
-            },
-            CustomTraceAttribute {
-                key: "bad.int".to_string(),
-                value_type: CustomTraceAttributeType::Int,
-                header: "x-bad-int".to_string(),
-            },
-            CustomTraceAttribute {
-                key: "missing.header".to_string(),
-                value_type: CustomTraceAttributeType::Str,
-                header: "x-missing".to_string(),
-            },
-        ];
+        let prefixes = vec!["x-katanemo-".to_string()];
+        let attrs = extract_custom_trace_attributes(&headers, Some(&prefixes));
 
-        let attrs = extract_custom_trace_attributes(&headers, Some(&custom_attributes));
-
-        assert_eq!(attrs.get("workspace.id"), Some(&"ws_123".to_string()));
         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_eq!(attrs.get("is.internal"), Some(&"true".to_string()));
-        assert_eq!(attrs.get("budget.value"), Some(&"42.5".to_string()));
-        assert!(!attrs.contains_key("bad.int"));
-        assert!(!attrs.contains_key("missing.header"));
+        assert!(!attrs.contains_key("other.id"));
     }
 }

crates/common/src/configuration.rs

@@ -90,24 +90,7 @@ pub struct Overrides {
 pub struct Tracing {
     pub sampling_rate: Option<f64>,
     pub trace_arch_internal: Option<bool>,
-    pub custom_attributes: Option<Vec<CustomTraceAttribute>>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct CustomTraceAttribute {
-    pub key: String,
-    #[serde(rename = "type")]
-    pub value_type: CustomTraceAttributeType,
-    pub header: String,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
-#[serde(rename_all = "lowercase")]
-pub enum CustomTraceAttributeType {
-    Str,
-    Bool,
-    Float,
-    Int,
+    pub custom_attribute_prefixes: Option<Vec<String>>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]

demos/use_cases/travel_agents/config.yaml

@@ -55,22 +55,5 @@ listeners:
 
 tracing:
   random_sampling: 100
-  custom_attributes:
-    - header: x-workspace-id
-      key: workspace.id
-      type: str
-    - header: x-tenant-id
-      key: tenant.id
-      type: str
-    - header: x-user-id
-      key: user.id
-      type: str
-    - header: x-admin-level
-      key: admin.level
-      type: int
-    - header: x-is-internal
-      key: is.internal
-      type: bool
-    - header: x-budget
-      key: budget.value
-      type: float
+  custom_attribute_prefixes:
+    - x-katanemo-

demos/use_cases/travel_agents/test.rest

@@ -3,15 +3,15 @@
 ### Travel Agent Chat Completion Request
 POST {{llm_endpoint}}/v1/chat/completions HTTP/1.1
 Content-Type: application/json
-X-Workspace-Id: ws_7e2c5d91b4224f59b0e6a4e0125c21b3
-X-Tenant-Id: ten_4102a8c7fa6542b084b395d2df184a9a
-X-User-Id: usr_19df7e6751b846f9ba026776e3c12abe
-X-Admin-Level: 3
-X-Is-Internal: true
-X-Budget: 42.5
+X-Katanemo-Workspace-Id: ws_7e2c5d91b4224f59b0e6a4e0125c21b3
+X-Katanemo-Tenant-Id: ten_4102a8c7fa6542b084b395d2df184a9a
+X-Katanemo-User-Id: usr_19df7e6751b846f9ba026776e3c12abe
+X-Katanemo-Admin-Level: 3
+X-Katanemo-Is-Internal: true
+X-Katanemo-Budget: 42.5
 
 {
-  "model": "gpt-4o",
+  "model": "gpt-5.2",
   "messages": [
     {
       "role": "user",
@@ -26,7 +26,28 @@ X-Budget: 42.5
       "content": "What is one Alaska flight that goes direct to Atlanta from Seattle?"
     }
   ],
-  "max_tokens": 1000,
+  "max_completion_tokens": 1000,
+  "stream": false,
+  "temperature": 1.0
+}
+
+
+### Travel Agent Request (prefix mismatch - ignored)
+POST {{llm_endpoint}}/v1/chat/completions HTTP/1.1
+Content-Type: application/json
+X-Other-Workspace-Id: ws_7e2c5d91b4224f59b0e6a4e0125c21b3
+X-Other-Tenant-Id: ten_4102a8c7fa6542b084b395d2df184a9a
+X-Other-User-Id: usr_19df7e6751b846f9ba026776e3c12abe
+
+{
+  "model": "gpt-5.2",
+  "messages": [
+    {
+      "role": "user",
+      "content": "What's the weather in Seattle?"
+    }
+  ],
+  "max_completion_tokens": 1000,
   "stream": false,
   "temperature": 1.0
 }