From aa64704ed18e1f2ecbbce8a19cbcc6ad497cfc32 Mon Sep 17 00:00:00 2001 From: Musa Date: Thu, 5 Feb 2026 17:03:20 -0800 Subject: [PATCH] refactor: prefix custom trace attributes and update schema handlers tests configs --- config/plano_config_schema.yaml | 26 +--- .../src/handlers/agent_chat_completions.rs | 2 +- crates/brightstaff/src/handlers/llm.rs | 4 +- .../src/tracing/custom_attributes.rs | 119 +++++------------- crates/common/src/configuration.rs | 19 +-- .../travel_agents/config.yaml | 21 +--- .../travel_agents/test.rest | 37 ++++-- 7 files changed, 73 insertions(+), 155 deletions(-) diff --git a/config/plano_config_schema.yaml b/config/plano_config_schema.yaml index 8fd98e2c..52820f24 100644 --- a/config/plano_config_schema.yaml +++ b/config/plano_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: diff --git a/crates/brightstaff/src/handlers/agent_chat_completions.rs b/crates/brightstaff/src/handlers/agent_chat_completions.rs index 32d9e91a..268933f0 100644 --- a/crates/brightstaff/src/handlers/agent_chat_completions.rs +++ b/crates/brightstaff/src/handlers/agent_chat_completions.rs @@ -188,7 +188,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(); diff --git a/crates/brightstaff/src/handlers/llm.rs b/crates/brightstaff/src/handlers/llm.rs index e3abcf50..4ac75869 100644 --- a/crates/brightstaff/src/handlers/llm.rs +++ b/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>, user_message_preview: Option, temperature: Option, - llm_providers: &Arc>>, + llm_providers: &Arc>, custom_attrs: &HashMap, ) -> common::traces::Span { use crate::tracing::{http, llm, OperationNameBuilder}; diff --git a/crates/brightstaff/src/tracing/custom_attributes.rs b/crates/brightstaff/src/tracing/custom_attributes.rs index 9f6d7c01..690cec0c 100644 --- a/crates/brightstaff/src/tracing/custom_attributes.rs +++ b/crates/brightstaff/src/tracing/custom_attributes.rs @@ -1,46 +1,45 @@ use std::collections::HashMap; -use common::configuration::{CustomTraceAttribute, CustomTraceAttributeType}; use common::traces::SpanBuilder; -use hyper::header::{HeaderMap, HeaderName}; +use hyper::header::HeaderMap; pub fn extract_custom_trace_attributes( headers: &HeaderMap, - custom_attributes: Option<&[CustomTraceAttribute]>, + span_attribute_header_prefixes: Option<&[String]>, ) -> HashMap { let mut attributes = HashMap::new(); - let Some(custom_attributes) = custom_attributes else { + let Some(span_attribute_header_prefixes) = span_attribute_header_prefixes else { return attributes; }; + if span_attribute_header_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 span_attribute_header_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::().ok().map(|v| v.to_string()), - CustomTraceAttributeType::Float => raw_value.parse::().ok().map(|v| v.to_string()), - CustomTraceAttributeType::Int => raw_value.parse::().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,9 +47,9 @@ pub fn extract_custom_trace_attributes( pub fn collect_custom_trace_attributes( headers: &HeaderMap, - custom_attributes: Option<&[CustomTraceAttribute]>, + span_attribute_header_prefixes: Option<&[String]>, ) -> HashMap { - extract_custom_trace_attributes(headers, custom_attributes) + extract_custom_trace_attributes(headers, span_attribute_header_prefixes) } pub fn append_span_attributes( @@ -66,72 +65,22 @@ pub fn append_span_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")); } } diff --git a/crates/common/src/configuration.rs b/crates/common/src/configuration.rs index b6888ac6..47fbb373 100644 --- a/crates/common/src/configuration.rs +++ b/crates/common/src/configuration.rs @@ -90,24 +90,7 @@ pub struct Overrides { pub struct Tracing { pub sampling_rate: Option, pub trace_arch_internal: Option, - pub custom_attributes: Option>, -} - -#[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>, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)] diff --git a/demos/agent_orchestration/travel_agents/config.yaml b/demos/agent_orchestration/travel_agents/config.yaml index 4bfc6ea5..2924cde8 100644 --- a/demos/agent_orchestration/travel_agents/config.yaml +++ b/demos/agent_orchestration/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- diff --git a/demos/agent_orchestration/travel_agents/test.rest b/demos/agent_orchestration/travel_agents/test.rest index 0d188104..7d7c5759 100644 --- a/demos/agent_orchestration/travel_agents/test.rest +++ b/demos/agent_orchestration/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 }