use crate::basic::compiler::BasicCompiler; use crate::config::ConfigManager; use crate::shared::state::AppState; use aws_sdk_s3::Client; use log::{info}; use std::collections::HashMap; use std::error::Error; use std::sync::Arc; use tokio::time::{interval, Duration}; #[derive(Debug, Clone)] pub struct FileState { pub etag: String, } pub struct DriveMonitor { state: Arc, bucket_name: String, file_states: Arc>>, bot_id: uuid::Uuid, } impl DriveMonitor { pub fn new(state: Arc, bucket_name: String, bot_id: uuid::Uuid) -> Self { Self { state, bucket_name, file_states: Arc::new(tokio::sync::RwLock::new(HashMap::new())), bot_id, } } pub fn spawn(self: Arc) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { info!("Drive Monitor service started for bucket: {}", self.bucket_name); let mut tick = interval(Duration::from_secs(30)); loop { tick.tick().await; if let Err(e) = self.check_for_changes().await { log::error!("Error checking for drive changes: {}", e); } } }) } async fn check_for_changes(&self) -> Result<(), Box> { let client = match &self.state.drive { Some(client) => client, None => return Ok(()), }; self.check_gbdialog_changes(client).await?; // TODO: Remove self.check_gbot(client).await?; Ok(()) } async fn check_gbdialog_changes(&self, client: &Client) -> Result<(), Box> { let prefix = ".gbdialog/"; let mut current_files = HashMap::new(); let mut continuation_token = None; loop { let list_objects = match tokio::time::timeout( Duration::from_secs(30), client .list_objects_v2() .bucket(&self.bucket_name.to_lowercase()) .set_continuation_token(continuation_token) .send() ).await { Ok(Ok(list)) => list, Ok(Err(e)) => return Err(e.into()), Err(_) => { log::error!("Timeout listing objects in bucket {}", self.bucket_name); return Ok(()); } }; for obj in list_objects.contents.unwrap_or_default() { let path = obj.key().unwrap_or_default().to_string(); let path_parts: Vec<&str> = path.split('/').collect(); if path_parts.len() < 2 || !path_parts[0].ends_with(".gbdialog") { continue; } if path.ends_with('/') || !path.ends_with(".bas") { continue; } let file_state = FileState { etag: obj.e_tag().unwrap_or_default().to_string(), }; current_files.insert(path, file_state); } if !list_objects.is_truncated.unwrap_or(false) { break; } continuation_token = list_objects.next_continuation_token; } let mut file_states = self.file_states.write().await; for (path, current_state) in current_files.iter() { if let Some(previous_state) = file_states.get(path) { if current_state.etag != previous_state.etag { if let Err(e) = self.compile_tool(client, path).await { log::error!("Failed to compile tool {}: {}", path, e); } } } else { if let Err(e) = self.compile_tool(client, path).await { log::error!("Failed to compile tool {}: {}", path, e); } } } let previous_paths: Vec = file_states .keys() .filter(|k| k.starts_with(prefix)) .cloned() .collect(); for path in previous_paths { if !current_files.contains_key(&path) { file_states.remove(&path); } } for (path, state) in current_files { file_states.insert(path, state); } Ok(()) } async fn check_gbot(&self, client: &Client) -> Result<(), Box> { let config_manager = ConfigManager::new(self.state.conn.clone()); let mut continuation_token = None; loop { let list_objects = match tokio::time::timeout( Duration::from_secs(30), client .list_objects_v2() .bucket(&self.bucket_name.to_lowercase()) .set_continuation_token(continuation_token) .send() ).await { Ok(Ok(list)) => list, Ok(Err(e)) => return Err(e.into()), Err(_) => { log::error!("Timeout listing objects in bucket {}", self.bucket_name); return Ok(()); } }; for obj in list_objects.contents.unwrap_or_default() { let path = obj.key().unwrap_or_default().to_string(); let path_parts: Vec<&str> = path.split('/').collect(); if path_parts.len() < 2 || !path_parts[0].ends_with(".gbot") { continue; } if !path.ends_with("config.csv") { continue; } match client.head_object().bucket(&self.bucket_name).key(&path).send().await { Ok(_head_res) => { let response = client.get_object().bucket(&self.bucket_name).key(&path).send().await?; let bytes = response.body.collect().await?.into_bytes(); let csv_content = String::from_utf8(bytes.to_vec()) .map_err(|e| format!("UTF-8 error in {}: {}", path, e))?; let llm_lines: Vec<_> = csv_content .lines() .filter(|line| line.trim_start().starts_with("llm-")) .collect(); if !llm_lines.is_empty() { use crate::llm::local::ensure_llama_servers_running; let mut restart_needed = false; for line in llm_lines { let parts: Vec<&str> = line.split(',').collect(); if parts.len() >= 2 { let key = parts[0].trim(); let new_value = parts[1].trim(); match config_manager.get_config(&self.bot_id, key, None) { Ok(old_value) => { if old_value != new_value { info!("Detected change in {} (old: {}, new: {})", key, old_value, new_value); restart_needed = true; } } Err(_) => { restart_needed = true; } } } } let _ = config_manager.sync_gbot_config(&self.bot_id, &csv_content); if restart_needed { if let Err(e) = ensure_llama_servers_running(Arc::clone(&self.state)).await { log::error!("Failed to restart LLaMA servers after llm- config change: {}", e); } } } else { let _ = config_manager.sync_gbot_config(&self.bot_id, &csv_content); } if csv_content.lines().any(|line| line.starts_with("theme-")) { self.broadcast_theme_change(&csv_content).await?; } } Err(e) => { log::error!("Config file {} not found or inaccessible: {}", path, e); } } } if !list_objects.is_truncated.unwrap_or(false) { break; } continuation_token = list_objects.next_continuation_token; } Ok(()) } async fn broadcast_theme_change(&self, csv_content: &str) -> Result<(), Box> { let mut theme_data = serde_json::json!({ "event": "change_theme", "data": {} }); for line in csv_content.lines() { let parts: Vec<&str> = line.split(',').collect(); if parts.len() >= 2 { let key = parts[0].trim(); let value = parts[1].trim(); match key { "theme-color1" => theme_data["data"]["color1"] = serde_json::Value::String(value.to_string()), "theme-color2" => theme_data["data"]["color2"] = serde_json::Value::String(value.to_string()), "theme-logo" => theme_data["data"]["logo_url"] = serde_json::Value::String(value.to_string()), "theme-title" => theme_data["data"]["title"] = serde_json::Value::String(value.to_string()), "theme-logo-text" => theme_data["data"]["logo_text"] = serde_json::Value::String(value.to_string()), _ => {} } } } let response_channels = self.state.response_channels.lock().await; for (session_id, tx) in response_channels.iter() { let theme_response = crate::shared::models::BotResponse { bot_id: self.bot_id.to_string(), user_id: "system".to_string(), session_id: session_id.clone(), channel: "web".to_string(), content: serde_json::to_string(&theme_data)?, message_type: 2, stream_token: None, is_complete: true, suggestions: Vec::new(), context_name: None, context_length: 0, context_max_length: 0, }; let _ = tx.try_send(theme_response); } Ok(()) } async fn compile_tool(&self, client: &Client, file_path: &str) -> Result<(), Box> { info!("Fetching object from Drive: bucket={}, key={}", &self.bucket_name, file_path); let response = match client.get_object().bucket(&self.bucket_name).key(file_path).send().await { Ok(res) => { info!("Successfully fetched object from Drive: bucket={}, key={}, size={}", &self.bucket_name, file_path, res.content_length().unwrap_or(0)); res } Err(e) => { log::error!("Failed to fetch object from Drive: bucket={}, key={}, error={:?}", &self.bucket_name, file_path, e); return Err(e.into()); } }; let bytes = response.body.collect().await?.into_bytes(); let source_content = String::from_utf8(bytes.to_vec())?; let tool_name = file_path .split('/') .last() .unwrap_or(file_path) .strip_suffix(".bas") .unwrap_or(file_path) .to_string(); let bot_name = self.bucket_name.strip_suffix(".gbai").unwrap_or(&self.bucket_name); let work_dir = format!("./work/{}.gbai/{}.gbdialog", bot_name, bot_name); let state_clone = Arc::clone(&self.state); let work_dir_clone = work_dir.clone(); let tool_name_clone = tool_name.clone(); let source_content_clone = source_content.clone(); let bot_id = self.bot_id; tokio::task::spawn_blocking(move || { std::fs::create_dir_all(&work_dir_clone)?; let local_source_path = format!("{}/{}.bas", work_dir_clone, tool_name_clone); std::fs::write(&local_source_path, &source_content_clone)?; let mut compiler = BasicCompiler::new(state_clone, bot_id); let result = compiler.compile_file(&local_source_path, &work_dir_clone)?; if let Some(mcp_tool) = result.mcp_tool { info!("MCP tool definition generated with {} parameters", mcp_tool.input_schema.properties.len()); } Ok::<(), Box>(()) }).await??; Ok(()) } }