add Custom Trace Attributes to extend observability (#708)

* add custom trace attributes

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

* refactor: rename custom_attribute_prefixes to span_attribute_header_prefixes in configuration and related handlers

* docs: add section on custom span attributes

* refactor: update tracing configuration to use span attributes and adjust related handlers

* docs: custom span attributes section to include static attributes and clarify configuration

* add custom trace attributes

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

* refactor: rename custom_attribute_prefixes to span_attribute_header_prefixes in configuration and related handlers

* docs: add section on custom span attributes

* refactor: update tracing configuration to use span attributes and adjust related handlers

* docs: custom span attributes section to include static attributes and clarify configuration

* refactor: remove TraceCollector usage and enhance logging with structured attributes

* refactor: custom trace attribute extraction to improve clarity

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

config/plano_config_schema.yaml

@@ -388,9 +388,19 @@ properties:
         type: integer
       trace_arch_internal:
         type: boolean
-      opentracing_grpc_endpoint:
+      span_attributes:
-        type: string
+        type: object
-      additionalProperties: false
+        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:

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<tokio::sync::RwLock<Option<Vec<common::configuration::Agent>>>>,
     listeners: Arc<tokio::sync::RwLock<Vec<common::configuration::Listener>>>,
+    span_attributes: Arc<Option<SpanAttributes>>,
     llm_providers: Arc<RwLock<LlmProviders>>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, 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<tokio::sync::RwLock<Vec<common::configuration::Listener>>>,
     llm_providers: Arc<RwLock<LlmProviders>>,
     request_id: String,
+    custom_attrs: std::collections::HashMap<String, String>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, 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

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<Option<HashMap<String, ModelAlias>>>,
     llm_providers: Arc<RwLock<LlmProviders>>,
+    span_attributes: Arc<Option<SpanAttributes>>,
     state_storage: Option<Arc<dyn StateStorage>>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, 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<Option<HashMap<String, ModelAlias>>>,
     llm_providers: Arc<RwLock<LlmProviders>>,
+    custom_attrs: HashMap<String, String>,
     state_storage: Option<Arc<dyn StateStorage>>,
     request_id: String,
     request_path: String,
@@ -94,6 +101,11 @@ async fn llm_chat_inner(
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, 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

crates/brightstaff/src/main.rs

@@ -114,6 +114,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     ));
 
     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<dyn std::error::Error + Send + Sync>> {
         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<dyn std::error::Error + Send + Sync>> {
             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<dyn std::error::Error + Send + Sync>> {
                             fully_qualified_url,
                             agents_list,
                             listeners,
+                            span_attributes,
                             llm_providers,
                         )
                         .with_context(parent_cx)
@@ -220,6 +229,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
                             fully_qualified_url,
                             model_aliases,
                             llm_providers,
+                            span_attributes,
                             state_storage,
                         )
                         .with_context(parent_cx)

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<String, String> {
+    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<String, String>,
+) -> 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()));
+    }
+}

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;

crates/common/src/configuration.rs

@@ -92,6 +92,14 @@ pub struct Tracing {
     pub trace_arch_internal: Option<bool>,
     pub random_sampling: Option<u32>,
     pub opentracing_grpc_endpoint: Option<String>,
+    pub span_attributes: Option<SpanAttributes>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct SpanAttributes {
+    pub header_prefixes: Option<Vec<String>>,
+    #[serde(rename = "static")]
+    pub static_attributes: Option<HashMap<String, String>>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]

demos/agent_orchestration/travel_agents/config.yaml

@@ -55,3 +55,6 @@ listeners:
 
 tracing:
   random_sampling: 100
+  span_attributes:
+    header_prefixes:
+      - x-acme-

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
 }

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
 -----------------------------------------