diff --git a/.gitignore b/.gitignore index 0b188bc..e2aab4f 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,8 @@ target/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + + +# Added by cargo + +/target diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..52d3681 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "llm-router-companion" +version = "0.1.0" +edition = "2024" + +[dependencies] +clap = { version = "4.6.7", features = ["derive"] } +dirs = "7.0.0" +ksni = "0.3.7" +reqwest = { version = "0.13.5", features = ["json"] } +serde = { version = "1.0.229", features = ["derive"] } +tokio = { version = "1.54.1", features = ["full"] } +toml = "1.1.8" diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..8cfe7da --- /dev/null +++ b/src/main.rs @@ -0,0 +1,378 @@ +use clap::Parser; +use ksni::{Tray, TrayMethods}; +use serde::{Deserialize, Serialize}; +use std::{ + fs, + io::{self, Write}, + path::PathBuf, + process::Command, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +#[derive(Parser, Debug)] +#[command( + name = "llm-router-tray", + version, + about = "LLM Server Status Tray Monitor with Systemd Integration and Preferred Backend" +)] +struct Args { + /// Interactively setup configuration and install the systemd user service + #[arg(long)] + install: bool, + + /// Start the installed systemd user service + #[arg(long)] + start_service: bool, + + /// Uninstall the service, remove configuration, and clean up systemd files + #[arg(long)] + uninstall: bool, +} + +#[derive(Serialize, Deserialize, Clone)] +struct Config { + backend_urls: Vec, + #[serde(default)] + preferred_backend: Option, +} + +#[derive(Clone)] +struct Backend { + url: String, + healthy: Arc, +} + +#[derive(Debug)] +struct LLMTray { + backend_data: Vec<(String, bool)>, + preferred_url: Option, +} + +impl Tray for LLMTray { + fn id(&self) -> String { + "llm-router-tray".into() + } + + fn icon_name(&self) -> String { + "cpu".into() + } + + fn title(&self) -> String { + let online_count = self.backend_data.iter().filter(|(_, h)| *h).count(); + format!("LLM Monitor ({}/{})", online_count, self.backend_data.len()) + } + + fn tool_tip(&self) -> ksni::ToolTip { + let mut desc = String::from("Monitored LLM Endpoints:\n"); + for (url, healthy) in &self.backend_data { + let status = if *healthy { + "🟢 Online" + } else { + "🔴 Offline" + }; + let is_pref = self.preferred_url.as_ref() == Some(url); + let pref_tag = if is_pref { " [Preferred]" } else { "" }; + desc.push_str(&format!("• {}{}: {}\n", url, pref_tag, status)); + } + ksni::ToolTip { + title: "LLM Server Status".into(), + description: desc.trim().into(), + ..Default::default() + } + } + + fn menu(&self) -> Vec> { + use ksni::menu::*; + let mut items = vec![]; + items.push( + StandardItem { + label: "Live Endpoint Statuses:".into(), + enabled: false, + ..Default::default() + } + .into(), + ); + items.push(MenuItem::Separator); + + for (url, healthy) in &self.backend_data { + let status_str = if *healthy { + "🟢 Online" + } else { + "🔴 Offline" + }; + let is_pref = self.preferred_url.as_ref() == Some(url); + let label = if is_pref { + format!("★ {} [{}] (Preferred)", url, status_str) + } else { + format!("{} [{}]", url, status_str) + }; + items.push( + StandardItem { + label, + enabled: false, + ..Default::default() + } + .into(), + ); + } + + items.push(MenuItem::Separator); + items.push( + StandardItem { + label: "Quit Monitor".into(), + icon_name: "application-exit".into(), + activate: Box::new(|_| std::process::exit(0)), + ..Default::default() + } + .into(), + ); + items + } +} + +fn get_config_path() -> PathBuf { + let mut path = dirs::config_dir().unwrap_or_else(|| PathBuf::from(".")); + path.push("llm-router-tray"); + fs::create_dir_all(&path).unwrap(); + path.push("config.toml"); + path +} + +fn get_service_path() -> PathBuf { + let mut path = dirs::config_dir().unwrap_or_else(|| PathBuf::from(".")); + path.push("systemd"); + path.push("user"); + path.push("llm-router-tray.service"); + path +} + +fn uninstall_service() { + println!("--- Uninstalling LLM Router Tray ---"); + let _ = Command::new("systemctl") + .args(["--user", "stop", "llm-router-tray.service"]) + .status(); + println!("Stopped systemd service (if running)."); + + let _ = Command::new("systemctl") + .args(["--user", "disable", "llm-router-tray.service"]) + .status(); + println!("Disabled systemd service."); + + let service_path = get_service_path(); + if service_path.exists() { + if let Err(e) = fs::remove_file(&service_path) { + println!("Failed to remove service file: {}", e); + } else { + println!("Deleted service file: {}", service_path.display()); + } + } + + let _ = Command::new("systemctl") + .args(["--user", "daemon-reload"]) + .status(); + + let mut config_dir = dirs::config_dir().unwrap_or_else(|| PathBuf::from(".")); + config_dir.push("llm-router-tray"); + if config_dir.exists() { + if let Err(e) = fs::remove_dir_all(&config_dir) { + println!("Failed to remove config folder: {}", e); + } else { + println!("Removed configuration folder: {}", config_dir.display()); + } + } + println!("\nUninstallation complete!"); + std::process::exit(0); +} + +fn start_systemd_service() { + match Command::new("systemctl") + .args(["--user", "start", "llm-router-tray.service"]) + .status() + { + Ok(status) if status.success() => { + println!("Successfully started llm-router-tray.service!"); + } + _ => { + println!("Failed to start systemd service. Ensure it has been installed first."); + } + } + std::process::exit(0); +} + +fn install_systemd_service() { + let exe_path = match std::env::current_exe() { + Ok(path) => path, + Err(e) => { + println!("Failed to get current executable path: {}", e); + return; + } + }; + let service_path = get_service_path(); + if let Some(parent) = service_path.parent() { + let _ = fs::create_dir_all(parent); + } + + let service_content = format!( + r#"[Unit] +Description=LLM Router Tray Monitor Service +After=network.target graphical-session.target + +[Service] +Type=simple +ExecStart={} +Restart=on-failure +RestartSec=5 + +[Install] +WantedBy=default.target +"#, + exe_path.display() + ); + + if let Err(e) = fs::write(&service_path, service_content) { + println!("Failed to write systemd service file: {}", e); + return; + } + println!("Created systemd service at: {}", service_path.display()); + + let _ = Command::new("systemctl") + .args(["--user", "daemon-reload"]) + .status(); + + match Command::new("systemctl") + .args(["--user", "enable", "--now", "llm-router-tray.service"]) + .status() + { + Ok(status) if status.success() => { + println!("Successfully enabled and started llm-router-tray.service as a user daemon!"); + } + _ => { + println!("Warning: Could not automatically start systemd service."); + } + } +} + +fn install_interactive() { + let path = get_config_path(); + println!("--- LLM Router Tray Interactive Setup Wizard ---"); + let mut backend_urls = Vec::new(); + println!("Enter your LLM server backend URLs (leave blank and press Enter when finished):"); + loop { + print!("Backend URL [e.g., http://192.168.1.50:11434]: "); + io::stdout().flush().unwrap(); + let mut url = String::new(); + io::stdin().read_line(&mut url).unwrap(); + let url = url.trim().to_string(); + + if url.is_empty() { + if backend_urls.is_empty() { + println!("Error: You must provide at least one backend server URL."); + continue; + } + break; + } + backend_urls.push(url); + } + + let mut preferred_backend = None; + if backend_urls.len() > 1 { + print!("Enter a preferred backend URL from the list above (leave blank for none): "); + io::stdout().flush().unwrap(); + let mut pref = String::new(); + io::stdin().read_line(&mut pref).unwrap(); + let pref = pref.trim().to_string(); + if !pref.is_empty() { + if backend_urls.contains(&pref) { + preferred_backend = Some(pref); + } else { + println!( + "Warning: Provided preferred URL was not in the backend list. Ignoring preference." + ); + } + } + } + let config = Config { + backend_urls, + preferred_backend, + }; + let toml_content = toml::to_string_pretty(&config).expect("Failed to serialize config"); + fs::write(&path, toml_content).expect("Failed to write config file"); + println!("\nConfiguration successfully saved to {}\n", path.display()); + install_systemd_service(); + std::process::exit(0); +} + +fn load_config() -> Config { + let path = get_config_path(); + if path.exists() { + let content = fs::read_to_string(&path).expect("Failed to read config file"); + toml::from_str(&content).expect("Failed to parse config TOML") + } else { + println!("No configuration file found at {}.", path.display()); + println!("Running interactive setup wizard..."); + install_interactive(); + unreachable!() + } +} + +#[tokio::main] +async fn main() { + let args = Args::parse(); + if args.uninstall { + uninstall_service(); + } + if args.start_service { + start_systemd_service(); + } + if args.install { + install_interactive(); + } + let config = load_config(); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(2)) + .build() + .unwrap(); + let mut backends = Vec::new(); + let mut initial_backend_data = Vec::new(); + for url in config.backend_urls { + backends.push(Backend { + url: url.clone(), + healthy: Arc::new(AtomicBool::new(false)), + }); + initial_backend_data.push((url, false)); + } + let tray = LLMTray { + backend_data: initial_backend_data, + preferred_url: config.preferred_backend.clone(), + }; + let tray_handle = tray.spawn().await.expect("Failed to start system tray"); + let preferred_url_clone = config.preferred_backend.clone(); + + tokio::spawn(async move { + loop { + for backend in &backends { + let health_url = format!("{}/v1/models", backend.url); + let res = client.get(&health_url).send().await; + let is_healthy = res.map_or(false, |r| r.status().is_success()); + backend.healthy.store(is_healthy, Ordering::Relaxed); + } + let snapshot: Vec<(String, bool)> = backends + .iter() + .map(|b| (b.url.clone(), b.healthy.load(Ordering::Relaxed))) + .collect(); + let pref_url_ref = preferred_url_clone.clone(); + tray_handle + .update(move |tray: &mut LLMTray| { + tray.backend_data = snapshot; + tray.preferred_url = pref_url_ref; + }) + .await; + tokio::time::sleep(Duration::from_secs(5)).await; + } + }); + std::future::pending::<()>().await; +}