Skip to main content

aster_cli/
logging.rs

1use anyhow::{Context, Result};
2use std::sync::Arc;
3use std::sync::Once;
4use tokio::sync::Mutex;
5use tracing_appender::rolling::Rotation;
6use tracing_subscriber::{
7    filter::LevelFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer,
8    Registry,
9};
10
11use aster::tracing::{langfuse_layer, otlp_layer};
12use aster_bench::bench_session::BenchAgentError;
13use aster_bench::error_capture::ErrorCaptureLayer;
14
15// Used to ensure we only set up tracing once
16static INIT: Once = Once::new();
17
18/// Sets up the logging infrastructure for the application.
19/// This includes:
20/// - File-based logging with JSON formatting (DEBUG level)
21/// - No console output (all logs go to files only)
22/// - Optional Langfuse integration (DEBUG level)
23/// - Optional error capture layer for benchmarking
24pub fn setup_logging(
25    name: Option<&str>,
26    error_capture: Option<Arc<Mutex<Vec<BenchAgentError>>>>,
27) -> Result<()> {
28    setup_logging_internal(name, error_capture, false)
29}
30
31/// Internal function that allows bypassing the Once check for testing
32fn setup_logging_internal(
33    name: Option<&str>,
34    error_capture: Option<Arc<Mutex<Vec<BenchAgentError>>>>,
35    force: bool,
36) -> Result<()> {
37    let mut result = Ok(());
38
39    // Register the error vector if provided
40    if let Some(errors) = error_capture {
41        ErrorCaptureLayer::register_error_vector(errors);
42    }
43
44    let mut setup = || {
45        result = (|| {
46            let log_dir = aster::logging::prepare_log_directory("cli", true)?;
47            let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S").to_string();
48            let log_filename = if let Some(n) = name {
49                format!("{}-{}.log", timestamp, n)
50            } else {
51                format!("{}.log", timestamp)
52            };
53            let file_appender = tracing_appender::rolling::RollingFileAppender::new(
54                Rotation::NEVER, // we do manual rotation via file naming and cleanup_old_logs
55                log_dir,
56                log_filename,
57            );
58
59            // Create JSON file logging layer with all logs (DEBUG and above)
60            let file_layer = fmt::layer()
61                .with_target(true)
62                .with_level(true)
63                .with_writer(file_appender)
64                .with_ansi(false)
65                .json();
66
67            // Base filter
68            let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
69                // Set default levels for different modules
70                EnvFilter::new("")
71                    // Set mcp-client to DEBUG
72                    .add_directive("mcp_client=debug".parse().unwrap())
73                    // Set aster module to DEBUG
74                    .add_directive("aster=debug".parse().unwrap())
75                    // Set aster-cli to INFO
76                    .add_directive("aster_cli=info".parse().unwrap())
77                    // Set everything else to WARN
78                    .add_directive(LevelFilter::WARN.into())
79            });
80
81            // Start building the subscriber
82            let mut layers = vec![
83                file_layer.with_filter(env_filter).boxed(),
84                // Console logging disabled for CLI - all logs go to files only
85            ];
86
87            // Only add ErrorCaptureLayer if not in test mode
88            if !force {
89                layers.push(ErrorCaptureLayer::new().boxed());
90            }
91
92            if !force {
93                if let Ok((otlp_tracing_layer, otlp_metrics_layer, otlp_logs_layer)) =
94                    otlp_layer::init_otlp()
95                {
96                    layers.push(
97                        otlp_tracing_layer
98                            .with_filter(otlp_layer::create_otlp_tracing_filter())
99                            .boxed(),
100                    );
101                    layers.push(
102                        otlp_metrics_layer
103                            .with_filter(otlp_layer::create_otlp_metrics_filter())
104                            .boxed(),
105                    );
106                    layers.push(
107                        otlp_logs_layer
108                            .with_filter(otlp_layer::create_otlp_logs_filter())
109                            .boxed(),
110                    );
111                }
112            }
113
114            if let Some(langfuse) = langfuse_layer::create_langfuse_observer() {
115                layers.push(langfuse.with_filter(LevelFilter::DEBUG).boxed());
116            }
117
118            // Build the subscriber
119            let subscriber = Registry::default().with(layers);
120
121            if force {
122                // For testing, just create and use the subscriber without setting it globally
123                // Write a test log to ensure the file is created
124                let _guard = subscriber.set_default();
125                tracing::warn!("Test log entry from setup");
126                tracing::info!("Another test log entry from setup");
127                // Flush the output
128                std::thread::sleep(std::time::Duration::from_millis(100));
129                Ok(())
130            } else {
131                // For normal operation, set the subscriber globally
132                subscriber
133                    .try_init()
134                    .context("Failed to set global subscriber")?;
135                Ok(())
136            }
137        })();
138    };
139
140    if force {
141        setup();
142    } else {
143        INIT.call_once(setup);
144    }
145
146    result
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use std::env;
153    use tempfile::TempDir;
154
155    fn setup_temp_home() -> TempDir {
156        let temp_dir = TempDir::new().unwrap();
157        if cfg!(windows) {
158            env::set_var("USERPROFILE", temp_dir.path());
159        } else {
160            env::set_var("HOME", temp_dir.path());
161        }
162        temp_dir
163    }
164
165    #[test]
166    fn test_log_directory_creation() {
167        let _temp_dir = setup_temp_home();
168        let log_dir = aster::logging::prepare_log_directory("cli", true).unwrap();
169        assert!(log_dir.exists());
170        assert!(log_dir.is_dir());
171
172        // Verify directory structure
173        let path_components: Vec<_> = log_dir.components().collect();
174        assert!(path_components.iter().any(|c| c.as_os_str() == "aster"));
175        assert!(path_components.iter().any(|c| c.as_os_str() == "logs"));
176        assert!(path_components.iter().any(|c| c.as_os_str() == "cli"));
177    }
178
179    #[tokio::test]
180    async fn test_langfuse_layer_creation() {
181        let _temp_dir = setup_temp_home();
182
183        // Store original environment variables (both sets)
184        let original_vars = [
185            ("LANGFUSE_PUBLIC_KEY", env::var("LANGFUSE_PUBLIC_KEY").ok()),
186            ("LANGFUSE_SECRET_KEY", env::var("LANGFUSE_SECRET_KEY").ok()),
187            ("LANGFUSE_URL", env::var("LANGFUSE_URL").ok()),
188            (
189                "LANGFUSE_INIT_PROJECT_PUBLIC_KEY",
190                env::var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY").ok(),
191            ),
192            (
193                "LANGFUSE_INIT_PROJECT_SECRET_KEY",
194                env::var("LANGFUSE_INIT_PROJECT_SECRET_KEY").ok(),
195            ),
196        ];
197
198        // Clear all Langfuse environment variables
199        for (var, _) in &original_vars {
200            env::remove_var(var);
201        }
202
203        // Test without any environment variables
204        assert!(langfuse_layer::create_langfuse_observer().is_none());
205
206        // Test with standard Langfuse variables
207        env::set_var("LANGFUSE_PUBLIC_KEY", "test_public_key");
208        env::set_var("LANGFUSE_SECRET_KEY", "test_secret_key");
209        assert!(langfuse_layer::create_langfuse_observer().is_some());
210
211        // Clear and test with init project variables
212        env::remove_var("LANGFUSE_PUBLIC_KEY");
213        env::remove_var("LANGFUSE_SECRET_KEY");
214        env::set_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY", "test_public_key");
215        env::set_var("LANGFUSE_INIT_PROJECT_SECRET_KEY", "test_secret_key");
216        assert!(langfuse_layer::create_langfuse_observer().is_some());
217
218        // Test fallback behavior
219        env::remove_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY");
220        assert!(langfuse_layer::create_langfuse_observer().is_none());
221
222        // Restore original environment variables
223        for (var, value) in original_vars {
224            match value {
225                Some(val) => env::set_var(var, val),
226                None => env::remove_var(var),
227            }
228        }
229    }
230}