apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-06-17 15:25:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

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

Browse source
This commit is contained in:
Musa 2026-02-05 17:03:20 -08:00
parent 3a663aa780
commit aa64704ed1
No known key found for this signature in database
7 changed files with 73 additions and 155 deletions
Download patch file Download diff file Expand all files Collapse all files

2
crates/brightstaff/src/handlers/agent_chat_completions.rs
View file

@ -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();

4
crates/brightstaff/src/handlers/llm.rs
View file

@ -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};

119
crates/brightstaff/src/tracing/custom_attributes.rs
View file

@ -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<String, String> {
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::<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,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<String, String> {
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"));
}
}

19
crates/common/src/configuration.rs
View file

@ -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)]