diff --git a/Cargo.lock b/Cargo.lock index 80ecd4a..6f25ec2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "ahash" version = "0.8.12" @@ -138,6 +144,15 @@ version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bumpalo" version = "3.19.1" @@ -282,6 +297,44 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "dirs" version = "5.0.1" @@ -397,6 +450,17 @@ dependencies = [ "winapi", ] +[[package]] +name = "filetime" +version = "0.2.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f98844151eee8917efc50bd9e8318cb963ae8b297431495d3f758616ea5c57db" +dependencies = [ + "cfg-if", + "libc", + "libredox", +] + [[package]] name = "find-msvc-tools" version = "0.1.8" @@ -412,6 +476,16 @@ dependencies = [ "glob", ] +[[package]] +name = "flate2" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b375d6465b98090a5f25b1c7703f3859783755aa9a80433b36e0379a3ec2f369" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -482,6 +556,16 @@ dependencies = [ "slab", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -891,6 +975,7 @@ checksum = "3d0b95e02c851351f877147b7deea7b1afb1df71b63aa5f8270716e0c5720616" dependencies = [ "bitflags 2.10.0", "libc", + "redox_syscall 0.7.0", ] [[package]] @@ -985,6 +1070,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.1.1" @@ -1071,7 +1166,7 @@ checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" dependencies = [ "cfg-if", "libc", - "redox_syscall", + "redox_syscall 0.5.18", "smallvec", "windows-link", ] @@ -1272,6 +1367,15 @@ dependencies = [ "bitflags 2.10.0", ] +[[package]] +name = "redox_syscall" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49f3fe0889e69e2ae9e41f4d6c4c0181701d00e4697b356fb1f74173a5e0ee27" +dependencies = [ + "bitflags 2.10.0", +] + [[package]] name = "redox_users" version = "0.4.6" @@ -1568,6 +1672,17 @@ dependencies = [ "serial-core", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shared_library" version = "0.1.9" @@ -1622,6 +1737,12 @@ dependencies = [ "tokio", ] +[[package]] +name = "simd-adler32" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" + [[package]] name = "slab" version = "0.4.11" @@ -1702,6 +1823,17 @@ dependencies = [ "syn", ] +[[package]] +name = "tar" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d863878d212c87a19c1a610eb53bb01fe12951c0501cf5a0d65f724914a667a" +dependencies = [ + "filetime", + "libc", + "xattr", +] + [[package]] name = "tempfile" version = "3.24.0" @@ -1974,6 +2106,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + [[package]] name = "unicode-ident" version = "1.0.22" @@ -2061,14 +2199,18 @@ dependencies = [ "anyhow", "chrono", "clap", + "flate2", + "futures-util", "indicatif", "libc", "portable-pty", + "reqwest", "rusqlite", "serde", "serde_json", "signal-hook", "signal-hook-tokio", + "tar", "tempfile", "term_size", "tokio", @@ -2098,6 +2240,7 @@ dependencies = [ name = "wake-llm" version = "0.6.0" dependencies = [ + "chrono", "dirs", "futures-util", "indicatif", @@ -2105,9 +2248,11 @@ dependencies = [ "reqwest", "serde", "serde_json", + "sha2", "thiserror 1.0.69", "tokio", "tracing", + "wake-core", ] [[package]] @@ -2607,6 +2752,16 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" +[[package]] +name = "xattr" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" +dependencies = [ + "libc", + "rustix", +] + [[package]] name = "yoke" version = "0.8.1" diff --git a/README.md b/README.md index 63a59e1..61a06cf 100644 --- a/README.md +++ b/README.md @@ -111,7 +111,7 @@ Old sessions are automatically pruned on each `wake shell` start. Wake automatically summarizes command outputs using a local LLM (Qwen3-0.6B). Summaries appear in `wake_list_commands` output, helping Claude quickly understand what happened without reading full output. -**Enabled by default.** On first run, the model (~468MB) downloads automatically. +**Enabled by default.** On first run, the model (~380MB) downloads automatically. - **CPU-friendly** — The small model runs efficiently without a GPU - **Privacy** — All inference happens locally, nothing leaves your machine diff --git a/crates/wake-cli/Cargo.toml b/crates/wake-cli/Cargo.toml index d986fdb..c8f3782 100644 --- a/crates/wake-cli/Cargo.toml +++ b/crates/wake-cli/Cargo.toml @@ -27,6 +27,12 @@ anyhow.workspace = true term_size.workspace = true libc.workspace = true +# For update command +reqwest = { version = "0.12", default-features = false, features = ["stream", "rustls-tls", "json"] } +futures-util = "0.3" +flate2 = "1" +tar = "0.4" + [dev-dependencies] tempfile = "3" rusqlite = { workspace = true } diff --git a/crates/wake-cli/src/commands/llm.rs b/crates/wake-cli/src/commands/llm.rs index c82aa6e..84ff801 100644 --- a/crates/wake-cli/src/commands/llm.rs +++ b/crates/wake-cli/src/commands/llm.rs @@ -1,7 +1,9 @@ //! LLM model management commands use anyhow::Result; -use wake_llm::{WakeLlm, DEFAULT_MODEL_SIZE}; +use std::fs; +use wake_core::ModelManifest; +use wake_llm::{WakeLlm, DEFAULT_MODEL, DEFAULT_MODEL_SIZE}; /// Show LLM model status pub async fn status() -> Result<()> { @@ -83,3 +85,155 @@ pub async fn download() -> Result<()> { Ok(()) } + +/// Represents a model file that can be cleaned +struct CleanableModel { + filename: String, + path: std::path::PathBuf, + size: u64, + reason: CleanReason, +} + +enum CleanReason { + /// Model is tracked in manifest but not the current default model + OrphanedModel, + /// Model file exists but is not tracked in manifest + UntrackedFile, +} + +impl CleanableModel { + fn reason_str(&self) -> &'static str { + match self.reason { + CleanReason::OrphanedModel => "orphaned (from previous version)", + CleanReason::UntrackedFile => "untracked .gguf file", + } + } +} + +/// Clean orphaned models from previous versions +pub async fn clean(dry_run: bool, force: bool) -> Result<()> { + let models_dir = ModelManifest::models_dir()?; + + if !models_dir.exists() { + println!("No models directory found at: {}", models_dir.display()); + return Ok(()); + } + + // Load manifest + let mut manifest = ModelManifest::load()?; + + // Find all .gguf files in models directory + let mut cleanable: Vec = Vec::new(); + + for entry in fs::read_dir(&models_dir)? { + let entry = entry?; + let path = entry.path(); + + // Skip non-files and non-gguf files + if !path.is_file() { + continue; + } + + let filename = match path.file_name().and_then(|n| n.to_str()) { + Some(name) => name.to_string(), + None => continue, + }; + + // Skip the manifest file itself + if filename == "manifest.toml" { + continue; + } + + // Skip the current default model + if filename == DEFAULT_MODEL { + continue; + } + + // Check if it's a .gguf file + if !filename.ends_with(".gguf") { + continue; + } + + let size = entry.metadata()?.len(); + + let reason = if manifest.has_model(&filename) { + CleanReason::OrphanedModel + } else { + CleanReason::UntrackedFile + }; + + cleanable.push(CleanableModel { filename, path, size, reason }); + } + + if cleanable.is_empty() { + println!("No orphaned models found. Models directory is clean."); + return Ok(()); + } + + // Calculate total size + let total_size: u64 = cleanable.iter().map(|m| m.size).sum(); + let total_mb = total_size as f64 / (1024.0 * 1024.0); + + // Show what will be deleted + println!("Found {} orphaned model(s):", cleanable.len()); + println!(); + + for model in &cleanable { + let size_mb = model.size as f64 / (1024.0 * 1024.0); + println!(" {} ({:.1} MB) - {}", model.filename, size_mb, model.reason_str()); + } + + println!(); + println!("Total: {:.1} MB to be freed", total_mb); + println!(); + + if dry_run { + println!("Dry run: No files were deleted."); + return Ok(()); + } + + // Confirm deletion unless --force + if !force { + println!("Delete these files? [y/N] "); + use std::io::{self, BufRead}; + let stdin = io::stdin(); + let mut line = String::new(); + stdin.lock().read_line(&mut line)?; + + if !line.trim().eq_ignore_ascii_case("y") { + println!("Aborted."); + return Ok(()); + } + } + + // Delete files and update manifest + let mut deleted_count = 0; + let mut deleted_size: u64 = 0; + + for model in &cleanable { + match fs::remove_file(&model.path) { + Ok(()) => { + // Remove from manifest if tracked + manifest.remove_model(&model.filename); + deleted_count += 1; + deleted_size += model.size; + println!("Deleted: {}", model.filename); + } + Err(e) => { + eprintln!("Failed to delete {}: {}", model.filename, e); + } + } + } + + // Save updated manifest + manifest.save()?; + + println!(); + println!( + "Cleaned {} file(s), freed {:.1} MB", + deleted_count, + deleted_size as f64 / (1024.0 * 1024.0) + ); + + Ok(()) +} diff --git a/crates/wake-cli/src/commands/mod.rs b/crates/wake-cli/src/commands/mod.rs index 417cb84..21b3918 100644 --- a/crates/wake-cli/src/commands/mod.rs +++ b/crates/wake-cli/src/commands/mod.rs @@ -8,3 +8,4 @@ pub mod prune; pub mod search; pub mod shell; pub mod status; +pub mod update; diff --git a/crates/wake-cli/src/commands/update.rs b/crates/wake-cli/src/commands/update.rs new file mode 100644 index 0000000..71d7cdb --- /dev/null +++ b/crates/wake-cli/src/commands/update.rs @@ -0,0 +1,319 @@ +//! Self-update command for wake CLI + +use anyhow::{bail, Context, Result}; +use flate2::read::GzDecoder; +use futures_util::StreamExt; +use indicatif::{ProgressBar, ProgressStyle}; +use reqwest::Client; +use serde::Deserialize; +use std::env; +use std::fs::{self, File}; +use std::io::Write; +use std::path::{Path, PathBuf}; +use tar::Archive; + +const GITHUB_RELEASES_URL: &str = "https://api.github.com/repos/joemckenney/wake/releases/latest"; + +#[derive(Debug, Deserialize)] +struct GithubRelease { + tag_name: String, + assets: Vec, +} + +#[derive(Debug, Deserialize)] +struct GithubAsset { + name: String, + browser_download_url: String, + size: u64, +} + +/// Get the current platform target triple +fn get_target_triple() -> Result { + let os = match env::consts::OS { + "linux" => "unknown-linux-gnu", + "macos" => "apple-darwin", + "windows" => "pc-windows-msvc", + other => bail!("Unsupported OS: {}", other), + }; + + let arch = match env::consts::ARCH { + "x86_64" => "x86_64", + "aarch64" => "aarch64", + other => bail!("Unsupported architecture: {}", other), + }; + + Ok(format!("{}-{}", arch, os)) +} + +/// Parse version string, removing 'v' prefix if present +fn parse_version(version: &str) -> &str { + version.strip_prefix('v').unwrap_or(version) +} + +/// Compare two semver versions +/// Returns true if new_version > current_version +fn is_newer_version(current: &str, new: &str) -> bool { + let current = parse_version(current); + let new = parse_version(new); + + let parse = |v: &str| -> Vec { v.split('.').filter_map(|s| s.parse().ok()).collect() }; + + let current_parts = parse(current); + let new_parts = parse(new); + + for i in 0..3 { + let c = current_parts.get(i).copied().unwrap_or(0); + let n = new_parts.get(i).copied().unwrap_or(0); + if n > c { + return true; + } + if n < c { + return false; + } + } + false +} + +/// Fetch the latest release info from GitHub +async fn fetch_latest_release(client: &Client) -> Result { + let response = client + .get(GITHUB_RELEASES_URL) + .header("User-Agent", "wake-cli") + .header("Accept", "application/vnd.github.v3+json") + .send() + .await + .context("Failed to fetch release info")?; + + if !response.status().is_success() { + bail!( + "GitHub API returned error: {} {}", + response.status(), + response.status().canonical_reason().unwrap_or("Unknown") + ); + } + + response.json().await.context("Failed to parse release JSON") +} + +/// Find the asset matching our platform +fn find_platform_asset<'a>(release: &'a GithubRelease, target: &str) -> Option<&'a GithubAsset> { + release.assets.iter().find(|a| a.name.contains(target) && a.name.ends_with(".tar.gz")) +} + +/// Download and extract the update +async fn download_and_extract( + client: &Client, + asset: &GithubAsset, + dest_dir: &Path, +) -> Result { + let response = client + .get(&asset.browser_download_url) + .header("User-Agent", "wake-cli") + .send() + .await + .context("Failed to download asset")?; + + if !response.status().is_success() { + bail!("Download failed: HTTP {}", response.status()); + } + + // Create progress bar + let pb = ProgressBar::new(asset.size); + pb.set_style( + ProgressStyle::default_bar() + .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})") + .unwrap() + .progress_chars("#>-"), + ); + + // Download to temp file + let temp_tar = dest_dir.join("wake-update.tar.gz"); + let mut file = File::create(&temp_tar).context("Failed to create temp file")?; + + let mut downloaded: u64 = 0; + let mut stream = response.bytes_stream(); + + while let Some(chunk) = stream.next().await { + let chunk = chunk.context("Error reading download stream")?; + file.write_all(&chunk)?; + downloaded += chunk.len() as u64; + pb.set_position(downloaded); + } + + pb.finish_with_message("Download complete"); + drop(file); + + // Extract tarball + println!("Extracting..."); + let tar_gz = File::open(&temp_tar)?; + let tar = GzDecoder::new(tar_gz); + let mut archive = Archive::new(tar); + + // Find the wake binary in the archive + let extract_dir = dest_dir.join("wake-extract"); + fs::create_dir_all(&extract_dir)?; + archive.unpack(&extract_dir)?; + + // Find the wake binary (could be in root or in a subdirectory) + let binary_name = if cfg!(windows) { "wake.exe" } else { "wake" }; + let binary_path = find_binary_in_dir(&extract_dir, binary_name)?; + + // Clean up tarball + fs::remove_file(&temp_tar).ok(); + + Ok(binary_path) +} + +/// Recursively find the binary in extracted directory +fn find_binary_in_dir(dir: &PathBuf, name: &str) -> Result { + for entry in fs::read_dir(dir)? { + let entry = entry?; + let path = entry.path(); + if path.is_file() && path.file_name().map(|n| n == name).unwrap_or(false) { + return Ok(path); + } + if path.is_dir() { + if let Ok(found) = find_binary_in_dir(&path, name) { + return Ok(found); + } + } + } + bail!("Could not find {} binary in extracted archive", name) +} + +/// Perform atomic binary replacement +fn replace_binary(new_binary: &PathBuf, current_binary: &PathBuf) -> Result<()> { + // Write to temp location first + let temp_binary = current_binary.with_extension("tmp"); + + // Copy new binary to temp + fs::copy(new_binary, &temp_binary).context("Failed to copy new binary")?; + + // Make executable on Unix + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mut perms = fs::metadata(&temp_binary)?.permissions(); + perms.set_mode(0o755); + fs::set_permissions(&temp_binary, perms)?; + } + + // Atomic rename + fs::rename(&temp_binary, current_binary).context("Failed to replace binary")?; + + Ok(()) +} + +/// Check for updates without installing +pub async fn check() -> Result<()> { + let current_version = env!("CARGO_PKG_VERSION"); + println!("Current version: {}", current_version); + println!("Checking for updates..."); + + let client = Client::builder().user_agent("wake-cli").build()?; + + let release = fetch_latest_release(&client).await?; + let latest_version = parse_version(&release.tag_name); + + if is_newer_version(current_version, latest_version) { + println!(); + println!("New version available: {} -> {}", current_version, latest_version); + println!("Run 'wake update' to install"); + } else { + println!("You're on the latest version ({})", current_version); + } + + Ok(()) +} + +/// Update wake to the latest version +pub async fn run() -> Result<()> { + let current_version = env!("CARGO_PKG_VERSION"); + println!("Current version: {}", current_version); + println!("Checking for updates..."); + + let client = Client::builder().user_agent("wake-cli").build()?; + + let release = fetch_latest_release(&client).await?; + let latest_version = parse_version(&release.tag_name); + + if !is_newer_version(current_version, latest_version) { + println!("You're already on the latest version ({})", current_version); + return Ok(()); + } + + println!(); + println!("New version available: {} -> {}", current_version, latest_version); + + // Get target triple + let target = get_target_triple()?; + println!("Platform: {}", target); + + // Find matching asset + let asset = find_platform_asset(&release, &target) + .ok_or_else(|| anyhow::anyhow!("No release found for platform: {}", target))?; + + println!("Downloading: {}", asset.name); + println!(); + + // Get current binary path + let current_binary = env::current_exe().context("Failed to get current executable path")?; + + // Create temp directory for download + let temp_dir = env::temp_dir().join("wake-update"); + fs::create_dir_all(&temp_dir)?; + + // Download and extract + let new_binary = download_and_extract(&client, asset, &temp_dir).await?; + + // Replace the binary + println!("Installing..."); + replace_binary(&new_binary, ¤t_binary)?; + + // Clean up + fs::remove_dir_all(&temp_dir).ok(); + + println!(); + println!("Successfully updated to version {}", latest_version); + println!(); + println!("Tip: Run 'wake llm clean' to remove any orphaned models from previous versions."); + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_version() { + assert_eq!(parse_version("v1.2.3"), "1.2.3"); + assert_eq!(parse_version("1.2.3"), "1.2.3"); + assert_eq!(parse_version("v0.5.0"), "0.5.0"); + } + + #[test] + fn test_is_newer_version() { + assert!(is_newer_version("0.5.0", "0.6.0")); + assert!(is_newer_version("0.5.0", "1.0.0")); + assert!(is_newer_version("0.5.0", "0.5.1")); + assert!(!is_newer_version("0.6.0", "0.5.0")); + assert!(!is_newer_version("0.5.0", "0.5.0")); + assert!(!is_newer_version("1.0.0", "0.9.9")); + } + + #[test] + fn test_is_newer_version_with_v_prefix() { + assert!(is_newer_version("v0.5.0", "v0.6.0")); + assert!(is_newer_version("0.5.0", "v0.6.0")); + assert!(is_newer_version("v0.5.0", "0.6.0")); + } + + #[test] + fn test_get_target_triple() { + let result = get_target_triple(); + assert!(result.is_ok()); + let triple = result.unwrap(); + assert!(triple.contains("x86_64") || triple.contains("aarch64")); + } +} diff --git a/crates/wake-cli/src/main.rs b/crates/wake-cli/src/main.rs index f728d1d..8bba941 100644 --- a/crates/wake-cli/src/main.rs +++ b/crates/wake-cli/src/main.rs @@ -65,6 +65,12 @@ enum Commands { #[arg(long, value_name = "DAYS")] older_than: Option, }, + /// Update wake to the latest version + Update { + /// Only check for updates, don't install + #[arg(long)] + check: bool, + }, } #[derive(Subcommand)] @@ -73,6 +79,15 @@ enum LlmAction { Download, /// Show model status and path Status, + /// Clean orphaned models from previous versions + Clean { + /// Preview what would be deleted without deleting + #[arg(long)] + dry_run: bool, + /// Skip confirmation prompt + #[arg(short, long)] + force: bool, + }, } #[derive(Subcommand)] @@ -106,6 +121,7 @@ async fn main() -> anyhow::Result<()> { Commands::Llm { action } => match action { LlmAction::Download => commands::llm::download().await, LlmAction::Status => commands::llm::status().await, + LlmAction::Clean { dry_run, force } => commands::llm::clean(dry_run, force).await, }, Commands::Hook { event } => match event { HookEvent::CmdStart { cmd } => commands::hook::cmd_start(&cmd).await, @@ -121,5 +137,12 @@ async fn main() -> anyhow::Result<()> { }) .await } + Commands::Update { check } => { + if check { + commands::update::check().await + } else { + commands::update::run().await + } + } } } diff --git a/crates/wake-core/src/lib.rs b/crates/wake-core/src/lib.rs index 10df5dd..9130728 100644 --- a/crates/wake-core/src/lib.rs +++ b/crates/wake-core/src/lib.rs @@ -1,11 +1,13 @@ pub mod config; pub mod db; pub mod git; +pub mod manifest; pub mod output; pub mod protocol; pub use config::{Config, ConfigError, SummarizationConfig}; pub use db::{Annotation, Command, CommandMetadata, Database, DbError, PruneStats, Session}; pub use git::{GitCache, GitInfo}; +pub use manifest::{ManifestError, ModelEntry, ModelManifest}; pub use output::{OutputBuffer, OutputResult}; pub use protocol::{HookEvent, HookMessage}; diff --git a/crates/wake-core/src/manifest.rs b/crates/wake-core/src/manifest.rs new file mode 100644 index 0000000..7dcaf35 --- /dev/null +++ b/crates/wake-core/src/manifest.rs @@ -0,0 +1,159 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::path::PathBuf; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum ManifestError { + #[error("Could not determine home directory")] + NoHomeDir, + #[error("Failed to read manifest file: {0}")] + ReadError(#[from] std::io::Error), + #[error("Failed to parse manifest file: {0}")] + ParseError(#[from] toml::de::Error), + #[error("Failed to serialize manifest: {0}")] + SerializeError(#[from] toml::ser::Error), +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelEntry { + pub filename: String, + pub wake_version: String, + pub sha256: String, + pub downloaded_at: String, + pub url: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ModelManifest { + #[serde(default)] + pub models: HashMap, +} + +impl ModelManifest { + /// Load manifest from ~/.wake/models/manifest.toml + pub fn load() -> Result { + let path = Self::manifest_path()?; + if !path.exists() { + return Ok(Self::default()); + } + let content = std::fs::read_to_string(&path)?; + let manifest: ModelManifest = toml::from_str(&content)?; + Ok(manifest) + } + + /// Save manifest to ~/.wake/models/manifest.toml + pub fn save(&self) -> Result<(), ManifestError> { + let path = Self::manifest_path()?; + + // Ensure parent directory exists + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + let content = toml::to_string_pretty(self)?; + std::fs::write(&path, content)?; + Ok(()) + } + + /// Get the path to the manifest file + pub fn manifest_path() -> Result { + dirs::home_dir() + .map(|h| h.join(".wake").join("models").join("manifest.toml")) + .ok_or(ManifestError::NoHomeDir) + } + + /// Get the models directory path + pub fn models_dir() -> Result { + dirs::home_dir().map(|h| h.join(".wake").join("models")).ok_or(ManifestError::NoHomeDir) + } + + /// Add or update a model entry + pub fn add_model(&mut self, entry: ModelEntry) { + self.models.insert(entry.filename.clone(), entry); + } + + /// Remove a model entry by filename + pub fn remove_model(&mut self, filename: &str) -> Option { + self.models.remove(filename) + } + + /// Check if a model exists in the manifest + pub fn has_model(&self, filename: &str) -> bool { + self.models.contains_key(filename) + } + + /// Get a model entry by filename + pub fn get_model(&self, filename: &str) -> Option<&ModelEntry> { + self.models.get(filename) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_model_entry_serialization() { + let entry = ModelEntry { + filename: "test-model.gguf".to_string(), + wake_version: "0.5.0".to_string(), + sha256: "abc123".to_string(), + downloaded_at: "2024-01-01T00:00:00Z".to_string(), + url: "https://example.com/model.gguf".to_string(), + }; + + let mut manifest = ModelManifest::default(); + manifest.add_model(entry.clone()); + + let toml_str = toml::to_string_pretty(&manifest).unwrap(); + let parsed: ModelManifest = toml::from_str(&toml_str).unwrap(); + + assert!(parsed.has_model("test-model.gguf")); + let parsed_entry = parsed.get_model("test-model.gguf").unwrap(); + assert_eq!(parsed_entry.filename, entry.filename); + assert_eq!(parsed_entry.wake_version, entry.wake_version); + assert_eq!(parsed_entry.sha256, entry.sha256); + } + + #[test] + fn test_manifest_add_remove() { + let mut manifest = ModelManifest::default(); + + let entry = ModelEntry { + filename: "model.gguf".to_string(), + wake_version: "0.5.0".to_string(), + sha256: "hash".to_string(), + downloaded_at: "2024-01-01T00:00:00Z".to_string(), + url: "https://example.com/model.gguf".to_string(), + }; + + manifest.add_model(entry); + assert!(manifest.has_model("model.gguf")); + + let removed = manifest.remove_model("model.gguf"); + assert!(removed.is_some()); + assert!(!manifest.has_model("model.gguf")); + } + + #[test] + fn test_empty_manifest_parse() { + let toml_str = ""; + let manifest: ModelManifest = toml::from_str(toml_str).unwrap(); + assert!(manifest.models.is_empty()); + } + + #[test] + fn test_manifest_with_models_parse() { + let toml_str = r#" +[models.test-model] +filename = "test-model.gguf" +wake_version = "0.5.0" +sha256 = "abc123" +downloaded_at = "2024-01-01T00:00:00Z" +url = "https://example.com/model.gguf" +"#; + let manifest: ModelManifest = toml::from_str(toml_str).unwrap(); + assert!(manifest.has_model("test-model")); + } +} diff --git a/crates/wake-llm/Cargo.toml b/crates/wake-llm/Cargo.toml index 96cd3e6..696faaf 100644 --- a/crates/wake-llm/Cargo.toml +++ b/crates/wake-llm/Cargo.toml @@ -22,9 +22,12 @@ serde.workspace = true serde_json.workspace = true # Utilities +wake-core.workspace = true dirs.workspace = true thiserror.workspace = true tracing = "0.1" +sha2 = "0.10" +chrono.workspace = true [features] default = [] diff --git a/crates/wake-llm/src/lib.rs b/crates/wake-llm/src/lib.rs index 027ad7c..3eec460 100644 --- a/crates/wake-llm/src/lib.rs +++ b/crates/wake-llm/src/lib.rs @@ -18,10 +18,13 @@ pub use download::DownloadProgress; pub use model::ModelStatus; pub use summarize::SummarizeError; -use std::path::PathBuf; +use sha2::{Digest, Sha256}; +use std::io::Read; +use std::path::{Path, PathBuf}; use std::sync::Arc; use thiserror::Error; use tokio::sync::RwLock; +use wake_core::{ModelEntry, ModelManifest}; /// Default model to use for summarization (Qwen3-0.6B - small and fast) pub const DEFAULT_MODEL: &str = "Qwen3-0.6B-Q4_K_M.gguf"; @@ -41,6 +44,23 @@ pub enum LlmError { Summarize(#[from] summarize::SummarizeError), #[error("IO error: {0}")] Io(#[from] std::io::Error), + #[error("Manifest error: {0}")] + Manifest(#[from] wake_core::ManifestError), +} + +/// Compute SHA256 hash of a file +pub fn compute_sha256(path: &Path) -> std::io::Result { + let mut file = std::fs::File::open(path)?; + let mut hasher = Sha256::new(); + let mut buffer = [0u8; 8192]; + loop { + let bytes_read = file.read(&mut buffer)?; + if bytes_read == 0 { + break; + } + hasher.update(&buffer[..bytes_read]); + } + Ok(format!("{:x}", hasher.finalize())) } /// Main interface for LLM operations @@ -108,6 +128,31 @@ impl WakeLlm { } download::download_model(&self.model_url, &self.model_path, progress_callback).await?; + + // Record in manifest + self.record_in_manifest()?; + + Ok(()) + } + + /// Record the current model in the manifest + fn record_in_manifest(&self) -> Result<(), LlmError> { + let mut manifest = ModelManifest::load()?; + + let sha256 = compute_sha256(&self.model_path)?; + let downloaded_at = chrono::Utc::now().to_rfc3339(); + + let entry = ModelEntry { + filename: DEFAULT_MODEL.to_string(), + wake_version: env!("CARGO_PKG_VERSION").to_string(), + sha256, + downloaded_at, + url: self.model_url.clone(), + }; + + manifest.add_model(entry); + manifest.save()?; + Ok(()) }