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
15static INIT: Once = Once::new();
17
18pub 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
31fn 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 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, log_dir,
56 log_filename,
57 );
58
59 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 let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
69 EnvFilter::new("")
71 .add_directive("mcp_client=debug".parse().unwrap())
73 .add_directive("aster=debug".parse().unwrap())
75 .add_directive("aster_cli=info".parse().unwrap())
77 .add_directive(LevelFilter::WARN.into())
79 });
80
81 let mut layers = vec![
83 file_layer.with_filter(env_filter).boxed(),
84 ];
86
87 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 let subscriber = Registry::default().with(layers);
120
121 if force {
122 let _guard = subscriber.set_default();
125 tracing::warn!("Test log entry from setup");
126 tracing::info!("Another test log entry from setup");
127 std::thread::sleep(std::time::Duration::from_millis(100));
129 Ok(())
130 } else {
131 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 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 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 for (var, _) in &original_vars {
200 env::remove_var(var);
201 }
202
203 assert!(langfuse_layer::create_langfuse_observer().is_none());
205
206 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 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 env::remove_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY");
220 assert!(langfuse_layer::create_langfuse_observer().is_none());
221
222 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}