write the server

.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

Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "llm-router"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+axum = "0.7"
+tokio = { version = "1.0", features = ["full"] }
+reqwest = { version = "0.12", features = ["stream"] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+toml = "0.8"
+dirs = "5.0"
+futures-util = "0.3"
+ksni = "0.3"
+clap = { version = "4.4", features = ["derive"] }

src/main.rs

@@ -0,0 +1,469 @@
+use axum::{
+    body::Body,
+    extract::State,
+    http::{Request, StatusCode},
+    response::{IntoResponse, Response},
+    routing::any,
+    Router,
+};
+use clap::Parser;
+use futures_util::StreamExt;
+use ksni::{Tray, TrayMethods};
+use serde::{Deserialize, Serialize};
+use std::{
+    fs,
+    io::{self, Write},
+    path::PathBuf,
+    process::Command,
+    sync::{
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+#[derive(Parser, Debug)]
+#[command(
+    name = "llm-router",
+    version,
+    about = "LLM Router Proxy with System Tray and Systemd Integration"
+)]
+struct Args {
+    /// Run the proxy service in daemon mode (non-interactive, used by systemd)
+    #[arg(long)]
+    daemon: 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 {
+    listen_port: u16,
+    backend_urls: Vec<String>,
+}
+
+#[derive(Clone)]
+struct Backend {
+    url: String,
+    healthy: Arc<AtomicBool>,
+}
+
+#[derive(Clone)]
+struct AppState {
+    backends: Vec<Backend>,
+    client: reqwest::Client,
+    counter: Arc<AtomicUsize>,
+}
+#[derive(Debug)]
+struct LLMRouterTray {
+    backend_data: Vec<(String, bool)>,
+}
+
+impl Tray for LLMRouterTray {
+    fn id(&self) -> String {
+        "llm-router".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 Router ({}/{})", online_count, self.backend_data.len())
+    }
+
+    fn tool_tip(&self) -> ksni::ToolTip {
+        let mut desc = String::from("Managed LLM Endpoints:\n");
+        for (url, healthy) in &self.backend_data {
+            let status = if *healthy {
+                "🟢 Online"
+            } else {
+                "🔴 Offline"
+            };
+            desc.push_str(&format!("• {}: {}\n", url, status));
+        }
+        ksni::ToolTip {
+            title: "LLM Router Status".into(),
+            description: desc.trim().into(),
+            ..Default::default()
+        }
+    }
+
+    fn menu(&self) -> Vec<ksni::menu::MenuItem<Self>> {
+        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"
+            };
+            items.push(
+                StandardItem {
+                    label: format!("{} [{}]", url, status_str),
+                    enabled: false,
+                    ..Default::default()
+                }
+                .into(),
+            );
+        }
+        items.push(MenuItem::Separator);
+        items.push(
+            StandardItem {
+                label: "Quit Router".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");
+    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.service");
+    path
+}
+
+fn uninstall_service() {
+    println!("--- Uninstalling LLM Router ---");
+    let _ = Command::new("systemctl")
+        .args(["--user", "stop", "llm-router.service"])
+        .status();
+    println!("Stopped systemd service (if running).");
+    let _ = Command::new("systemctl")
+        .args(["--user", "disable", "llm-router.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");
+    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.service"])
+        .status()
+    {
+        Ok(status) if status.success() => {
+            println!("Successfully started llm-router.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 Proxy Service
+After=network.target
+
+[Service]
+Type=simple
+ExecStart={} --daemon
+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.service"])
+        .status()
+    {
+        Ok(status) if status.success() => {
+            println!("Successfully enabled and started llm-router.service as a user daemon!");
+        }
+        _ => {
+            println!("Warning: Could not automatically start systemd service.");
+        }
+    }
+}
+
+fn load_config(is_daemon: bool) -> 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 {
+        if is_daemon {
+            eprintln!("Error: Configuration file not found at {}.", path.display());
+            eprintln!("Please run the application interactively at least once without '--daemon' to complete setup.");
+            std::process::exit(1);
+        }
+
+        println!("No configuration file found at {}.", path.display());
+        println!("--- Interactive Setup Wizard ---");
+
+        print!("Enter port for proxy to listen on [default: 4000]: ");
+        io::stdout().flush().unwrap();
+        let mut port_str = String::new();
+        io::stdin().read_line(&mut port_str).unwrap();
+        let listen_port = port_str.trim().parse().unwrap_or(4000);
+
+        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 config = Config {
+            listen_port,
+            backend_urls,
+        };
+
+        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());
+
+        print!(
+            "Would you like to install and run this as a background systemd user service? [y/N]: "
+        );
+        io::stdout().flush().unwrap();
+        let mut install_ans = String::new();
+        io::stdin().read_line(&mut install_ans).unwrap();
+        if install_ans.trim().to_lowercase().starts_with('y') {
+            install_systemd_service();
+        }
+
+        config
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    let args = Args::parse();
+    if args.uninstall {
+        uninstall_service();
+    }
+    if args.start_service {
+        start_systemd_service();
+    }
+    let config = load_config(args.daemon);
+    let client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(120))
+        .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 = LLMRouterTray {
+        backend_data: initial_backend_data,
+    };
+    let tray_handle = tray.spawn().await.expect("Failed to start system tray");
+
+    let state = AppState {
+        backends: backends.clone(),
+        client: client.clone(),
+        counter: Arc::new(AtomicUsize::new(0)),
+    };
+    tokio::spawn(async move {
+        health_check_loop(backends, client, tray_handle).await;
+    });
+    let app = Router::new()
+        .route("/*path", any(proxy_handler))
+        .fallback(any(proxy_handler))
+        .with_state(state);
+
+    let bind_address = format!("0.0.0.0:{}", config.listen_port);
+    let listener = tokio::net::TcpListener::bind(&bind_address).await.unwrap();
+    println!(
+        "LLM Router running on http://0.0.0.0:{}",
+        config.listen_port
+    );
+    axum::serve(listener, app).await.unwrap();
+}
+
+async fn health_check_loop(
+    backends: Vec<Backend>,
+    client: reqwest::Client,
+    tray_handle: ksni::Handle<LLMRouterTray>,
+) {
+    loop {
+        for backend in &backends {
+            let health_url = format!("{}/v1/models", backend.url);
+            let res = client
+                .get(&health_url)
+                .timeout(Duration::from_secs(2))
+                .send()
+                .await;
+            let is_healthy = res.map_or(false, |r| r.status().is_success());
+            let prev = backend.healthy.load(Ordering::Relaxed);
+
+            if prev != is_healthy {
+                println!(
+                    "Backend [{}] status changed: {}",
+                    backend.url,
+                    if is_healthy { "ONLINE" } else { "OFFLINE" }
+                );
+            }
+            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();
+        tray_handle
+            .update(move |tray: &mut LLMRouterTray| {
+                tray.backend_data = snapshot;
+            })
+            .await;
+        tokio::time::sleep(Duration::from_secs(5)).await;
+    }
+}
+
+fn select_backend(backends: &[Backend], counter: &AtomicUsize) -> Option<String> {
+    let len = backends.len();
+    if len == 0 {
+        return None;
+    }
+    let start_idx = counter.fetch_add(1, Ordering::Relaxed);
+    for i in 0..len {
+        let idx = (start_idx + i) % len;
+        if backends[idx].healthy.load(Ordering::Relaxed) {
+            return Some(backends[idx].url.clone());
+        }
+    }
+    None
+}
+
+async fn proxy_handler(State(state): State<AppState>, req: Request<Body>) -> Response {
+    let backend_url = match select_backend(&state.backends, &state.counter) {
+        Some(url) => url,
+        None => {
+            return (
+                StatusCode::SERVICE_UNAVAILABLE,
+                "All LLM backends are offline",
+            )
+                .into_response()
+        }
+    };
+    let path = req.uri().path().to_string();
+    let query = req
+        .uri()
+        .query()
+        .map(|v| format!("?{}", v))
+        .unwrap_or_default();
+    let target_url = format!("{}{}{}", backend_url, path, query);
+    let method = req.method().clone();
+    let headers = req.headers().clone();
+    let body_bytes = match axum::body::to_bytes(req.into_body(), usize::MAX).await {
+        Ok(bytes) => bytes,
+        Err(_) => return (StatusCode::BAD_REQUEST, "Failed to read request body").into_response(),
+    };
+    let mut forward_req = state.client.request(method, &target_url).body(body_bytes);
+    for (key, value) in headers.iter() {
+        forward_req = forward_req.header(key, value);
+    }
+    match forward_req.send().await {
+        Ok(resp) => {
+            let status = resp.status();
+            let mut response_builder = Response::builder().status(status);
+            for (key, value) in resp.headers() {
+                response_builder = response_builder.header(key, value);
+            }
+            let stream = resp.bytes_stream().map(|result| {
+                result.map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err))
+            });
+            let body = Body::from_stream(stream);
+            response_builder.body(body).unwrap_or_else(|_| {
+                (StatusCode::INTERNAL_SERVER_ERROR, "Response build error").into_response()
+            })
+        }
+        Err(_) => (StatusCode::BAD_GATEWAY, "Failed to reach LLM backend").into_response(),
+    }
+}