use crate::tracing::observation_layer::{BatchManager, ObservationLayer, SpanTracker}; use chrono::Utc; use reqwest::{Client, StatusCode}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use std::env; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; use url::Url; use uuid::Uuid; const DEFAULT_LANGFUSE_URL: &str = "http://localhost:3000"; #[derive(Debug, Serialize, Deserialize)] struct LangfuseIngestionResponse { successes: Vec, errors: Vec, } #[derive(Debug, Serialize, Deserialize)] struct LangfuseIngestionSuccess { id: String, status: i32, } #[derive(Debug, Serialize, Deserialize)] struct LangfuseIngestionError { id: String, status: i32, message: Option, error: Option, } #[derive(Debug, Clone)] pub struct LangfuseBatchManager { pub batch: Vec, pub client: Client, pub base_url: String, pub public_key: String, pub secret_key: String, } impl LangfuseBatchManager { pub fn new(public_key: String, secret_key: String, base_url: String) -> Self { Self { batch: Vec::new(), client: Client::builder() .timeout(Duration::from_secs(10)) .build() .expect("Failed to create HTTP client"), base_url, public_key, secret_key, } } pub fn spawn_sender(manager: Arc>) { const BATCH_INTERVAL: Duration = Duration::from_secs(5); tokio::spawn(async move { loop { tokio::time::sleep(BATCH_INTERVAL).await; if let Err(e) = manager.lock().await.send() { tracing::error!( error.msg = %e, error.type = %std::any::type_name_of_val(&e), "Failed to send batch to Langfuse" ); } } }); } pub async fn send_async(&mut self) -> Result<(), Box> { if self.batch.is_empty() { return Ok(()); } let payload = json!({ "batch": self.batch }); let base_url = Url::parse(&self.base_url).map_err(|e| format!("Invalid base URL: {e}"))?; let url = base_url .join("api/public/ingestion") .map_err(|e| format!("Failed to construct endpoint URL: {e}"))?; let response = self .client .post(url) .basic_auth(&self.public_key, Some(&self.secret_key)) .json(&payload) .send() .await?; match response.status() { status if status.is_success() => { let response_body: LangfuseIngestionResponse = response.json().await?; for error in &response_body.errors { tracing::error!( id = %error.id, status = error.status, message = error.message.as_deref().unwrap_or("No message"), error = ?error.error, "Partial failure in batch ingestion" ); } if !response_body.successes.is_empty() { self.batch.clear(); } if response_body.successes.is_empty() && !response_body.errors.is_empty() { Err("Langfuse ingestion failed for all items".into()) } else { Ok(()) } } status @ (StatusCode::BAD_REQUEST | StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN | StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED) => { let err_text = response.text().await.unwrap_or_default(); Err(format!("Langfuse API error: {}: {}", status, err_text).into()) } status => { let err_text = response.text().await.unwrap_or_default(); Err(format!("Unexpected status code: {}: {}", status, err_text).into()) } } } } impl BatchManager for LangfuseBatchManager { fn add_event(&mut self, event_type: &str, body: Value) { self.batch.push(json!({ "id": Uuid::new_v4().to_string(), "timestamp": Utc::now().to_rfc3339(), "type": event_type, "body": body })); } fn send(&mut self) -> Result<(), Box> { tokio::task::block_in_place(|| { tokio::runtime::Handle::current().block_on(self.send_async()) }) } fn is_empty(&self) -> bool { self.batch.is_empty() } } pub fn create_langfuse_observer() -> Option { let public_key = env::var("LANGFUSE_PUBLIC_KEY") .or_else(|_| env::var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY")) .unwrap_or_default(); // Use empty string if not found let secret_key = env::var("LANGFUSE_SECRET_KEY") .or_else(|_| env::var("LANGFUSE_INIT_PROJECT_SECRET_KEY")) .unwrap_or_default(); // Use empty string if not found // Return None if either key is empty if public_key.is_empty() || secret_key.is_empty() { return None; } let base_url = env::var("LANGFUSE_URL") .or_else(|_| env::var("LANGFUSE_BASE_URL")) .or_else(|_| env::var("LANGFUSE_HOST")) .unwrap_or_else(|_| DEFAULT_LANGFUSE_URL.to_string()); let batch_manager = Arc::new(Mutex::new(LangfuseBatchManager::new( public_key, secret_key, base_url, ))); if !cfg!(test) { LangfuseBatchManager::spawn_sender(batch_manager.clone()); } Some(ObservationLayer { batch_manager, span_tracker: Arc::new(Mutex::new(SpanTracker::new())), }) } #[cfg(test)] mod tests { use super::*; use serde_json::json; use std::collections::HashMap; use tokio::sync::Mutex; use tracing::dispatcher; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; struct TestFixture { original_subscriber: Option, original_env_vars: HashMap, mock_server: Option, } impl TestFixture { async fn new() -> Self { Self { original_subscriber: Some(dispatcher::get_default(dispatcher::Dispatch::clone)), original_env_vars: Self::save_env_vars(), mock_server: None, } } fn save_env_vars() -> HashMap { [ "LANGFUSE_PUBLIC_KEY", "LANGFUSE_INIT_PROJECT_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_INIT_PROJECT_SECRET_KEY", "LANGFUSE_URL", ] .iter() .filter_map(|&var| env::var(var).ok().map(|val| (var.to_string(), val))) .collect() } async fn with_mock_server(mut self) -> Self { self.mock_server = Some(MockServer::start().await); self } fn mock_server_uri(&self) -> String { self.mock_server .as_ref() .expect("Mock server not initialized") .uri() } async fn mock_response(&self, status: u16, body: Value) { Mock::given(method("POST")) .and(path("/api/public/ingestion")) .respond_with(ResponseTemplate::new(status).set_body_json(body)) .mount(self.mock_server.as_ref().unwrap()) .await; } } impl Drop for TestFixture { fn drop(&mut self) { // Restore original subscriber if let Some(subscriber) = &self.original_subscriber { let _ = dispatcher::set_global_default(subscriber.clone()); } // Restore environment for var in [ "LANGFUSE_PUBLIC_KEY", "LANGFUSE_INIT_PROJECT_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_INIT_PROJECT_SECRET_KEY", "LANGFUSE_URL", ] { if let Some(value) = self.original_env_vars.get(var) { env::set_var(var, value); } else { env::remove_var(var); } } } } fn create_test_event() -> Value { json!({ "name": "test_span", "type": "SPAN" }) } #[tokio::test] async fn test_batch_manager_creation() { let _fixture = TestFixture::new().await; let manager = LangfuseBatchManager::new( "test-public".to_string(), "test-secret".to_string(), "http://test.local".to_string(), ); assert_eq!(manager.public_key, "test-public"); assert_eq!(manager.secret_key, "test-secret"); assert_eq!(manager.base_url, "http://test.local"); assert!(manager.batch.is_empty()); } #[tokio::test] async fn test_add_event() { let _fixture = TestFixture::new().await; let mut manager = LangfuseBatchManager::new( "test-public".to_string(), "test-secret".to_string(), "http://test.local".to_string(), ); manager.add_event("test-event", create_test_event()); assert_eq!(manager.batch.len(), 1); let event = &manager.batch[0]; assert_eq!(event["type"], "test-event"); assert_eq!(event["body"], create_test_event()); assert!(event["id"].as_str().is_some()); assert!(event["timestamp"].as_str().is_some()); } #[tokio::test] async fn test_batch_send_success() { let fixture = TestFixture::new().await.with_mock_server().await; fixture .mock_response( 200, json!({ "successes": [{"id": "1", "status": 200}], "errors": [] }), ) .await; let mut manager = LangfuseBatchManager::new( "test-public".to_string(), "test-secret".to_string(), fixture.mock_server_uri(), ); manager.add_event("test-event", create_test_event()); let result = manager.send_async().await; assert!(result.is_ok()); assert!(manager.batch.is_empty()); } #[tokio::test] async fn test_batch_send_partial_failure() { let fixture = TestFixture::new().await.with_mock_server().await; fixture .mock_response( 200, json!({ "successes": [{"id": "1", "status": 200}], "errors": [{"id": "2", "status": 400, "message": "Invalid data"}] }), ) .await; let mut manager = LangfuseBatchManager::new( "test-public".to_string(), "test-secret".to_string(), fixture.mock_server_uri(), ); manager.add_event("test-event", create_test_event()); let result = manager.send_async().await; assert!(result.is_ok()); assert!(manager.batch.is_empty()); } #[tokio::test] async fn test_batch_send_complete_failure() { let fixture = TestFixture::new().await.with_mock_server().await; fixture .mock_response( 200, json!({ "successes": [], "errors": [{"id": "1", "status": 400, "message": "Invalid data"}] }), ) .await; let mut manager = LangfuseBatchManager::new( "test-public".to_string(), "test-secret".to_string(), fixture.mock_server_uri(), ); manager.add_event("test-event", create_test_event()); let result = manager.send_async().await; assert!(result.is_err()); assert!(!manager.batch.is_empty()); } #[tokio::test] async fn test_create_langfuse_observer() { let fixture = TestFixture::new().await.with_mock_server().await; // Test 1: No environment variables set - remove all possible variables for var in &[ "LANGFUSE_PUBLIC_KEY", "LANGFUSE_INIT_PROJECT_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", "LANGFUSE_INIT_PROJECT_SECRET_KEY", "LANGFUSE_URL", ] { env::remove_var(var); } let observer = create_langfuse_observer(); assert!( observer.is_none(), "Observer should be None without environment variables" ); // Test 2: Only public key set (regular) env::set_var("LANGFUSE_PUBLIC_KEY", "test-public-key"); let observer = create_langfuse_observer(); assert!( observer.is_none(), "Observer should be None with only public key" ); env::remove_var("LANGFUSE_PUBLIC_KEY"); // Test 3: Only secret key set (regular) env::set_var("LANGFUSE_SECRET_KEY", "test-secret-key"); let observer = create_langfuse_observer(); assert!( observer.is_none(), "Observer should be None with only secret key" ); env::remove_var("LANGFUSE_SECRET_KEY"); // Test 4: Only public key set (init project) env::set_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY", "test-public-key"); let observer = create_langfuse_observer(); assert!( observer.is_none(), "Observer should be None with only init project public key" ); env::remove_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY"); // Test 5: Only secret key set (init project) env::set_var("LANGFUSE_INIT_PROJECT_SECRET_KEY", "test-secret-key"); let observer = create_langfuse_observer(); assert!( observer.is_none(), "Observer should be None with only init project secret key" ); env::remove_var("LANGFUSE_INIT_PROJECT_SECRET_KEY"); // Test 6: Both regular keys set (should succeed) env::set_var("LANGFUSE_PUBLIC_KEY", "test-public-key"); env::set_var("LANGFUSE_SECRET_KEY", "test-secret-key"); env::set_var("LANGFUSE_URL", fixture.mock_server_uri()); let observer = create_langfuse_observer(); assert!( observer.is_some(), "Observer should be Some with both regular keys set" ); // Clean up regular keys env::remove_var("LANGFUSE_PUBLIC_KEY"); env::remove_var("LANGFUSE_SECRET_KEY"); // Test 7: Both init project keys set (should succeed) env::set_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY", "test-public-key"); env::set_var("LANGFUSE_INIT_PROJECT_SECRET_KEY", "test-secret-key"); let observer = create_langfuse_observer(); assert!( observer.is_some(), "Observer should be Some with both init project keys set" ); // Verify the observer has an empty batch manager let batch_manager = observer.unwrap().batch_manager; assert!(batch_manager.lock().await.is_empty()); } #[tokio::test] async fn test_batch_manager_spawn_sender() { let fixture = TestFixture::new().await.with_mock_server().await; fixture .mock_response( 200, json!({ "successes": [{"id": "1", "status": 200}], "errors": [] }), ) .await; let manager = Arc::new(Mutex::new(LangfuseBatchManager::new( "test-public".to_string(), "test-secret".to_string(), fixture.mock_server_uri(), ))); manager .lock() .await .add_event("test-event", create_test_event()); // Instead of spawning the sender which uses blocking operations, // test the async send directly let result = manager.lock().await.send_async().await; assert!(result.is_ok()); assert!(manager.lock().await.batch.is_empty()); } }