opentelemetry_langfuse/
context.rs

1//! Langfuse context helpers for setting trace attributes.
2//!
3//! This module provides a `LangfuseContext` struct that can be used to store
4//! trace-level attributes. This is typically used with interceptors or middleware
5//! where each instance maintains its own context.
6//!
7//! # Example
8//!
9//! ```no_run
10//! use opentelemetry_langfuse::LangfuseContext;
11//!
12//! // Create a context instance
13//! let context = LangfuseContext::new();
14//!
15//! // Set session ID for grouping traces
16//! context.set_session_id("session-123");
17//!
18//! // Set user ID for attribution
19//! context.set_user_id("user-456");
20//!
21//! // Add tags for filtering
22//! context.add_tags(vec!["production".to_string(), "api-v2".to_string()]);
23//!
24//! // Get attributes to add to spans
25//! let attributes = context.get_attributes();
26//! ```
27
28use opentelemetry::KeyValue;
29use std::collections::HashMap;
30use std::sync::{Arc, RwLock};
31
32/// Langfuse-specific attribute keys.
33///
34/// These follow the Langfuse OpenTelemetry conventions documented at:
35/// <https://langfuse.com/integrations/native/opentelemetry>
36pub mod attributes {
37    /// Session ID attribute key (Langfuse expects `langfuse.session.id` or `session.id`)
38    pub const TRACE_SESSION_ID: &str = "langfuse.session.id";
39    /// User ID attribute key (Langfuse expects `langfuse.user.id` or `user.id`)
40    pub const TRACE_USER_ID: &str = "langfuse.user.id";
41    /// Tags attribute key - must be a string array (Langfuse expects `langfuse.trace.tags`)
42    pub const TRACE_TAGS: &str = "langfuse.trace.tags";
43    /// Metadata attribute key (JSON object string)
44    pub const TRACE_METADATA: &str = "langfuse.trace.metadata";
45    /// Trace name attribute key
46    pub const TRACE_NAME: &str = "langfuse.trace.name";
47}
48
49/// Thread-safe storage for Langfuse context attributes.
50///
51/// This context allows you to set attributes that will be automatically
52/// included in all spans created within the same context.
53#[derive(Clone)]
54pub struct LangfuseContext {
55    attributes: Arc<RwLock<HashMap<String, String>>>,
56}
57
58impl LangfuseContext {
59    /// Create a new empty context.
60    #[must_use]
61    pub fn new() -> Self {
62        Self {
63            attributes: Arc::new(RwLock::new(HashMap::new())),
64        }
65    }
66
67    /// Set the session ID for the current trace.
68    pub fn set_session_id(&self, session_id: impl Into<String>) -> &Self {
69        self.set_attribute(attributes::TRACE_SESSION_ID, session_id);
70        self
71    }
72
73    /// Set the user ID for the current trace.
74    pub fn set_user_id(&self, user_id: impl Into<String>) -> &Self {
75        self.set_attribute(attributes::TRACE_USER_ID, user_id);
76        self
77    }
78
79    /// Add tags to the current trace.
80    ///
81    /// Tags are stored as a JSON array string.
82    pub fn add_tags(&self, tags: Vec<String>) -> &Self {
83        let tags_json = serde_json::to_string(&tags).unwrap_or_else(|_| "[]".to_string());
84        self.set_attribute(attributes::TRACE_TAGS, tags_json);
85        self
86    }
87
88    /// Add a single tag.
89    pub fn add_tag(&self, tag: impl Into<String>) -> &Self {
90        let tag = tag.into();
91        let mut attrs = self.attributes.write().unwrap();
92
93        // Append to existing tags if present
94        if let Some(existing) = attrs.get(attributes::TRACE_TAGS) {
95            // Parse existing JSON array, add tag, and re-serialize
96            if let Ok(mut tags_vec) = serde_json::from_str::<Vec<String>>(existing) {
97                tags_vec.push(tag);
98                let tags_json =
99                    serde_json::to_string(&tags_vec).unwrap_or_else(|_| "[]".to_string());
100                attrs.insert(attributes::TRACE_TAGS.to_string(), tags_json);
101            } else {
102                // Fallback if existing isn't valid JSON
103                attrs.insert(attributes::TRACE_TAGS.to_string(), format!("[\"{}\"]", tag));
104            }
105        } else {
106            attrs.insert(attributes::TRACE_TAGS.to_string(), format!("[\"{}\"]", tag));
107        }
108        drop(attrs);
109        self
110    }
111
112    /// Set metadata as JSON string.
113    pub fn set_metadata(&self, metadata: serde_json::Value) -> &Self {
114        let metadata_str = metadata.to_string();
115        self.set_attribute(attributes::TRACE_METADATA, metadata_str);
116        self
117    }
118
119    /// Set a custom attribute.
120    pub fn set_attribute(&self, key: impl Into<String>, value: impl Into<String>) -> &Self {
121        let mut attrs = self.attributes.write().unwrap();
122        attrs.insert(key.into(), value.into());
123        self
124    }
125
126    /// Set the trace name.
127    pub fn set_trace_name(&self, name: impl Into<String>) -> &Self {
128        self.set_attribute(attributes::TRACE_NAME, name);
129        self
130    }
131
132    /// Clear all attributes.
133    pub fn clear(&self) {
134        let mut attrs = self.attributes.write().unwrap();
135        attrs.clear();
136    }
137
138    /// Get all current attributes as key-value pairs.
139    #[must_use]
140    pub fn get_attributes(&self) -> Vec<KeyValue> {
141        let attrs = self.attributes.read().unwrap();
142        attrs
143            .iter()
144            .map(|(k, v)| KeyValue::new(k.clone(), v.clone()))
145            .collect()
146    }
147
148    /// Check if a specific attribute is set.
149    #[must_use]
150    pub fn has_attribute(&self, key: &str) -> bool {
151        let attrs = self.attributes.read().unwrap();
152        attrs.contains_key(key)
153    }
154
155    /// Get a specific attribute value.
156    #[must_use]
157    pub fn get_attribute(&self, key: &str) -> Option<String> {
158        let attrs = self.attributes.read().unwrap();
159        attrs.get(key).cloned()
160    }
161}
162
163impl Default for LangfuseContext {
164    fn default() -> Self {
165        Self::new()
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn test_context_attributes() {
175        let ctx = LangfuseContext::new();
176        ctx.set_session_id("session-123");
177        ctx.set_user_id("user-456");
178
179        assert_eq!(
180            ctx.get_attribute(attributes::TRACE_SESSION_ID),
181            Some("session-123".to_string())
182        );
183        assert_eq!(
184            ctx.get_attribute(attributes::TRACE_USER_ID),
185            Some("user-456".to_string())
186        );
187    }
188
189    #[test]
190    fn test_tags() {
191        let ctx = LangfuseContext::new();
192        ctx.add_tags(vec!["tag1".to_string(), "tag2".to_string()]);
193
194        let tags_json = ctx.get_attribute(attributes::TRACE_TAGS).unwrap();
195        let tags: Vec<String> = serde_json::from_str(&tags_json).unwrap();
196        assert_eq!(tags, vec!["tag1", "tag2"]);
197    }
198
199    #[test]
200    fn test_add_single_tag() {
201        let ctx = LangfuseContext::new();
202        ctx.add_tag("tag1");
203        ctx.add_tag("tag2");
204
205        let tags_json = ctx.get_attribute(attributes::TRACE_TAGS).unwrap();
206        let tags: Vec<String> = serde_json::from_str(&tags_json).unwrap();
207        assert_eq!(tags, vec!["tag1", "tag2"]);
208    }
209
210    #[test]
211    fn test_fluent_api() {
212        let ctx = LangfuseContext::new();
213        ctx.set_session_id("session-123")
214            .set_user_id("user-456")
215            .add_tags(vec!["tag1".to_string()])
216            .set_trace_name("my-trace");
217
218        assert!(ctx.has_attribute(attributes::TRACE_SESSION_ID));
219        assert!(ctx.has_attribute(attributes::TRACE_USER_ID));
220        assert!(ctx.has_attribute(attributes::TRACE_TAGS));
221        assert!(ctx.has_attribute(attributes::TRACE_NAME));
222    }
223
224    #[test]
225    fn test_clear() {
226        let ctx = LangfuseContext::new();
227        ctx.set_session_id("session-123");
228        assert!(ctx.has_attribute(attributes::TRACE_SESSION_ID));
229
230        ctx.clear();
231        assert!(!ctx.has_attribute(attributes::TRACE_SESSION_ID));
232    }
233}
opentelemetry_langfuse/context.rs

opentelemetry_langfuse/
context.rs
