diff --git a/config/plano_config_schema.yaml b/config/plano_config_schema.yaml index cd736eb6..7c35d2cf 100644 --- a/config/plano_config_schema.yaml +++ b/config/plano_config_schema.yaml @@ -388,9 +388,19 @@ properties: type: integer trace_arch_internal: type: boolean - opentracing_grpc_endpoint: - type: string - additionalProperties: false + 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: @@ -403,7 +413,7 @@ properties: type: string model: type: string - additionalProperties: false + additionalProperties: false state_storage: type: object properties: diff --git a/crates/brightstaff/src/handlers/agent_chat_completions.rs b/crates/brightstaff/src/handlers/agent_chat_completions.rs index dea736e3..513e0ef2 100644 --- a/crates/brightstaff/src/handlers/agent_chat_completions.rs +++ b/crates/brightstaff/src/handlers/agent_chat_completions.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Instant; use bytes::Bytes; +use common::configuration::SpanAttributes; use common::errors::BrightStaffError; use common::llm_providers::LlmProviders; use hermesllm::apis::OpenAIMessage; @@ -20,7 +21,7 @@ use super::agent_selector::{AgentSelectionError, AgentSelector}; use super::pipeline_processor::{PipelineError, PipelineProcessor}; use super::response_handler::ResponseHandler; use crate::router::plano_orchestrator::OrchestratorService; -use crate::tracing::{operation_component, set_service_name}; +use crate::tracing::{collect_custom_trace_attributes, operation_component, set_service_name}; /// Main errors for agent chat completions #[derive(Debug, thiserror::Error)] @@ -43,8 +44,11 @@ pub async fn agent_chat( _: String, agents_list: Arc>>>, listeners: Arc>>, + span_attributes: Arc>, llm_providers: Arc>, ) -> Result>, hyper::Error> { + let custom_attrs = + collect_custom_trace_attributes(request.headers(), span_attributes.as_ref().as_ref()); // Extract request_id from headers or generate a new one let request_id: String = match request .headers() @@ -77,6 +81,7 @@ pub async fn agent_chat( listeners, llm_providers, request_id, + custom_attrs, ) .await { @@ -164,6 +169,7 @@ async fn handle_agent_chat_inner( listeners: Arc>>, llm_providers: Arc>, request_id: String, + custom_attrs: std::collections::HashMap, ) -> Result>, AgentFilterChainError> { // Initialize services let agent_selector = AgentSelector::new(orchestrator_service); @@ -186,6 +192,9 @@ async fn handle_agent_chat_inner( get_active_span(|span| { span.update_name(listener.name.to_string()); + for (key, value) in &custom_attrs { + span.set_attribute(opentelemetry::KeyValue::new(key.clone(), value.clone())); + } }); info!(listener = %listener.name, "handling request"); @@ -348,6 +357,9 @@ async fn handle_agent_chat_inner( set_service_name(operation_component::AGENT); get_active_span(|span| { span.update_name(format!("{} /v1/chat/completions", agent_name)); + for (key, value) in &custom_attrs { + span.set_attribute(opentelemetry::KeyValue::new(key.clone(), value.clone())); + } }); pipeline_processor diff --git a/crates/brightstaff/src/handlers/llm.rs b/crates/brightstaff/src/handlers/llm.rs index 8e8f9661..ee41dd2d 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; +use common::configuration::{ModelAlias, SpanAttributes}; use common::consts::{ ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, REQUEST_ID_HEADER, TRACE_PARENT_HEADER, }; @@ -28,7 +28,9 @@ use crate::state::response_state_processor::ResponsesStateProcessor; use crate::state::{ extract_input_items, retrieve_and_combine_input, StateStorage, StateStorageError, }; -use crate::tracing::{llm as tracing_llm, operation_component, set_service_name}; +use crate::tracing::{ + collect_custom_trace_attributes, llm as tracing_llm, operation_component, set_service_name, +}; use common::errors::BrightStaffError; @@ -38,6 +40,7 @@ pub async fn llm_chat( full_qualified_llm_provider_url: String, model_aliases: Arc>>, llm_providers: Arc>, + span_attributes: Arc>, state_storage: Option>, ) -> Result>, hyper::Error> { let request_path = request.uri().path().to_string(); @@ -50,6 +53,8 @@ pub async fn llm_chat( Some(id) => id, None => uuid::Uuid::new_v4().to_string(), }; + let custom_attrs = + collect_custom_trace_attributes(&request_headers, span_attributes.as_ref().as_ref()); // Create a span with request_id that will be included in all log lines let request_span = info_span!( @@ -71,6 +76,7 @@ pub async fn llm_chat( full_qualified_llm_provider_url, model_aliases, llm_providers, + custom_attrs, state_storage, request_id, request_path, @@ -87,6 +93,7 @@ async fn llm_chat_inner( full_qualified_llm_provider_url: String, model_aliases: Arc>>, llm_providers: Arc>, + custom_attrs: HashMap, state_storage: Option>, request_id: String, request_path: String, @@ -94,6 +101,11 @@ async fn llm_chat_inner( ) -> Result>, hyper::Error> { // Set service name for LLM operations set_service_name(operation_component::LLM); + get_active_span(|span| { + for (key, value) in &custom_attrs { + span.set_attribute(opentelemetry::KeyValue::new(key.clone(), value.clone())); + } + }); // Extract or generate traceparent - this establishes the trace context for all spans let traceparent: String = match request_headers diff --git a/crates/brightstaff/src/main.rs b/crates/brightstaff/src/main.rs index 87deda6a..97345556 100644 --- a/crates/brightstaff/src/main.rs +++ b/crates/brightstaff/src/main.rs @@ -114,6 +114,12 @@ async fn main() -> Result<(), Box> { )); let model_aliases = Arc::new(plano_config.model_aliases.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 @@ -173,6 +179,7 @@ async fn main() -> Result<(), Box> { let llm_providers = llm_providers.clone(); let agents_list = combined_agents_filters_list.clone(); let listeners = listeners.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); @@ -183,6 +190,7 @@ async fn main() -> Result<(), Box> { let model_aliases = Arc::clone(&model_aliases); let agents_list = agents_list.clone(); let listeners = listeners.clone(); + let span_attributes = span_attributes.clone(); let state_storage = state_storage.clone(); async move { @@ -202,6 +210,7 @@ async fn main() -> Result<(), Box> { fully_qualified_url, agents_list, listeners, + span_attributes, llm_providers, ) .with_context(parent_cx) @@ -220,6 +229,7 @@ async fn main() -> Result<(), Box> { fully_qualified_url, model_aliases, llm_providers, + 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 new file mode 100644 index 00000000..24abc72b --- /dev/null +++ b/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 { + 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, +) -> 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())); + } +} diff --git a/crates/brightstaff/src/tracing/mod.rs b/crates/brightstaff/src/tracing/mod.rs index 7332170c..1fa8a7e2 100644 --- a/crates/brightstaff/src/tracing/mod.rs +++ b/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; diff --git a/crates/common/src/configuration.rs b/crates/common/src/configuration.rs index 0a683b8b..f4e2b7b4 100644 --- a/crates/common/src/configuration.rs +++ b/crates/common/src/configuration.rs @@ -92,6 +92,14 @@ pub struct Tracing { pub trace_arch_internal: Option, pub random_sampling: Option, pub opentracing_grpc_endpoint: 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 2cb24d71..911baf89 100644 --- a/demos/agent_orchestration/travel_agents/config.yaml +++ b/demos/agent_orchestration/travel_agents/config.yaml @@ -55,3 +55,6 @@ listeners: tracing: random_sampling: 100 + span_attributes: + header_prefixes: + - x-acme- diff --git a/demos/agent_orchestration/travel_agents/test.rest b/demos/agent_orchestration/travel_agents/test.rest index f3ecaf66..b6348f28 100644 --- a/demos/agent_orchestration/travel_agents/test.rest +++ b/demos/agent_orchestration/travel_agents/test.rest @@ -3,9 +3,16 @@ ### Travel Agent Chat Completion Request POST {{llm_endpoint}}/v1/chat/completions HTTP/1.1 Content-Type: application/json +X-Acme-Workspace-Id: ws_7e2c5d91b4224f59b0e6a4e0125c21b3 +X-Acme-Tenant-Id: ten_4102a8c7fa6542b084b395d2df184a9a +X-Acme-User-Id: usr_19df7e6751b846f9ba026776e3c12abe +X-Acme-Admin-Level: 3 +X-Acme-Environment: production +X-Acme-Is-Internal: false +X-Acme-Cost-Center: HD100 { - "model": "gpt-4o", + "model": "gpt-5.2", "messages": [ { "role": "user", @@ -20,7 +27,28 @@ Content-Type: application/json "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 } diff --git a/docs/source/guides/observability/tracing.rst b/docs/source/guides/observability/tracing.rst index 9e07483a..950befd2 100644 --- a/docs/source/guides/observability/tracing.rst +++ b/docs/source/guides/observability/tracing.rst @@ -142,6 +142,109 @@ In your observability platform (Jaeger, Grafana Tempo, Datadog, etc.), filter tr For complete details on all available signals, detection methods, and best practices, see the :doc:`../../concepts/signals` guide. +Custom Span Attributes +------------------------------------------- + +Plano can automatically attach **custom span attributes** derived from request headers and **static** attributes +defined in configuration. This lets you stamp +traces with identifiers like workspace, tenant, or user IDs without changing application code or adding +custom instrumentation. + +**Why This Is Useful** + +- **Tenant-aware debugging**: Filter traces by ``workspace.id`` or ``tenant.id``. +- **Customer-specific visibility**: Attribute performance or errors to a specific customer. +- **Low overhead**: No code changes in agents or clients—just headers. + +How It Works +~~~~~~~~~~~~ + +You configure one or more header prefixes. Any incoming HTTP header whose name starts with one of these +prefixes is captured as a span attribute. You can also provide static attributes that are always injected. + +- The **prefix is only for matching**, not the resulting attribute key. +- The attribute key is the header name **with the prefix removed**, then hyphens converted to dots. + +.. note:: + + Custom span attributes are attached to LLM spans when handling ``/v1/...`` requests via ``llm_chat``. For orchestrator requests to ``/agents/...``, + these attributes are added to both the orchestrator selection span and to each agent span created by ``agent_chat``. + +**Example** + +Configured prefix:: + + tracing: + span_attributes: + header_prefixes: + - x-katanemo- + +Incoming headers:: + + X-Katanemo-Workspace-Id: ws_123 + X-Katanemo-Tenant-Id: ten_456 + +Resulting span attributes:: + + workspace.id = "ws_123" + tenant.id = "ten_456" + +Configuration +~~~~~~~~~~~~~ + +Add the prefix list under ``tracing`` in your config: + +.. code-block:: yaml + + tracing: + random_sampling: 100 + span_attributes: + header_prefixes: + - x-katanemo- + static: + environment: production + service.version: "1.0.0" + +Static attributes are always injected alongside any header-derived attributes. If a header-derived +attribute key matches a static key, the header value overrides the static value. + +You can provide multiple prefixes: + +.. code-block:: yaml + + tracing: + span_attributes: + header_prefixes: + - x-katanemo- + - x-tenant- + static: + environment: production + service.version: "1.0.0" + +Notes and Examples +~~~~~~~~~~~~~~~~~~ + +- **Prefix must match exactly**: ``katanemo-`` does not match ``x-katanemo-`` headers. +- **Trailing dash is recommended**: Without it, ``x-katanemo`` would also match ``x-katanemo-foo`` and + ``x-katanemofoo``. +- **Keys are always strings**: Values are captured as string attributes. + +**Prefix mismatch example** + +Config:: + + tracing: + span_attributes: + header_prefixes: + - x-katanemo- + +Request headers:: + + X-Other-User-Id: usr_999 + +Result: no attributes are captured from ``X-Other-User-Id``. + + Benefits of Using ``Traceparent`` Headers -----------------------------------------