From 20e9953a8e2eca2ca8d5570dbd118d1a5dbd9134 Mon Sep 17 00:00:00 2001 From: Musa Date: Fri, 6 Feb 2026 11:53:19 -0800 Subject: [PATCH] refactor: update tracing configuration to use span attributes and adjust related handlers --- config/plano_config_schema.yaml | 17 +++- .../src/handlers/agent_chat_completions.rs | 83 +++++++++---------- crates/brightstaff/src/handlers/llm.rs | 10 +-- crates/brightstaff/src/main.rs | 17 ++-- .../src/tracing/custom_attributes.rs | 46 +++++++++- crates/common/src/configuration.rs | 9 +- .../travel_agents/config.yaml | 5 +- 7 files changed, 122 insertions(+), 65 deletions(-) diff --git a/config/plano_config_schema.yaml b/config/plano_config_schema.yaml index 5f93c5ca..55310891 100644 --- a/config/plano_config_schema.yaml +++ b/config/plano_config_schema.yaml @@ -382,10 +382,19 @@ properties: type: integer trace_arch_internal: type: boolean - span_attribute_header_prefixes: - type: array - items: - type: string + span_attributes: + type: object + properties: + header_prefixes: + type: array + items: + type: string + static: + type: object + additionalProperties: + type: string + additionalProperties: false + additionalProperties: false 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 dc18e093..09dc72a7 100644 --- a/crates/brightstaff/src/handlers/agent_chat_completions.rs +++ b/crates/brightstaff/src/handlers/agent_chat_completions.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use std::time::{Instant, SystemTime}; use bytes::Bytes; -use common::configuration::Tracing; +use common::configuration::SpanAttributes; use common::consts::TRACE_PARENT_HEADER; use common::traces::{generate_random_span_id, parse_traceparent, SpanBuilder, SpanKind}; use hermesllm::apis::OpenAIMessage; @@ -46,7 +46,7 @@ pub async fn agent_chat( agents_list: Arc>>>, listeners: Arc>>, trace_collector: Arc, - tracing_config: Arc>, + span_attributes: Arc>, ) -> Result>, hyper::Error> { match handle_agent_chat( request, @@ -54,7 +54,7 @@ pub async fn agent_chat( agents_list, listeners, trace_collector, - tracing_config, + span_attributes, ) .await { @@ -133,7 +133,7 @@ async fn handle_agent_chat( agents_list: Arc>>>, listeners: Arc>>, trace_collector: Arc, - tracing_config: Arc>, + span_attributes: Arc>, ) -> Result>, AgentFilterChainError> { // Initialize services let agent_selector = AgentSelector::new(orchestrator_service); @@ -185,10 +185,7 @@ async fn handle_agent_chat( }; let custom_attrs = collect_custom_trace_attributes( &request_headers, - tracing_config - .as_ref() - .as_ref() - .and_then(|tracing| tracing.span_attribute_header_prefixes.as_deref()), + span_attributes.as_ref().as_ref(), ); let chat_request_bytes = request.collect().await?.to_bytes(); @@ -264,26 +261,26 @@ async fn handle_agent_chat( let mut selection_span_builder = append_span_attributes( SpanBuilder::new(&selection_operation_name) - .with_span_id(selection_span_id) - .with_kind(SpanKind::Internal) - .with_start_time(selection_start_time) - .with_end_time(selection_end_time) - .with_attribute(http::METHOD, "POST") - .with_attribute(http::TARGET, "/agents/select") - .with_attribute("selection.listener", listener.name.clone()) - .with_attribute("selection.agent_count", selected_agents.len().to_string()) - .with_attribute( - "selection.agents", - selected_agents - .iter() - .map(|a| a.id.as_str()) - .collect::>() - .join(","), - ) - .with_attribute( - "duration_ms", - format!("{:.2}", selection_elapsed.as_secs_f64() * 1000.0), - ), + .with_span_id(selection_span_id) + .with_kind(SpanKind::Internal) + .with_start_time(selection_start_time) + .with_end_time(selection_end_time) + .with_attribute(http::METHOD, "POST") + .with_attribute(http::TARGET, "/agents/select") + .with_attribute("selection.listener", listener.name.clone()) + .with_attribute("selection.agent_count", selected_agents.len().to_string()) + .with_attribute( + "selection.agents", + selected_agents + .iter() + .map(|a| a.id.as_str()) + .collect::>() + .join(","), + ) + .with_attribute( + "duration_ms", + format!("{:.2}", selection_elapsed.as_secs_f64() * 1000.0), + ), &custom_attrs, ); @@ -362,21 +359,21 @@ async fn handle_agent_chat( let mut span_builder = append_span_attributes( SpanBuilder::new(&operation_name) - .with_span_id(span_id) - .with_kind(SpanKind::Internal) - .with_start_time(agent_start_time) - .with_end_time(agent_end_time) - .with_attribute(http::METHOD, "POST") - .with_attribute(http::TARGET, full_path) - .with_attribute("agent.name", agent_name.clone()) - .with_attribute( - "agent.sequence", - format!("{}/{}", agent_index + 1, agent_count), - ) - .with_attribute( - "duration_ms", - format!("{:.2}", agent_elapsed.as_secs_f64() * 1000.0), - ), + .with_span_id(span_id) + .with_kind(SpanKind::Internal) + .with_start_time(agent_start_time) + .with_end_time(agent_end_time) + .with_attribute(http::METHOD, "POST") + .with_attribute(http::TARGET, full_path) + .with_attribute("agent.name", agent_name.clone()) + .with_attribute( + "agent.sequence", + format!("{}/{}", agent_index + 1, agent_count), + ) + .with_attribute( + "duration_ms", + format!("{:.2}", agent_elapsed.as_secs_f64() * 1000.0), + ), &custom_attrs, ); diff --git a/crates/brightstaff/src/handlers/llm.rs b/crates/brightstaff/src/handlers/llm.rs index 4b64733f..6d5b236a 100644 --- a/crates/brightstaff/src/handlers/llm.rs +++ b/crates/brightstaff/src/handlers/llm.rs @@ -1,5 +1,5 @@ use bytes::Bytes; -use common::configuration::{ModelAlias, Tracing}; +use common::configuration::{ModelAlias, SpanAttributes}; use common::consts::{ ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, REQUEST_ID_HEADER, TRACE_PARENT_HEADER, }; @@ -34,7 +34,6 @@ fn full>(chunk: T) -> BoxBody { .boxed() } -// ! we reached the limit of the number of arguments for a function #[allow(clippy::too_many_arguments)] pub async fn llm_chat( request: Request, @@ -43,17 +42,14 @@ pub async fn llm_chat( model_aliases: Arc>>, llm_providers: Arc>, trace_collector: Arc, - tracing_config: Arc>, // ! right here + span_attributes: Arc>, state_storage: Option>, ) -> Result>, hyper::Error> { let request_path = request.uri().path().to_string(); let request_headers = request.headers().clone(); let custom_attrs = collect_custom_trace_attributes( &request_headers, - tracing_config - .as_ref() - .as_ref() - .and_then(|tracing| tracing.span_attribute_header_prefixes.as_deref()), + span_attributes.as_ref().as_ref(), ); let request_id: String = match request_headers .get(REQUEST_ID_HEADER) diff --git a/crates/brightstaff/src/main.rs b/crates/brightstaff/src/main.rs index 5e8af979..9d577c7f 100644 --- a/crates/brightstaff/src/main.rs +++ b/crates/brightstaff/src/main.rs @@ -115,12 +115,17 @@ async fn main() -> Result<(), Box> { )); let model_aliases = Arc::new(plano_config.model_aliases.clone()); - let tracing_config = Arc::new(plano_config.tracing.clone()); + let span_attributes = Arc::new( + plano_config + .tracing + .as_ref() + .and_then(|tracing| tracing.span_attributes.clone()), + ); // Initialize trace collector and start background flusher // Tracing is enabled if the tracing config is present in plano_config.yaml // Pass Some(true/false) to override, or None to use env var OTEL_TRACING_ENABLED - let tracing_enabled = if tracing_config.is_some() { + let tracing_enabled = if plano_config.tracing.is_some() { info!("Tracing configuration found in plano_config.yaml"); Some(true) } else { @@ -180,7 +185,7 @@ async fn main() -> Result<(), Box> { let agents_list = combined_agents_filters_list.clone(); let listeners = listeners.clone(); let trace_collector = trace_collector.clone(); - let tracing_config = tracing_config.clone(); + let span_attributes = span_attributes.clone(); let state_storage = state_storage.clone(); let service = service_fn(move |req| { let router_service = Arc::clone(&router_service); @@ -192,7 +197,7 @@ async fn main() -> Result<(), Box> { let agents_list = agents_list.clone(); let listeners = listeners.clone(); let trace_collector = trace_collector.clone(); - let tracing_config = tracing_config.clone(); + let span_attributes = span_attributes.clone(); let state_storage = state_storage.clone(); async move { @@ -213,7 +218,7 @@ async fn main() -> Result<(), Box> { agents_list, listeners, trace_collector, - tracing_config, + span_attributes, ) .with_context(parent_cx) .await; @@ -232,7 +237,7 @@ async fn main() -> Result<(), Box> { model_aliases, llm_providers, trace_collector, - tracing_config, + span_attributes, state_storage, ) .with_context(parent_cx) diff --git a/crates/brightstaff/src/tracing/custom_attributes.rs b/crates/brightstaff/src/tracing/custom_attributes.rs index 690cec0c..90e60d7a 100644 --- a/crates/brightstaff/src/tracing/custom_attributes.rs +++ b/crates/brightstaff/src/tracing/custom_attributes.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; +use common::configuration::SpanAttributes; use common::traces::SpanBuilder; use hyper::header::HeaderMap; @@ -47,9 +48,24 @@ pub fn extract_custom_trace_attributes( pub fn collect_custom_trace_attributes( headers: &HeaderMap, - span_attribute_header_prefixes: Option<&[String]>, + span_attributes: Option<&SpanAttributes>, ) -> HashMap { - extract_custom_trace_attributes(headers, span_attribute_header_prefixes) + 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()); + } + } + + attributes.extend(extract_custom_trace_attributes( + headers, + span_attributes.header_prefixes.as_deref(), + )); + attributes } pub fn append_span_attributes( @@ -83,4 +99,30 @@ mod tests { 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 = extract_custom_trace_attributes(&headers, None); + assert!(attrs_none.is_empty()); + + let empty_prefixes: Vec = Vec::new(); + let attrs_empty = extract_custom_trace_attributes(&headers, Some(&empty_prefixes)); + 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 prefixes = vec!["x-katanemo-".to_string(), "x-tenant-".to_string()]; + let attrs = extract_custom_trace_attributes(&headers, Some(&prefixes)); + + assert_eq!(attrs.get("tenant.id"), Some(&"ten_456".to_string())); + assert_eq!(attrs.get("user.id"), Some(&"usr_789".to_string())); + } } diff --git a/crates/common/src/configuration.rs b/crates/common/src/configuration.rs index 6ce5daeb..55045450 100644 --- a/crates/common/src/configuration.rs +++ b/crates/common/src/configuration.rs @@ -90,7 +90,14 @@ pub struct Overrides { pub struct Tracing { pub sampling_rate: Option, pub trace_arch_internal: Option, - pub span_attribute_header_prefixes: Option>, + pub span_attributes: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct SpanAttributes { + pub header_prefixes: Option>, + #[serde(rename = "static")] + pub static_attributes: 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 cb2a632c..ce9b8f44 100644 --- a/demos/agent_orchestration/travel_agents/config.yaml +++ b/demos/agent_orchestration/travel_agents/config.yaml @@ -55,5 +55,6 @@ listeners: tracing: random_sampling: 100 - span_attribute_header_prefixes: - - x-katanemo- + span_attributes: + header_prefixes: + - x-katanemo-