diff --git a/crates/desktop/artcraft/src/core/commands/app_preferences/get_app_preferences_command.rs b/crates/desktop/artcraft/src/core/commands/app_preferences/get_app_preferences_command.rs index a42668a668..2777ffa578 100644 --- a/crates/desktop/artcraft/src/core/commands/app_preferences/get_app_preferences_command.rs +++ b/crates/desktop/artcraft/src/core/commands/app_preferences/get_app_preferences_command.rs @@ -1,5 +1,6 @@ use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager; use crate::core::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory; +use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename; use errors::AnyhowResult; use log::{error, info}; use serde_derive::Serialize; @@ -14,6 +15,8 @@ pub struct GetAppPreferencesResponse { pub struct AppPreferencesPayload { /// The downloads directory to use when a user downloads a file. pub preferred_download_directory: PreferredDownloadDirectory, + pub preferred_download_filename: PreferredDownloadFilename, + pub auto_download: bool, /// Play sounds on events. pub play_sounds: bool, @@ -63,6 +66,8 @@ async fn get_prefs(app_prefs: &AppPreferencesManager) -> AnyhowResult { + match request.value { + Some(ValueType::Bool(value)) => prefs.auto_download = value, + _ => return Err(anyhow!("Auto Download must be a boolean")), + } + } + PreferenceName::PreferredDownloadFilename => { + prefs.preferred_download_filename = filename_value(request.value)?; + } PreferenceName::PreferredDownloadDirectory => { match request.value { Some(ValueType::DownloadDirectory(dir)) => @@ -113,8 +126,8 @@ async fn update_prefs( } } - app_prefs.set_clone(&prefs)?; app_data_root.settings_dir().write_app_preferences(&prefs)?; + app_prefs.set_clone(&prefs)?; Ok(()) } @@ -125,3 +138,48 @@ fn string_value(value: &ValueType) -> AnyhowResult { _ => Err(anyhow!("Invalid value type: {:?}", value)), } } + +fn filename_value(value: Option) -> AnyhowResult { + // The untagged value enum decodes string variants as String first. + let filename = match value { + Some(ValueType::String(value)) if value == "artcraft_convention" => PreferredDownloadFilename::ArtcraftConvention, + Some(ValueType::DownloadFilename(filename)) => filename, + _ => return Err(anyhow!("Invalid download filename preference")), + }; + if let PreferredDownloadFilename::Custom(format) = &filename { + PreferredDownloadFilename::validate_custom_format(format).map_err(|err| anyhow!(err))?; + } + Ok(filename) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::state::app_preferences::app_preferences::AppPreferences; + + #[tokio::test] + async fn download_preferences_accept_wire_values_and_survive_reloading() { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path()).unwrap(); + let manager = AppPreferencesManager::with_owned(AppPreferences::default()); + for request in [ + serde_json::json!({ "preference": "preferred_download_directory", "value": {"custom": "/existing/downloads"} }), + serde_json::json!({ "preference": "auto_download", "value": true }), + serde_json::json!({ "preference": "preferred_download_filename", "value": {"custom_format": "my_{model}_{date}"} }), + ] { + update_prefs(serde_json::from_value(request).unwrap(), &manager, &root).await.unwrap(); + } + let prefs = AppPreferences::load_from_file_or_default(&root); + assert!(prefs.auto_download); + assert_eq!(prefs.preferred_download_directory, PreferredDownloadDirectory::Custom("/existing/downloads".into())); + assert_eq!(prefs.preferred_download_filename, PreferredDownloadFilename::Custom("my_{model}_{date}".into())); + + let invalid = serde_json::json!({ "preference": "preferred_download_filename", "value": {"custom_format": "../escape"} }); + assert!(update_prefs(serde_json::from_value(invalid).unwrap(), &manager, &root).await.is_err()); + assert_eq!(AppPreferences::load_from_file_or_default(&root).preferred_download_filename, prefs.preferred_download_filename); + let reset = serde_json::json!({ "preference": "preferred_download_filename", "value": "artcraft_convention" }); + update_prefs(serde_json::from_value(reset).unwrap(), &manager, &root).await.unwrap(); + assert_eq!(AppPreferences::load_from_file_or_default(&root).preferred_download_filename, PreferredDownloadFilename::ArtcraftConvention); + assert!(manager.get_clone().unwrap().auto_download); + } +} diff --git a/crates/desktop/artcraft/src/core/commands/download/download_url_command.rs b/crates/desktop/artcraft/src/core/commands/download/download_url_command.rs index cfcb4c071c..dc26b3963e 100644 --- a/crates/desktop/artcraft/src/core/commands/download/download_url_command.rs +++ b/crates/desktop/artcraft/src/core/commands/download/download_url_command.rs @@ -10,7 +10,8 @@ use crate::core::state::app_preferences::app_preferences::AppPreferences; use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::utils::download_url_to_temp_dir::download_url_to_temp_dir; -use crate::core::utils::download_url_to_user_download_dir::download_url_to_user_download_dir; +use crate::core::utils::download_url_to_user_download_dir::{download_url_with_filename, suggested_download_filename}; +use chrono::Local; use anyhow::anyhow; use log::{error, info}; use serde_derive::{Deserialize, Serialize}; @@ -23,6 +24,13 @@ use url::Url; #[derive(Deserialize, Debug)] pub struct DownloadUrlRequest { pub url: Url, + #[serde(rename = "model")] + pub maybe_model: Option, + #[serde(rename = "batch_index")] + pub maybe_batch_index: Option, + /// Set only when the user explicitly chooses a folder for a batch download. + #[serde(rename = "directory")] + pub maybe_directory: Option, } #[derive(Serialize)] @@ -113,11 +121,9 @@ pub async fn handle_request( let app_prefs = app_prefs.get_clone()?; - let download_path = download_url_to_user_download_dir( - &request.url, - app_data_root, - &app_prefs - ).await?; + let filename = suggested_download_filename(&request.url, &app_prefs, request.maybe_model.as_deref(), request.maybe_batch_index, Local::now()); + let directory = request.maybe_directory.unwrap_or_else(|| app_prefs.preferred_download_directory.download_directory(app_data_root)); + let download_path = download_url_with_filename(&request.url, &directory, &filename).await?; info!("downloaded to: {:?}", download_path); diff --git a/crates/desktop/artcraft/src/core/commands/download/get_download_path_command.rs b/crates/desktop/artcraft/src/core/commands/download/get_download_path_command.rs new file mode 100644 index 0000000000..9a9b5af00d --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/download/get_download_path_command.rs @@ -0,0 +1,46 @@ +use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::utils::download_url_to_user_download_dir::suggested_download_filename; +use chrono::Local; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use tauri::State; +use url::Url; + +#[derive(Deserialize)] +pub struct GetDownloadPathRequest { + pub url: Url, + #[serde(rename = "model")] + pub maybe_model: Option, + #[serde(rename = "batch_index")] + pub maybe_batch_index: Option, +} + +#[derive(Serialize)] +pub struct GetDownloadPathResponse { + pub filename: String, + pub path: PathBuf, +} + +/// Suggest a name for a native save dialog using the same formatter as +/// automatic downloads. The dialog may still let the user override it. +#[tauri::command] +pub fn get_download_path_command( + request: GetDownloadPathRequest, + app_prefs: State<'_, AppPreferencesManager>, + app_data_root: State<'_, AppDataRoot>, +) -> Result { + let prefs = app_prefs.get_clone().map_err(|err| err.to_string())?; + let filename = suggested_download_filename( + &request.url, + &prefs, + request.maybe_model.as_deref(), + request.maybe_batch_index, + Local::now(), + ); + let path = prefs + .preferred_download_directory + .download_directory(&app_data_root) + .join(&filename); + Ok(GetDownloadPathResponse { filename, path }) +} diff --git a/crates/desktop/artcraft/src/core/commands/download/mod.rs b/crates/desktop/artcraft/src/core/commands/download/mod.rs index ad9c6330d3..41523b63b8 100644 --- a/crates/desktop/artcraft/src/core/commands/download/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/download/mod.rs @@ -1,3 +1,4 @@ pub mod download_directory_reveal_command; pub mod download_media_file_command; pub mod download_url_command; +pub mod get_download_path_command; diff --git a/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences.rs b/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences.rs index 8178bfca2e..50b341ce1e 100644 --- a/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences.rs +++ b/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences.rs @@ -1,5 +1,6 @@ use crate::core::state::app_preferences::app_preferences_serializable::AppPreferencesSerializable; use crate::core::state::app_preferences::preferred_download_directory::{PreferredDownloadDirectory, SystemDownloadDirectory}; +use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename; use crate::core::state::data_dir::app_data_root::AppDataRoot; #[derive(Clone)] @@ -7,6 +8,11 @@ pub struct AppPreferences { /// The downloads directory to use when a user downloads a file. pub preferred_download_directory: PreferredDownloadDirectory, + pub preferred_download_filename: PreferredDownloadFilename, + + /// Save generated results before marking their task complete. + pub auto_download: bool, + /// Play sounds on events. pub play_sounds: bool, @@ -36,6 +42,8 @@ impl Default for AppPreferences { fn default() -> Self { Self { preferred_download_directory: PreferredDownloadDirectory::System(SystemDownloadDirectory::Downloads), + preferred_download_filename: PreferredDownloadFilename::default(), + auto_download: false, play_sounds: true, // NB: These are defined in the frontend. enqueue_success_sound: Some("done".to_string()), diff --git a/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences_serializable.rs b/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences_serializable.rs index e795d44cc2..2963de0fbd 100644 --- a/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences_serializable.rs +++ b/crates/desktop/artcraft/src/core/state/app_preferences/app_preferences_serializable.rs @@ -1,5 +1,6 @@ use crate::core::state::app_preferences::app_preferences::AppPreferences; use crate::core::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory; +use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename; use crate::core::state::data_dir::app_data_root::AppDataRoot; use errors::AnyhowResult; use serde_derive::{Deserialize, Serialize}; @@ -7,7 +8,8 @@ use serde_derive::{Deserialize, Serialize}; /// Vector clock versioning string rather than semver. /// - Version 1 - initial version. /// - Version 2 - added "delete_file_sound", marked optionals "skip_serializing_if" -const CURRENT_VERSION: &str = "2"; +/// - Version 3 - added auto download and filename convention. +const CURRENT_VERSION: &str = "3"; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AppPreferencesSerializable { @@ -17,6 +19,12 @@ pub struct AppPreferencesSerializable { /// The downloads directory to use when a user downloads a file. pub preferred_download_directory: Option, + #[serde(rename = "preferred_download_filename", skip_serializing_if = "Option::is_none")] + pub maybe_preferred_download_filename: Option, + + #[serde(rename = "auto_download", skip_serializing_if = "Option::is_none")] + pub maybe_auto_download: Option, + /// Play sounds on events. #[serde(skip_serializing_if = "Option::is_none")] pub play_sounds: Option, @@ -54,6 +62,8 @@ impl AppPreferencesSerializable { Self { version: CURRENT_VERSION.to_string(), preferred_download_directory: Some(preferences.preferred_download_directory.clone()), + maybe_preferred_download_filename: Some(preferences.preferred_download_filename.clone()), + maybe_auto_download: Some(preferences.auto_download), play_sounds: Some(preferences.play_sounds), delete_file_sound: preferences.delete_file_sound.clone(), generation_success_sound: preferences.generation_success_sound.clone(), @@ -73,6 +83,13 @@ impl AppPreferencesSerializable { preferences.play_sounds = play_sounds; } + if let Some(filename) = &self.maybe_preferred_download_filename { + preferences.preferred_download_filename = filename.clone(); + } + if let Some(auto_download) = self.maybe_auto_download { + preferences.auto_download = auto_download; + } + preferences.delete_file_sound = self.delete_file_sound.clone(); preferences.generation_success_sound = self.generation_success_sound.clone(); preferences.generation_failure_sound = self.generation_failure_sound.clone(); @@ -81,3 +98,40 @@ impl AppPreferencesSerializable { preferences } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn old_preferences_keep_the_directory_and_default_auto_download_off() { + let old: AppPreferencesSerializable = serde_json::from_value(serde_json::json!({ + "version": "2", "preferred_download_directory": { "custom": "/my/downloads" }, + "play_sounds": false, "generation_success_sound": "special_flower", + })).unwrap(); + let prefs = old.to_preferences(); + assert!(!prefs.auto_download); + assert_eq!(prefs.preferred_download_filename, PreferredDownloadFilename::ArtcraftConvention); + assert_eq!(prefs.preferred_download_directory, PreferredDownloadDirectory::Custom("/my/downloads".into())); + assert!(!prefs.play_sounds); + assert_eq!(prefs.generation_success_sound.as_deref(), Some("special_flower")); + } + + #[test] + fn download_preferences_round_trip_without_changing_existing_fields() { + let prefs = AppPreferences { + auto_download: true, + preferred_download_filename: PreferredDownloadFilename::Custom("{model}_{date}".into()), + preferred_download_directory: PreferredDownloadDirectory::Custom("/my/downloads".into()), + ..AppPreferences::default() + }; + let json = serde_json::to_value(prefs.to_serializable()).unwrap(); + assert_eq!(json["auto_download"], true); + assert_eq!(json["preferred_download_filename"]["custom_format"], "{model}_{date}"); + let restored = serde_json::from_value::(json).unwrap().to_preferences(); + assert!(restored.auto_download); + assert_eq!(restored.preferred_download_filename, prefs.preferred_download_filename); + assert_eq!(restored.preferred_download_directory, prefs.preferred_download_directory); + assert_eq!(restored.generation_success_sound, prefs.generation_success_sound); + } +} diff --git a/crates/desktop/artcraft/src/core/state/app_preferences/mod.rs b/crates/desktop/artcraft/src/core/state/app_preferences/mod.rs index d2385da909..0340633aeb 100644 --- a/crates/desktop/artcraft/src/core/state/app_preferences/mod.rs +++ b/crates/desktop/artcraft/src/core/state/app_preferences/mod.rs @@ -1,4 +1,5 @@ pub mod app_preferences; pub mod app_preferences_manager; pub mod app_preferences_serializable; -pub mod preferred_download_directory; \ No newline at end of file +pub mod preferred_download_directory; +pub mod preferred_download_filename; diff --git a/crates/desktop/artcraft/src/core/state/app_preferences/preferred_download_filename.rs b/crates/desktop/artcraft/src/core/state/app_preferences/preferred_download_filename.rs new file mode 100644 index 0000000000..9cc0e736b5 --- /dev/null +++ b/crates/desktop/artcraft/src/core/state/app_preferences/preferred_download_filename.rs @@ -0,0 +1,351 @@ +use chrono::{DateTime, Datelike, Local, Timelike}; +use serde_derive::{Deserialize, Serialize}; + +/// Characters that may never appear in a download filename (or a custom +/// filename format): path separators, traversal, quoting, and shell-unsafe +/// characters. +const UNSAFE_CHARACTERS: &[char] = &['/', '\\', '\'', '"', '`', '%', '<', '>', '|', ':', '*', '?', '\0']; + +/// How downloaded generation files are named on disk. +/// +/// Stored in the app preferences file. NEVER change existing serialized +/// values; only add new ones. +#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum PreferredDownloadFilename { + /// `{model}_{date}(_{batch_index}).{ext}` + /// - model slug, eg. "seedance_2.0", "flux_pro_1.1" + /// - date is "YYYY-MM-DD-HH-MM-SS" (download time) + /// - batch index is 1-n when the job produced more than one file + /// - ext matches the cloud file (jpg, png, mp4, ...) + /// + /// NB: Serializes as `"artcraft_convention"` in JSON. + ArtcraftConvention, + + /// A user-defined format string. Supported tokens: + /// - `{model}` — the model slug + /// - `{date}` — "YYYY-MM-DD-HH-MM-SS" (download time) + /// - `{YYYY}`, `{YY}`, `{MM}`, `{DD}`, `{HH}`, `{mm}`, `{SS}` — date parts + /// - `{batch_index}` — 1-n (appended automatically when the job produced + /// more than one file and the format doesn't include it) + /// + /// Slash, quote, and other unsafe characters are rejected. + /// + /// NB: Serializes as `{"custom_format": "..."}` in JSON. (The key is NOT + /// "custom" so it can't be confused with `PreferredDownloadDirectory`'s + /// custom variant in loosely-typed payloads.) + #[serde(rename = "custom_format")] + Custom(String), +} + +impl Default for PreferredDownloadFilename { + fn default() -> Self { + Self::ArtcraftConvention + } +} + +/// Everything needed to render a download filename. +pub struct DownloadFilenameParts<'a> { + /// Model slug, eg. "seedance_2.0". Use [`model_slug_from_model_type_str`]. + pub model_slug: &'a str, + + /// The download time. + pub download_time: DateTime, + + /// 1-based index when the job produced more than one file. + pub maybe_batch_index: Option, + + /// File extension without the leading dot, matching the cloud file. + pub extension: &'a str, +} + +impl PreferredDownloadFilename { + /// Render the filename (including extension) for one downloaded file. + /// The output is sanitized: unsafe characters can never appear. + pub fn build_filename(&self, parts: &DownloadFilenameParts<'_>) -> String { + let stem = match self { + Self::ArtcraftConvention => { + let mut stem = format!("{}_{}", parts.model_slug, full_date_string(&parts.download_time),); + if let Some(index) = parts.maybe_batch_index { + stem.push_str(&format!("_{index}")); + } + stem + }, + Self::Custom(format) => { + let mut stem = expand_format_tokens(format, parts); + let has_batch_token = format.contains("{batch_index}"); + if !has_batch_token { + if let Some(index) = parts.maybe_batch_index { + stem.push_str(&format!("_{index}")); + } + } + stem + }, + }; + + // Leave room for an extension and a collision suffix on common filesystems. + let mut stem = sanitize(&stem); + while stem.len() > 200 { + stem.pop(); + } + if stem.is_empty() { + stem = "artcraft".to_string(); + } + let device_name = stem.split('.').next().unwrap_or_default().to_ascii_uppercase(); + if matches!(device_name.as_str(), "CON" | "PRN" | "AUX" | "NUL") + || (device_name.len() == 4 + && (device_name.starts_with("COM") || device_name.starts_with("LPT")) + && matches!(device_name.as_bytes()[3], b'1'..=b'9')) + { + stem.insert_str(0, "artcraft_"); + } + let extension = sanitize(parts.extension); + let extension = if extension.is_empty() { "bin" } else { &extension }; + + format!("{stem}.{extension}") + } + + /// Vet a user-supplied custom format string. + pub fn validate_custom_format(format: &str) -> Result<(), String> { + if format.trim().is_empty() { + return Err("Filename format cannot be empty".to_string()); + } + if format.contains("..") { + return Err("Filename format cannot contain '..'".to_string()); + } + if let Some(bad) = format.chars().find(|c| UNSAFE_CHARACTERS.contains(c) || c.is_control()) { + return Err(format!("Filename format cannot contain {bad:?}")); + } + Ok(()) + } +} + +/// Derive a friendly model slug from a `TaskModelType` string value: +/// version markers like "2p0" become "2.0" (eg. "seedance_2p0" -> +/// "seedance_2.0", "flux_pro_1p1" -> "flux_pro_1.1"). +pub fn model_slug_from_model_type_str(model_type: &str) -> String { + let mut slug = String::with_capacity(model_type.len()); + let chars = model_type.chars().collect::>(); + for (i, c) in chars.iter().enumerate() { + let is_version_dot = + *c == 'p' && i > 0 && i + 1 < chars.len() && chars[i - 1].is_ascii_digit() && chars[i + 1].is_ascii_digit(); + slug.push(if is_version_dot { '.' } else { *c }); + } + slug +} + +fn full_date_string(time: &DateTime) -> String { + format!( + "{:04}-{:02}-{:02}-{:02}-{:02}-{:02}", + time.year(), + time.month(), + time.day(), + time.hour(), + time.minute(), + time.second(), + ) +} + +fn expand_format_tokens(format: &str, parts: &DownloadFilenameParts<'_>) -> String { + let time = &parts.download_time; + format + .replace("{model}", parts.model_slug) + .replace("{date}", &full_date_string(time)) + .replace("{YYYY}", &format!("{:04}", time.year())) + .replace("{YY}", &format!("{:02}", time.year() % 100)) + .replace("{MM}", &format!("{:02}", time.month())) + .replace("{DD}", &format!("{:02}", time.day())) + .replace("{HH}", &format!("{:02}", time.hour())) + .replace("{mm}", &format!("{:02}", time.minute())) + .replace("{SS}", &format!("{:02}", time.second())) + .replace( + "{batch_index}", + &parts.maybe_batch_index.map(|i| i.to_string()).unwrap_or_default(), + ) +} + +/// Strip anything unsafe out of a rendered filename component. +fn sanitize(value: &str) -> String { + value + .chars() + .filter(|c| !UNSAFE_CHARACTERS.contains(c) && !c.is_control()) + .collect::() + .replace("..", "_") + .trim_matches(|c: char| c.is_whitespace() || c == '.') + .to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + + mod artcraft_convention { + use super::*; + + #[test] + fn single_file() { + let name = PreferredDownloadFilename::ArtcraftConvention.build_filename(&DownloadFilenameParts { + model_slug: "seedance_2.0", + download_time: test_time(), + maybe_batch_index: None, + extension: "mp4", + }); + assert_eq!(name, "seedance_2.0_2026-08-05-14-30-09.mp4"); + } + + #[test] + fn batch_file() { + let name = PreferredDownloadFilename::ArtcraftConvention.build_filename(&DownloadFilenameParts { + model_slug: "flux_pro_1.1", + download_time: test_time(), + maybe_batch_index: Some(3), + extension: "png", + }); + assert_eq!(name, "flux_pro_1.1_2026-08-05-14-30-09_3.png"); + } + } + + mod custom_format { + use super::*; + + #[test] + fn tokens_expand() { + let format = PreferredDownloadFilename::Custom("{model}-{YYYY}{MM}{DD}-{HH}{mm}{SS}".to_string()); + let name = format.build_filename(&DownloadFilenameParts { + model_slug: "nano_banana", + download_time: test_time(), + maybe_batch_index: None, + extension: "jpg", + }); + assert_eq!(name, "nano_banana-20260805-143009.jpg"); + } + + #[test] + fn batch_index_token() { + let format = PreferredDownloadFilename::Custom("gen_{batch_index}".to_string()); + let name = format.build_filename(&DownloadFilenameParts { + model_slug: "unused", + download_time: test_time(), + maybe_batch_index: Some(2), + extension: "png", + }); + assert_eq!(name, "gen_2.png"); + } + + #[test] + fn batch_index_appended_when_absent_from_format() { + let format = PreferredDownloadFilename::Custom("gen".to_string()); + let name = format.build_filename(&DownloadFilenameParts { + model_slug: "unused", + download_time: test_time(), + maybe_batch_index: Some(2), + extension: "png", + }); + assert_eq!(name, "gen_2.png"); + } + + #[test] + fn empty_expansions_reserved_names_and_long_unicode_names_are_portable() { + let parts = DownloadFilenameParts { + model_slug: "midjourney_8", + download_time: test_time(), + maybe_batch_index: None, + extension: "png", + }; + assert_eq!( + PreferredDownloadFilename::Custom("{batch_index}".into()).build_filename(&parts), + "artcraft.png" + ); + assert_eq!( + PreferredDownloadFilename::Custom("CON".into()).build_filename(&parts), + "artcraft_CON.png" + ); + let name = PreferredDownloadFilename::Custom("画".repeat(150)).build_filename(&parts); + assert!(name.len() < 240); + assert!(name.ends_with(".png")); + } + + #[test] + fn unsafe_characters_stripped_from_output() { + let format = PreferredDownloadFilename::Custom("weird{model}".to_string()); + let name = format.build_filename(&DownloadFilenameParts { + model_slug: "a/b\\c\"d", + download_time: test_time(), + maybe_batch_index: None, + extension: "png", + }); + assert_eq!(name, "weirdabcd.png"); + } + } + + mod validation { + use super::*; + + #[test] + fn accepts_reasonable_formats() { + assert!(PreferredDownloadFilename::validate_custom_format("{model}_{date}").is_ok()); + assert!(PreferredDownloadFilename::validate_custom_format("my file {YYYY}").is_ok()); + } + + #[test] + fn rejects_unsafe_formats() { + assert!(PreferredDownloadFilename::validate_custom_format("").is_err()); + assert!(PreferredDownloadFilename::validate_custom_format("a/b").is_err()); + assert!(PreferredDownloadFilename::validate_custom_format("a\\b").is_err()); + assert!(PreferredDownloadFilename::validate_custom_format("a\"b").is_err()); + assert!(PreferredDownloadFilename::validate_custom_format("a'b").is_err()); + assert!(PreferredDownloadFilename::validate_custom_format("../etc").is_err()); + } + } + + mod model_slug { + use super::*; + + #[test] + fn version_markers_become_dots() { + assert_eq!(model_slug_from_model_type_str("seedance_2p0"), "seedance_2.0"); + assert_eq!(model_slug_from_model_type_str("flux_pro_1p1"), "flux_pro_1.1"); + assert_eq!(model_slug_from_model_type_str("seedream_5p0_pro_u"), "seedream_5.0_pro_u"); + } + + #[test] + fn non_version_names_unchanged() { + assert_eq!(model_slug_from_model_type_str("nano_banana"), "nano_banana"); + assert_eq!(model_slug_from_model_type_str("midjourney_8"), "midjourney_8"); + assert_eq!(model_slug_from_model_type_str("marble_1p1_plus"), "marble_1.1_plus"); + } + } + + mod serialization { + use super::*; + + #[test] + fn artcraft_convention_serializes_as_string() { + let val = PreferredDownloadFilename::ArtcraftConvention; + assert_eq!(serde_json::to_string(&val).unwrap(), "\"artcraft_convention\""); + } + + #[test] + fn custom_serializes_with_custom_format_key() { + let val = PreferredDownloadFilename::Custom("{model}_{date}".to_string()); + assert_eq!(serde_json::to_string(&val).unwrap(), "{\"custom_format\":\"{model}_{date}\"}"); + } + + #[test] + fn round_trips() { + for val in [ + PreferredDownloadFilename::ArtcraftConvention, + PreferredDownloadFilename::Custom("{model}".to_string()), + ] { + let json = serde_json::to_string(&val).unwrap(); + let back: PreferredDownloadFilename = serde_json::from_str(&json).unwrap(); + assert_eq!(back, val); + } + } + } + + fn test_time() -> DateTime { + Local.with_ymd_and_hms(2026, 8, 5, 14, 30, 9).unwrap() + } +} diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs index de3ddd51fe..620d19aec0 100644 --- a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs @@ -5,6 +5,7 @@ use crate::core::threads::third_party_task_polling_thread::events::notify_fronte use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::data_dir::trait_data_subdir::DataSubdir; use crate::core::state::task_database::TaskDatabase; +use crate::core::utils::auto_download::{auto_download_task_urls, clear_auto_download_checkpoint}; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet; @@ -89,6 +90,10 @@ async fn handle_fal_complete_inner( return Ok(()); } + // Save local copies before uploading or announcing completion. + let download_urls = urls.iter().map(|url| reqwest::Url::parse(url)).collect::, _>>()?; + auto_download_task_urls(app_handle, task, &download_urls).await?; + // Use the prompt token from the task record (created at generation time). let maybe_prompt_token = task.prompt_token.as_ref() .map(|s| PromptToken::new_from_str(s)); @@ -163,6 +168,7 @@ async fn handle_fal_complete_inner( }).await?; if updated { + clear_auto_download_checkpoint(app_handle, task); if let Some(primary_token) = maybe_primary_media_file_token { let completion = CompletionData { primary_media_file_token: primary_token, @@ -341,4 +347,3 @@ async fn try_upload( Ok(media_token) } - diff --git a/crates/desktop/artcraft/src/core/utils/auto_download.rs b/crates/desktop/artcraft/src/core/utils/auto_download.rs new file mode 100644 index 0000000000..3b4e2f3e31 --- /dev/null +++ b/crates/desktop/artcraft/src/core/utils/auto_download.rs @@ -0,0 +1,309 @@ +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::app_preferences::app_preferences::AppPreferences; +use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; +use crate::core::utils::download_url_to_user_download_dir::{download_url_with_filename, suggested_download_filename}; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use anyhow::{anyhow, ensure}; +use artcraft_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; +use chrono::{DateTime, Local}; +use errors::AnyhowResult; +use serde::{Deserialize, Serialize}; +use sqlite_tasks::queries::task::Task; +use std::collections::BTreeMap; +use std::io::Write; +use std::path::PathBuf; +use std::time::Duration; +use tauri::{AppHandle, Manager}; +use tempfile::NamedTempFile; +use tokens::tokens::batch_generations::BatchGenerationToken; +use url::Url; + +#[derive(Serialize, Deserialize)] +struct AutoDownloadCheckpoint { + download_time: DateTime, + saved: BTreeMap, +} + +/// Await this before updating task status or emitting ANY completion events. +/// A failed download leaves the task pending for the next polling pass. +pub async fn auto_download_task( + app: &AppHandle, + task: &Task, + maybe_batch_token: Option<&BatchGenerationToken>, + maybe_primary_url: Option<&str>, + maybe_expected_count: Option, +) -> AnyhowResult<()> { + let prefs = app.state::().get_clone()?; + if !prefs.auto_download { + return Ok(()); + } + let urls = if let Some(batch) = maybe_batch_token { + let config = app.state::(); + let credentials = app.state::().get_credentials()?; + let result = tokio::time::timeout( + Duration::from_secs(60), + list_batch_generated_redux_media_files(&config.storyteller_host, credentials.as_ref(), batch), + ) + .await??; + ensure!(!result.media_files.is_empty(), "No downloadable results for batch {}", batch); + // Stable order across retries, even if the backend reorders the response. + let mut files = result.media_files; + files.sort_by(|a, b| a.token.as_str().cmp(b.token.as_str())); + files.into_iter().map(|file| file.media_links.cdn_url).collect::>() + } else { + vec![Url::parse( + maybe_primary_url.ok_or_else(|| anyhow!("Completed task has no download URL"))?, + )?] + }; + if let Some(expected) = maybe_expected_count { + ensure!( + urls.len() == expected, + "Waiting for all {} download results (currently {})", + expected, + urls.len() + ); + } + download_generated_results(&app.state::(), &prefs, task, &urls).await +} + +/// For providers that already expose all result URLs before upload. +pub async fn auto_download_task_urls(app: &AppHandle, task: &Task, urls: &[Url]) -> AnyhowResult<()> { + let prefs = app.state::().get_clone()?; + download_generated_results(&app.state::(), &prefs, task, urls).await +} + +pub fn clear_auto_download_checkpoint(app: &AppHandle, task: &Task) { + let path = checkpoint_path(&app.state::(), task); + if let Err(err) = std::fs::remove_file(path) { + if err.kind() != std::io::ErrorKind::NotFound { + log::warn!("Could not remove automatic download checkpoint: {}", err); + } + } +} + +async fn download_generated_results( + root: &AppDataRoot, + prefs: &AppPreferences, + task: &Task, + urls: &[Url], +) -> AnyhowResult<()> { + if !prefs.auto_download { + return Ok(()); + } + ensure!(!urls.is_empty(), "Completed task has no downloadable results"); + let path = checkpoint_path(root, task); + let mut checkpoint: AutoDownloadCheckpoint = match std::fs::read(&path) { + Ok(bytes) => serde_json::from_slice(&bytes)?, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + AutoDownloadCheckpoint { download_time: Local::now(), saved: BTreeMap::new() } + }, + Err(err) => return Err(err.into()), + }; + let directory = prefs.preferred_download_directory.download_directory(root); + for (index, url) in urls.iter().enumerate() { + if checkpoint.saved.get(url.as_str()).is_some_and(|path| path.is_file()) { + continue; + } + let filename = suggested_download_filename( + url, + prefs, + task.model_type.as_ref().map(|model| model.to_str()), + (urls.len() > 1).then_some(index + 1), + checkpoint.download_time, + ); + let saved = download_url_with_filename(url, &directory, &filename).await?; + checkpoint.saved.insert(url.to_string(), saved); + let mut file = NamedTempFile::new_in(root.state_dir().path())?; + file.write_all(&serde_json::to_vec(&checkpoint)?)?; + file.as_file().sync_all()?; + file.persist(&path)?; + } + Ok(()) +} + +fn checkpoint_path(root: &AppDataRoot, task: &Task) -> PathBuf { + root.state_dir().path().join(format!("auto_download_{}.json", task.id.as_str())) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory; + use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename; + use enums::common::generation_provider::GenerationProvider; + use enums::tauri::tasks::{task_model_type::TaskModelType, task_status::TaskStatus, task_type::TaskType}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use tokio::sync::oneshot; + use tokens::tokens::sqlite::tasks::TaskId; + + #[tokio::test] + async fn disabled_auto_download_does_not_access_network_or_create_files() { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path()).unwrap(); + let mut prefs = preferences(&root); + prefs.auto_download = false; + download_generated_results( + &root, + &prefs, + &task(GenerationProvider::Artcraft), + &[Url::parse("http://127.0.0.1:1/never.png").unwrap()], + ) + .await + .unwrap(); + assert!(!dir.path().join("chosen").exists()); + } + + #[tokio::test] + async fn all_results_are_saved_for_artcraft_and_midjourney_and_retries_skip_saved_files() { + for provider in [GenerationProvider::Artcraft, GenerationProvider::Midjourney] { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path()).unwrap(); + let prefs = preferences(&root); + let task = task(provider); + let (base, server) = server(vec![(200, "image"), (200, "video"), (200, "audio"), (200, "splat")]).await; + let urls = ["png", "mp4", "mp3", "spz"].map(|ext| Url::parse(&format!("{base}/result.{ext}")).unwrap()); + download_generated_results(&root, &prefs, &task, &urls).await.unwrap(); + server.await.unwrap(); + for (index, (ext, bytes)) in [("png", "image"), ("mp4", "video"), ("mp3", "audio"), ("spz", "splat")] + .iter() + .enumerate() + { + assert_eq!( + std::fs::read_to_string(dir.path().join(format!("chosen/midjourney_8_result_{}.{ext}", index + 1))).unwrap(), + *bytes + ); + } + // The server is gone; resuming from disk must not redownload anything. + download_generated_results(&root, &prefs, &task, &urls).await.unwrap(); + assert_eq!(std::fs::read_dir(dir.path().join("chosen")).unwrap().count(), 4); + } + } + + #[tokio::test] + async fn partial_failure_keeps_finished_files_and_retries_only_the_failed_result() { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path()).unwrap(); + let prefs = preferences(&root); + let task = task(GenerationProvider::Midjourney); + let (base, server) = server(vec![(200, "first"), (500, "error"), (200, "second")]).await; + let urls = ["first", "second"].map(|name| Url::parse(&format!("{base}/{name}.png")).unwrap()); + assert!(download_generated_results(&root, &prefs, &task, &urls).await.is_err()); + assert_eq!(std::fs::read_dir(dir.path().join("chosen")).unwrap().count(), 1); + download_generated_results(&root, &prefs, &task, &urls).await.unwrap(); + server.await.unwrap(); + assert_eq!( + std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result_1.png")).unwrap(), + "first" + ); + assert_eq!( + std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result_2.png")).unwrap(), + "second" + ); + } + + #[tokio::test] + async fn matching_custom_filenames_never_overwrite_previous_generations() { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path()).unwrap(); + let prefs = preferences(&root); + let (base, server) = server(vec![(200, "first"), (200, "second")]).await; + let url = Url::parse(&format!("{base}/file.png")).unwrap(); + for _ in 0..2 { + download_generated_results(&root, &prefs, &task(GenerationProvider::Artcraft), &[url.clone()]) + .await + .unwrap(); + } + server.await.unwrap(); + assert_eq!( + std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result.png")).unwrap(), + "first" + ); + assert_eq!( + std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result_1.png")).unwrap(), + "second" + ); + } + + #[tokio::test] + async fn completion_waits_until_the_entire_file_is_on_disk() { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path()).unwrap(); + let prefs = preferences(&root); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = Url::parse(&format!("http://{}/file.png", listener.local_addr().unwrap())).unwrap(); + let (started_tx, started_rx) = oneshot::channel(); + let (finish_tx, finish_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + socket.read(&mut request).await.unwrap(); + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 6\r\nConnection: close\r\n\r\nabc") + .await + .unwrap(); + started_tx.send(()).unwrap(); + finish_rx.await.unwrap(); + socket.write_all(b"def").await.unwrap(); + }); + let completion = tokio::spawn(async move { + download_generated_results(&root, &prefs, &task(GenerationProvider::Artcraft), &[url]) + .await + .unwrap(); + // Callers may mark success and emit notifications only after this await. + std::fs::read_to_string(root.path().join("chosen/midjourney_8_result.png")).unwrap() + }); + started_rx.await.unwrap(); + assert!(!completion.is_finished()); + assert!(!dir.path().join("chosen/midjourney_8_result.png").exists()); + finish_tx.send(()).unwrap(); + assert_eq!(completion.await.unwrap(), "abcdef"); + server.await.unwrap(); + } + + fn preferences(root: &AppDataRoot) -> AppPreferences { + AppPreferences { + auto_download: true, + preferred_download_directory: PreferredDownloadDirectory::Custom(root.path().join("chosen")), + preferred_download_filename: PreferredDownloadFilename::Custom("{model}_result".into()), + ..AppPreferences::default() + } + } + + fn task(provider: GenerationProvider) -> Task { + Task { + id: TaskId::generate(), + status: TaskStatus::Pending, + task_type: TaskType::ImageGeneration, + model_type: Some(TaskModelType::Midjourney8), + provider, + provider_job_id: None, + queue_status_url: None, + queue_response_url: None, + prompt_token: None, + frontend_caller: None, + frontend_subscriber_id: None, + frontend_subscriber_payload: None, + } + } + + async fn server(responses: Vec<(u16, &'static str)>) -> (String, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + for (status, body) in responses { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + socket.read(&mut request).await.unwrap(); + let response = format!( + "HTTP/1.1 {status} Test\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + socket.write_all(response.as_bytes()).await.unwrap(); + } + }); + (base, server) + } +} diff --git a/crates/desktop/artcraft/src/core/utils/download_url_to_user_download_dir.rs b/crates/desktop/artcraft/src/core/utils/download_url_to_user_download_dir.rs index 7196ef45ae..086fcdd2f3 100644 --- a/crates/desktop/artcraft/src/core/utils/download_url_to_user_download_dir.rs +++ b/crates/desktop/artcraft/src/core/utils/download_url_to_user_download_dir.rs @@ -1,12 +1,15 @@ use crate::core::artcraft_error::ArtcraftError; use crate::core::state::app_preferences::app_preferences::AppPreferences; +use crate::core::state::app_preferences::preferred_download_filename::{model_slug_from_model_type_str, DownloadFilenameParts, PreferredDownloadFilename}; use crate::core::state::data_dir::app_data_root::AppDataRoot; use anyhow::anyhow; use errors::AnyhowResult; use log::{error, info}; -use std::fs::OpenOptions; -use std::io::Write; -use std::path::PathBuf; +use chrono::{DateTime, Local}; +use std::path::{Path, PathBuf}; +use std::time::Duration; +use tempfile::NamedTempFile; +use tokio::io::AsyncWriteExt; use url::Url; pub async fn download_url_to_user_download_dir( @@ -14,67 +17,64 @@ pub async fn download_url_to_user_download_dir( app_data_root: &AppDataRoot, app_prefs: &AppPreferences, ) -> Result { - - let url_file_name = { - let path_segments = url.path_segments() - .ok_or_else(|| ArtcraftError::AnyhowError(anyhow!("URL does not have path segments (1): {:?}", url)))?; - - let mut url_file_name = path_segments.last() - .ok_or_else(|| ArtcraftError::AnyhowError(anyhow!("URL does not have path segments (2): {:?}", url)))? - .trim() - .to_string(); - - if !url_file_name.to_lowercase().starts_with("artcraft") { - url_file_name = format!("artcraft_{}", url_file_name); - } - - url_file_name - }; - - check_url_file_name_for_downloadability(&url_file_name)?; - + let filename = suggested_download_filename(url, app_prefs, None, None, Local::now()); let download_directory = app_prefs .preferred_download_directory .download_directory(app_data_root); + download_url_with_filename(url, &download_directory, &filename).await +} - info!("Pushing path `{}` onto downloaded directory `{:?}`", url_file_name, download_directory); +pub fn suggested_download_filename( + url: &Url, + prefs: &AppPreferences, + maybe_model: Option<&str>, + maybe_batch_index: Option, + download_time: DateTime, +) -> String { + let extension = Path::new(url.path()).extension().and_then(|ext| ext.to_str()).unwrap_or("bin"); + let model_slug = model_slug_from_model_type_str(maybe_model.unwrap_or("artcraft")); + prefs.preferred_download_filename.build_filename(&DownloadFilenameParts { + model_slug: &model_slug, + download_time, + maybe_batch_index, + extension, + }) +} - let download_filename = { - let mut download_filename = download_directory.clone(); - download_filename.push(&url_file_name); - download_filename - }; - - info!("Download filename: `{:?}`", download_filename); - - if download_filename == download_directory { - return Err(ArtcraftError::AnyhowError(anyhow!("Download filename resolved to directory: {:?}", download_filename)))?; +/// Stream to a temporary file in the destination directory, then publish the +/// complete file atomically. Never overwrite another download, even if two +/// generations finish in the same second or use the same custom name. +pub async fn download_url_with_filename(url: &Url, directory: &Path, filename: &str) -> Result { + check_url_file_name_for_downloadability(filename)?; + tokio::fs::create_dir_all(directory).await?; + let temporary = NamedTempFile::new_in(directory)?; + let mut file = tokio::fs::File::from_std(temporary.reopen()?); + let client = reqwest::Client::builder().connect_timeout(Duration::from_secs(30)).timeout(Duration::from_secs(600)).build()?; + let mut response = client.get(url.clone()).send().await?.error_for_status()?; + while let Some(chunk) = response.chunk().await? { + file.write_all(&chunk).await?; } + file.flush().await?; + file.sync_all().await?; + drop(file); + let path = persist_download(temporary, directory, filename)?; + info!("Downloaded generation to {:?}", path); + Ok(path) +} - if download_filename.exists() { - if download_filename.is_dir() { - return Err(ArtcraftError::AnyhowError(anyhow!("Download path exists and resolved to directory: {:?}", download_filename))); - } else { - return Err(ArtcraftError::CannotDownloadFilePathAlreadyExists { path: download_filename }.into()); +fn persist_download(mut file: NamedTempFile, directory: &Path, filename: &str) -> Result { + let stem = Path::new(filename).file_stem().and_then(|stem| stem.to_str()).unwrap_or("artcraft"); + let extension = Path::new(filename).extension().and_then(|ext| ext.to_str()).unwrap_or("bin"); + for suffix in 0..10_000 { + let name = if suffix == 0 { filename.to_string() } else { format!("{stem}_{suffix}.{extension}") }; + let path = directory.join(name); + match file.persist_noclobber(&path) { + Ok(_) => return Ok(path), + Err(err) if err.error.kind() == std::io::ErrorKind::AlreadyExists => file = err.file, + Err(err) => return Err(err.error.into()), } } - - info!("Downloading bytes from url: {:?}", url); - - let response = reqwest::get(url.clone()).await?; - let response_bytes = response.bytes().await?; - - info!("Writing file: {:?}", download_filename); - - let mut file = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(&download_filename)?; - - file.write_all(&response_bytes)?; - - Ok(download_filename) + Err(anyhow!("Too many downloads with the same filename").into()) } // TODO: This likely needs more protections. @@ -84,10 +84,7 @@ fn check_url_file_name_for_downloadability(filename: &str) -> Result<(), Artcraf return Err(ArtcraftError::BadDownloadFilename { path: "".into() }.into()); } - if !filename.is_ascii() { - error!("Download filename is not completely ASCII: {:?}", filename); - return Err(ArtcraftError::BadDownloadFilename { path: filename.into() }.into()); - } + PreferredDownloadFilename::validate_custom_format(filename).map_err(|err| anyhow!(err))?; if filename.contains("/") || filename.contains("..") @@ -108,6 +105,8 @@ fn check_url_file_name_for_downloadability(filename: &str) -> Result<(), Artcraf return Err(ArtcraftError::BadDownloadFilename { path: filename.into() }.into()); } + let filename_lower = filename.to_ascii_lowercase(); + let filename = filename_lower.as_str(); if filename.ends_with(".apk") || filename.ends_with(".app") || filename.ends_with(".bat") diff --git a/crates/desktop/artcraft/src/core/utils/mod.rs b/crates/desktop/artcraft/src/core/utils/mod.rs index 5ea5eb93de..9ca7a8f6eb 100644 --- a/crates/desktop/artcraft/src/core/utils/mod.rs +++ b/crates/desktop/artcraft/src/core/utils/mod.rs @@ -1,4 +1,5 @@ pub mod clear_all_webview_cookies; +pub mod auto_download; pub mod download_media_file_to_temp_dir; pub mod download_url_to_temp_dir; pub mod download_url_to_user_download_dir; @@ -12,4 +13,4 @@ pub mod simple_http_download; pub mod simple_http_download_to_tempfile; pub mod task_database_pending_statuses; pub mod upload_bytes_to_media_file; -pub mod window; \ No newline at end of file +pub mod window; diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index f1e56503c9..3b678d13dd 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -19,6 +19,7 @@ use crate::core::commands::generate::models::video::list_video_models_command::l use crate::core::commands::download::download_directory_reveal_command::download_directory_reveal_command; use crate::core::commands::download::download_media_file_command::download_media_file_command; use crate::core::commands::download::download_url_command::download_url_command; +use crate::core::commands::download::get_download_path_command::get_download_path_command; use crate::core::commands::enqueue::image_bg_removal::enqueue_image_bg_removal_command::enqueue_image_bg_removal_command; use crate::core::commands::enqueue::image_to_gaussian::enqueue_image_to_gaussian_command::enqueue_image_to_gaussian_command; use crate::core::commands::enqueue::image_to_object::enqueue_image_to_3d_object_command::enqueue_image_to_3d_object_command; @@ -220,6 +221,7 @@ pub fn run() { download_directory_reveal_command, download_media_file_command, download_url_command, + get_download_path_command, enqueue_image_bg_removal_command, enqueue_image_to_3d_object_command, enqueue_image_to_gaussian_command, diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs index 178002053d..64ee22c99d 100644 --- a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs @@ -2,6 +2,7 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; use crate::core::events::functional_events::text_to_image_generation_complete_event::{GeneratedImage, TextToImageGenerationCompleteEvent}; use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint}; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::task_database::TaskDatabase; @@ -248,8 +249,6 @@ async fn upload_images_to_storyteller( model: None, }; - event.send_infallible(&app_handle); - let task = get_task_by_provider_and_provider_job_id(GetTaskByProviderAndProviderJobIdArgs { db: task_database.get_connection(), provider: GenerationProvider::Grok, @@ -286,6 +285,8 @@ async fn upload_images_to_storyteller( } } + auto_download_task(app_handle, &task, Some(&batch_token), maybe_cdn_url.as_deref(), Some(images.images.len())).await?; + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { db: task_database.get_connection(), task_id: &task.id, @@ -300,6 +301,8 @@ async fn upload_images_to_storyteller( return Ok(()); // If anything breaks with queries, don't spam events. } + clear_auto_download_checkpoint(app_handle, &task); + event.send_infallible(app_handle); send_frontend_ui_update( app_handle, app_env_configs, diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs index c622cfde68..bc7807b607 100644 --- a/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs @@ -1,6 +1,7 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint}; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::data_dir::trait_data_subdir::DataSubdir; @@ -326,6 +327,8 @@ async fn upload_grok_video( } } + auto_download_task(app_handle, local_task, None, maybe_cdn_url.as_deref(), None).await?; + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { db: task_database.get_connection(), task_id: &local_task.id, @@ -340,6 +343,7 @@ async fn upload_grok_video( return Ok(()); // If anything breaks with queries, don't spam events. } + clear_auto_download_checkpoint(app_handle, local_task); let event = GenerationCompleteEvent { //media_file_token: result.media_file_token, action: Some(GenerationAction::GenerateVideo), @@ -365,4 +369,3 @@ async fn upload_grok_video( Ok(()) } - diff --git a/crates/desktop/artcraft/src/services/midjourney/threads/upload_midjourney_batch.rs b/crates/desktop/artcraft/src/services/midjourney/threads/upload_midjourney_batch.rs index d399e33269..8b0448e12a 100644 --- a/crates/desktop/artcraft/src/services/midjourney/threads/upload_midjourney_batch.rs +++ b/crates/desktop/artcraft/src/services/midjourney/threads/upload_midjourney_batch.rs @@ -5,6 +5,7 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::data_dir::trait_data_subdir::DataSubdir; use crate::core::state::task_database::TaskDatabase; +use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint}; use crate::services::midjourney::threads::events::maybe_handle_text_to_image_complete_event::maybe_handle_text_to_image_complete_event; use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image; use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; @@ -88,6 +89,10 @@ pub(super) async fn upload_midjourney_batch(maybe_app_handle: Option<&AppHandle> } } + if let Some(app) = maybe_app_handle { + auto_download_task(app, local_task, Some(&batch_token), maybe_cdn_url.as_deref(), Some(checkpoint.uploaded.len())).await?; + } + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { db: task_database.get_connection(), task_id: &local_task.id, maybe_batch_token: Some(&batch_token), maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), maybe_primary_media_file_class: Some(TaskMediaFileClass::Image), maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(), maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref() }).await?; if !updated { @@ -95,6 +100,7 @@ pub(super) async fn upload_midjourney_batch(maybe_app_handle: Option<&AppHandle> } if let Some(app_handle) = maybe_app_handle { + clear_auto_download_checkpoint(app_handle, local_task); let event = GenerationCompleteEvent { //media_file_token: result.media_file_token, action: Some(GenerationAction::GenerateImage), diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs index 64bf2042c8..b3bf451457 100644 --- a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs @@ -1,6 +1,7 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::utils::auto_download::{auto_download_task_urls, clear_auto_download_checkpoint}; use crate::core::events::generation_events::generation_failed_event::GenerationFailedEvent; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; @@ -73,6 +74,10 @@ pub async fn handle_classic_successful_generations( info!("Task succeeded: {:?}", task_id); + let local_task = &sqlite_tasks_by_sora_task_id[task_id.as_str()]; + let urls = generation.items.iter().map(|item| Url::parse(&item.url)).collect::, _>>()?; + auto_download_task_urls(app_handle, local_task, &urls).await?; + let generation_type = match generation.model_type { CommonModelType::GptImage1 => GenerationType::Image, CommonModelType::Sora2 => GenerationType::Video, @@ -169,6 +174,7 @@ pub async fn handle_classic_successful_generations( }).await?; if updated { + clear_auto_download_checkpoint(app_handle, local_task); // If anything breaks with queries, don't spam events. let event = GenerationCompleteEvent { action: Some(match generation_type { diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/handle_storyteller_successful_job.rs b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/handle_storyteller_successful_job.rs index 3393bb955c..2d6551f16a 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/handle_storyteller_successful_job.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/handle_storyteller_successful_job.rs @@ -6,6 +6,7 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::task_database::TaskDatabase; use crate::core::utils::enum_conversion::generation_provider::to_generation_service_provider; use crate::core::utils::enum_conversion::task_type::to_generation_action; +use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint}; use super::events::maybe_handle_frontend_caller_notification::maybe_handle_frontend_caller_notification; use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem; use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; @@ -32,6 +33,13 @@ pub async fn handle_successful_job( .as_ref() .map(|result| MediaFileToken::new_from_str(&result.entity_token)); + auto_download_task( + app_handle, task, + job.maybe_result.as_ref().and_then(|result| result.maybe_batch_token.as_ref()), + job.maybe_result.as_ref().map(|result| result.media_links.cdn_url.as_str()), + None, + ).await?; + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { db: task_database.get_connection(), task_id: &task.id, @@ -51,6 +59,7 @@ pub async fn handle_successful_job( return Ok(()); // If anything breaks with queries, don't spam events. } + clear_auto_download_checkpoint(app_handle, task); send_additional_success_events(app_handle, app_env_configs, creds, job, task).await; let service = to_generation_service_provider(task.provider); diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/storyteller_task_polling_thread.rs b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/storyteller_task_polling_thread.rs index ee94c674e6..7a5aac74ea 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/storyteller_task_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/storyteller_task_polling_thread.rs @@ -109,7 +109,10 @@ async fn polling_loop( TaskStatus::CompleteSuccess => continue, // NB: We're done with this task. _ => {} } - handle_successful_job(app_handle, app_env_configs, creds.as_ref(), job, task, task_database).await?; + if let Err(err) = handle_successful_job(app_handle, app_env_configs, creds.as_ref(), job, task, task_database).await { + // One failed local download must not delay other completed jobs. + error!("Completion for task {} will be retried: {}", task.id.as_str(), err); + } } JobStatusPlus::CompleteFailure => { match task.status { diff --git a/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs b/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs index 74b1d61fb9..1b52c9f615 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs @@ -2,6 +2,7 @@ use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCrede use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint}; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::data_dir::trait_data_subdir::DataSubdir; @@ -316,6 +317,8 @@ async fn upload_spz_splat( } } + auto_download_task(app_handle, local_task, None, maybe_cdn_url.as_deref(), None).await?; + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { db: task_database.get_connection(), task_id: &local_task.id, @@ -330,6 +333,7 @@ async fn upload_spz_splat( return Ok(()); // If anything breaks with queries, don't spam events. } + clear_auto_download_checkpoint(app_handle, local_task); let event = GenerationCompleteEvent { //media_file_token: result.media_file_token, action: Some(GenerationAction::GenerateGaussian), diff --git a/frontend/apps/artcraft/app/src/components/generation-feed/DesktopGenerationGallery.tsx b/frontend/apps/artcraft/app/src/components/generation-feed/DesktopGenerationGallery.tsx index fbfa3a9b71..69fef3c575 100644 --- a/frontend/apps/artcraft/app/src/components/generation-feed/DesktopGenerationGallery.tsx +++ b/frontend/apps/artcraft/app/src/components/generation-feed/DesktopGenerationGallery.tsx @@ -271,7 +271,7 @@ function ItemActions({ if (!item.fullImage || isDownloading) return; setIsDownloading(true); try { - await downloadMediaFileToDisk(item.fullImage, item.mediaClass); + await downloadMediaFileToDisk(item.fullImage, item.mediaClass, item.modelId); } finally { setIsDownloading(false); } diff --git a/frontend/apps/artcraft/app/src/components/generation-feed/desktopMediaActions.ts b/frontend/apps/artcraft/app/src/components/generation-feed/desktopMediaActions.ts index e2e1e64646..1a658ea51c 100644 --- a/frontend/apps/artcraft/app/src/components/generation-feed/desktopMediaActions.ts +++ b/frontend/apps/artcraft/app/src/components/generation-feed/desktopMediaActions.ts @@ -43,34 +43,10 @@ import { useTabStore } from "~/pages/Stores/TabState"; export const SHARE_URL_BASE = "https://getartcraft.com/media/"; -const EXT_BY_MEDIA_CLASS: Record = { - image: "png", - video: "mp4", - audio: "mp3", - // "dimensional" is the deprecated pre-split 3D class; mesh/splat replace it. - dimensional: "glb", - mesh: "glb", - splat: "spz", -}; - -function extensionForUrl(url: string, mediaClass?: string): string { - try { - const pathname = new URL(url).pathname; - const match = pathname.match(/\.([a-z0-9]{2,5})$/i); - if (match) return match[1].toLowerCase(); - } catch { - // ignore — fall through to mediaClass default - } - if (mediaClass && EXT_BY_MEDIA_CLASS[mediaClass]) { - return EXT_BY_MEDIA_CLASS[mediaClass]; - } - return "bin"; -} - /** Download a media file, prompting for a location when configured to. */ -export async function downloadMediaFileToDisk(url: string, mediaClass?: string) { +export async function downloadMediaFileToDisk(url: string, mediaClass?: string, model?: string) { try { - const chosenPath = await promptDownloadLocationIfNeeded(url); + const chosenPath = await promptDownloadLocationIfNeeded(url, model); if (chosenPath === null) { // User dismissed the picker. return; @@ -78,7 +54,7 @@ export async function downloadMediaFileToDisk(url: string, mediaClass?: string) if (typeof chosenPath === "string") { await downloadUrlToPath(url, chosenPath); } else { - await DownloadUrl(url); + await DownloadUrl(url, { model }); } if (mediaClass === FilterMediaClasses.DIMENSIONAL) { toast.success(`Downloaded 3D model`); @@ -94,7 +70,7 @@ export async function downloadMediaFileToDisk(url: string, mediaClass?: string) /** * Batch download: prompt for a directory, then save each item as its own - * file (`artcraft-{token}.{ext}`). Progress and the outcome are surfaced via + * file using the preferred naming scheme. Progress and the outcome are surfaced via * a single updating toast. Returns true when at least one file was saved * (false on dismiss / nothing downloadable). */ @@ -120,11 +96,11 @@ export async function downloadMediaFilesToFolder( id: toastId, }); try { - const ext = extensionForUrl(item.fullImage, item.mediaClass); - await downloadUrlToPath( - item.fullImage, - `${dir}/artcraft-${item.id}.${ext}`, - ); + await DownloadUrl(item.fullImage, { + directory: dir, + model: item.modelId, + batch_index: downloadable.length > 1 ? i + 1 : undefined, + }); } catch (error) { console.error(">>> Failed to save file:", error); failedCount++; diff --git a/frontend/libs/api/src/lib/LocalApi.ts b/frontend/libs/api/src/lib/LocalApi.ts index 8de8f7759c..f9e92c9780 100644 --- a/frontend/libs/api/src/lib/LocalApi.ts +++ b/frontend/libs/api/src/lib/LocalApi.ts @@ -1,7 +1,7 @@ import { download } from "@tauri-apps/plugin-upload"; -import { downloadDir } from "@tauri-apps/api/path"; import { open, save } from "@tauri-apps/plugin-dialog"; +import { DownloadUrl, GetDownloadPath } from "@storyteller/tauri-api"; const ASK_LOCATION_BEFORE_DOWNLOAD_KEY = "artcraft_ask_location_before_download"; @@ -27,17 +27,6 @@ export const setAskLocationBeforeDownload = (enabled: boolean): void => { } }; -const deriveDownloadFilename = (url: string): string => { - try { - const urlObj = new URL(url); - const last = urlObj.pathname.split("/").pop(); - if (last && last.length > 0) return last; - } catch { - // fall through - } - return "downloaded_file"; -}; - /** * Prompts the user with a native save dialog if the * "Ask location before download" setting is on. @@ -49,10 +38,11 @@ const deriveDownloadFilename = (url: string): string => { */ export const promptDownloadLocationIfNeeded = async ( url: string, + model?: string, ): Promise => { if (!getAskLocationBeforeDownload()) return undefined; - const filename = deriveDownloadFilename(url); - const chosen = await save({ defaultPath: filename }); + const { path } = await GetDownloadPath(url, { model }); + const chosen = await save({ defaultPath: path }); return chosen ?? null; }; @@ -70,29 +60,19 @@ export const downloadUrlToPath = async (url: string, path: string) => { await download(url, path); }; -export const downloadFileFromUrl = async (url: string) => { +export const downloadFileFromUrl = async (url: string, _mediaClass?: string, model?: string) => { console.log("GOT THE URL", url); try { - const filename = deriveDownloadFilename(url); - - let filePath: string; - const chosen = await promptDownloadLocationIfNeeded(url); + const chosen = await promptDownloadLocationIfNeeded(url, model); if (chosen === null) { // User dismissed the picker. return; } if (typeof chosen === "string") { - filePath = chosen; + await download(url, chosen); } else { - const downloadsPath = await downloadDir(); - filePath = `${downloadsPath}/${filename}`; + await DownloadUrl(url, { model }); } - - await download(url, filePath); - - console.log( - `File downloaded and saved to ${filePath}`, - ); } catch (error) { console.error("Error downloading file:", error); throw error; diff --git a/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx b/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx index 3b7cdc68ce..4158d9a2bd 100644 --- a/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx +++ b/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx @@ -280,7 +280,7 @@ interface GalleryModalProps { * (e.g. open the login modal on desktop, navigate to /login on web) instead * of a Retry link. */ onLoginClick?: () => void; - onDownloadClicked?: (url: string, mediaClass?: string) => Promise; + onDownloadClicked?: (url: string, mediaClass?: string, model?: string) => Promise; onAddToSceneClicked?: ( url: string, media_id: string | undefined, diff --git a/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx b/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx index bd35c08ee4..be9821c90d 100644 --- a/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx +++ b/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx @@ -60,7 +60,7 @@ interface LightboxModalProps { additionalInfo?: ReactNode; downloadUrl?: string; mediaId?: string; - onDownloadClicked?: (url: string, mediaClass?: string) => Promise; + onDownloadClicked?: (url: string, mediaClass?: string, model?: string) => Promise; onAddToSceneClicked?: ( url: string, media_id: string | undefined, @@ -1183,7 +1183,7 @@ export function LightboxModal({ onClick={async (e) => { e.stopPropagation(); gtagEvent("download_clicked"); - await onDownloadClicked(actionUrl, mediaClass); + await onDownloadClicked(actionUrl, mediaClass, modelType ?? undefined); }} > Download diff --git a/frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.spec.tsx b/frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.spec.tsx new file mode 100644 index 0000000000..c94fd06d86 --- /dev/null +++ b/frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.spec.tsx @@ -0,0 +1,106 @@ +import { cleanup, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, expect, it, vi } from "vitest"; +import { DownloadsSettingsPane } from "./DownloadsSettingsPane"; + +const api = vi.hoisted(() => ({ get: vi.fn(), update: vi.fn() })); + +vi.mock("@storyteller/tauri-api", () => ({ + GetAppPreferences: api.get, + UpdateAppPreferences: api.update, + DownloadDirectoryReveal: vi.fn(), + PreferenceName: { + AutoDownload: "auto_download", + PreferredDownloadDirectory: "preferred_download_directory", + PreferredDownloadFilename: "preferred_download_filename", + }, +})); +vi.mock("@tauri-apps/plugin-dialog", () => ({ open: vi.fn() })); +vi.mock("@storyteller/api", () => ({ + getAskLocationBeforeDownload: () => false, + setAskLocationBeforeDownload: vi.fn(), +})); + +beforeEach(() => { + vi.stubGlobal("ResizeObserver", class { + observe() {} + unobserve() {} + disconnect() {} + }); + let preferences = { + preferred_download_directory: { custom: "/existing/downloads" }, + auto_download: false, + preferred_download_filename: "artcraft_convention" as string | { custom_format: string }, + }; + api.get.mockImplementation(async () => ({ preferences: { ...preferences } })); + api.update.mockImplementation(async ({ preference, value }) => { + preferences = { ...preferences, [preference]: value }; + return { success: true }; + }); +}); + +afterEach(() => { + cleanup(); + vi.clearAllMocks(); + vi.unstubAllGlobals(); +}); + +it("persists Auto Download using Tauri and keeps the existing directory", async () => { + const user = userEvent.setup(); + render(); + await screen.findByText("/existing/downloads"); + const checkbox = screen.getByRole("checkbox", { name: "Auto Download" }) as HTMLInputElement; + expect(checkbox.checked).toBe(false); + expect(screen.getByText("Downloads to your system the minute generations complete")).toBeTruthy(); + + await user.click(checkbox); + await waitFor(() => expect(checkbox.checked).toBe(true)); + expect(api.update).toHaveBeenCalledWith({ preference: "auto_download", value: true }); + cleanup(); + render(); + await screen.findByText("/existing/downloads"); + expect((screen.getByRole("checkbox", { name: "Auto Download" }) as HTMLInputElement).checked).toBe(true); +}); + +it("saves a custom format explicitly, reloads it, and restores the app convention", async () => { + const user = userEvent.setup(); + render(); + await screen.findByText("/existing/downloads"); + await user.click(screen.getByRole("button", { name: "Preferred download naming scheme" })); + await user.click(screen.getByRole("option", { name: "Custom format" })); + const input = screen.getByRole("textbox", { name: "Custom format" }); + await user.clear(input); + // paste preserves braces, which userEvent.type otherwise treats as keys. + await user.paste("project_{model}_{YYYY}_{batch_index}"); + expect(api.update).not.toHaveBeenCalled(); + await user.click(screen.getByRole("button", { name: "Save format" })); + await waitFor(() => expect(api.update).toHaveBeenCalledWith({ + preference: "preferred_download_filename", + value: { custom_format: "project_{model}_{YYYY}_{batch_index}" }, + })); + cleanup(); + render(); + expect((await screen.findByRole("textbox", { name: "Custom format" }) as HTMLInputElement).value).toBe("project_{model}_{YYYY}_{batch_index}"); + await user.click(screen.getByRole("button", { name: "Preferred download naming scheme" })); + await user.click(screen.getByRole("option", { name: "ArtCraft convention" })); + await waitFor(() => expect(screen.queryByRole("textbox", { name: "Custom format" })).toBeNull()); + expect(api.update).toHaveBeenLastCalledWith({ preference: "preferred_download_filename", value: "artcraft_convention" }); +}); + +it("rejects unsafe names and shows persistence failures without enabling Auto Download", async () => { + const user = userEvent.setup(); + render(); + await screen.findByText("/existing/downloads"); + await user.click(screen.getByRole("button", { name: "Preferred download naming scheme" })); + await user.click(screen.getByRole("option", { name: "Custom format" })); + await user.clear(screen.getByRole("textbox", { name: "Custom format" })); + await user.paste("../escape"); + await user.click(screen.getByRole("button", { name: "Save format" })); + expect(await screen.findByRole("alert")).toBeTruthy(); + expect(api.update).not.toHaveBeenCalled(); + + api.update.mockRejectedValueOnce(new Error("Cannot save preferences")); + await user.click(screen.getByRole("checkbox", { name: "Auto Download" })); + await screen.findByText("Error: Cannot save preferences"); + expect((screen.getByRole("checkbox", { name: "Auto Download" }) as HTMLInputElement).checked).toBe(false); +}); diff --git a/frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.tsx b/frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.tsx index 3d78e4ba89..5821589a43 100644 --- a/frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.tsx +++ b/frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.tsx @@ -6,22 +6,35 @@ import { DownloadDirectoryReveal, GetAppPreferences, PreferenceName, + PreferredDownloadFilename, SystemDirectory, UpdateAppPreferences, } from "@storyteller/tauri-api"; import { open } from "@tauri-apps/plugin-dialog"; import { Label } from "@storyteller/ui-label"; import { Switch } from "@storyteller/ui-switch"; +import { Select } from "@storyteller/ui-select"; import { FolderIcon, RotateCcwIcon, SearchIcon } from "lucide-react"; import { getAskLocationBeforeDownload, setAskLocationBeforeDownload, } from "@storyteller/api"; +const DEFAULT_CUSTOM_FORMAT = "{model}_{date}"; +const FILENAME_OPTIONS = [ + { value: "artcraft_convention", label: "ArtCraft convention" }, + { value: "custom", label: "Custom format" }, +]; + export const DownloadsSettingsPane = () => { const [preferences, setPreferences] = useState< AppPreferencesPayload | undefined >(undefined); + const [filenameMode, setFilenameMode] = useState("artcraft_convention"); + const [customFormat, setCustomFormat] = useState(DEFAULT_CUSTOM_FORMAT); + const [formatError, setFormatError] = useState(null); + const [settingsError, setSettingsError] = useState(null); + const [saving, setSaving] = useState(false); const [askLocationBeforeDownload, setAskLocationBeforeDownloadState] = useState(() => getAskLocationBeforeDownload()); @@ -34,10 +47,14 @@ export const DownloadsSettingsPane = () => { useEffect(() => { const fetchData = async () => { const prefs = await GetAppPreferences(); - console.log("prefs", prefs); setPreferences(prefs.preferences); + const filename = prefs.preferences.preferred_download_filename; + if (filename && typeof filename === "object") { + setFilenameMode("custom"); + setCustomFormat(filename.custom_format); + } }; - fetchData(); + fetchData().catch((error) => setSettingsError(String(error))); }, []); // NB: This might be a complex type. @@ -53,7 +70,6 @@ export const DownloadsSettingsPane = () => { const reloadPreferences = async () => { const prefs = await GetAppPreferences(); - console.log("prefs", prefs); setPreferences(prefs.preferences); }; @@ -89,8 +105,64 @@ export const DownloadsSettingsPane = () => { await DownloadDirectoryReveal(); }; + const toggleAutoDownload = async (enabled: boolean) => { + setSaving(true); + setSettingsError(null); + try { + await UpdateAppPreferences({ preference: PreferenceName.AutoDownload, value: enabled }); + await reloadPreferences(); + } catch (error) { + setSettingsError(String(error)); + } finally { + setSaving(false); + } + }; + + const saveFilenamePreference = async (value: PreferredDownloadFilename) => { + setSaving(true); + setFormatError(null); + try { + await UpdateAppPreferences({ preference: PreferenceName.PreferredDownloadFilename, value }); + await reloadPreferences(); + setFilenameMode(typeof value === "string" ? value : "custom"); + } catch (error) { + setFormatError(String(error)); + } finally { + setSaving(false); + } + }; + + const saveCustomFormat = async () => { + const format = customFormat.trim(); + if (!format) { + setFormatError("Format cannot be empty"); + return; + } + if (/[/\\'"`%<>|:*?\x00-\x1f\x7f]/.test(format) || format.includes("..")) { + setFormatError("Format cannot contain slashes, quotes, or other unsafe characters"); + return; + } + await saveFilenamePreference({ custom_format: format }); + }; + return (
+ {settingsError &&

{settingsError}

} +
+ +

+ Downloads to your system the minute generations complete +

+

@@ -132,6 +204,49 @@ export const DownloadsSettingsPane = () => { setEnabled={toggleAskLocationBeforeDownload} />

+
+ +

+ {"The ArtCraft convention is {model}_{date}.{ext}, with a batch number when a generation produces several files."} +

+
+ setCustomFormat(event.target.value)} + placeholder={DEFAULT_CUSTOM_FORMAT} + spellCheck={false} + className="min-w-0 flex-1 rounded-md border border-ui-controls-border bg-ui-controls px-3 py-2 font-mono text-xs" + /> + +
+

+ {"Tokens: {model}, {date}, {YYYY}, {YY}, {MM}, {DD}, {HH}, {mm}, {SS}, {batch_index}. The file extension is added automatically. Batch numbers are added automatically if omitted."} +

+

Changes apply after you click Save format.

+ + )} + {formatError &&

{formatError}

} +
); }; diff --git a/frontend/libs/tauri-api/src/lib/download/DownloadUrl.ts b/frontend/libs/tauri-api/src/lib/download/DownloadUrl.ts index b8d13bd2ad..aa268cf579 100644 --- a/frontend/libs/tauri-api/src/lib/download/DownloadUrl.ts +++ b/frontend/libs/tauri-api/src/lib/download/DownloadUrl.ts @@ -8,11 +8,21 @@ export interface DownloadUrlSuccess extends CommandResult { export interface DownloadUrlPayload { } -export const DownloadUrl = async (url: string) : Promise => { +export interface DownloadOptions { + model?: string; + batch_index?: number; + directory?: string; +} + +export const GetDownloadPath = (url: string, options: DownloadOptions = {}): Promise<{ filename: string; path: string }> => + invoke("get_download_path_command", { request: { url, ...options } }); + +export const DownloadUrl = async (url: string, options: DownloadOptions = {}) : Promise => { try { return await invoke("download_url_command", { request: { - url: url + url, + ...options, } }) as DownloadUrlSuccess; } catch (error) { diff --git a/frontend/libs/tauri-api/src/lib/settings/GetAppPreferences.ts b/frontend/libs/tauri-api/src/lib/settings/GetAppPreferences.ts index b887f36f21..41f0c396f6 100644 --- a/frontend/libs/tauri-api/src/lib/settings/GetAppPreferences.ts +++ b/frontend/libs/tauri-api/src/lib/settings/GetAppPreferences.ts @@ -7,6 +7,8 @@ export interface GetAppPreferencesResult { export interface AppPreferencesPayload { // Preferred download directory preferred_download_directory: PreferredDownloadDirectory, + preferred_download_filename: PreferredDownloadFilename, + auto_download: boolean, // Play sounds on events. play_sounds: boolean, @@ -23,6 +25,12 @@ export interface AppPreferencesPayload { export type PreferredDownloadDirectory = SystemDirectory | CustomDirectory; +export type PreferredDownloadFilename = "artcraft_convention" | CustomFilenameFormat; + +export interface CustomFilenameFormat { + custom_format: string; +} + export interface SystemDirectory { // If the directory is a system directory. system: string, diff --git a/frontend/libs/tauri-api/src/lib/settings/UpdateAppPreference.ts b/frontend/libs/tauri-api/src/lib/settings/UpdateAppPreference.ts index 57a5277df3..18feec3e06 100644 --- a/frontend/libs/tauri-api/src/lib/settings/UpdateAppPreference.ts +++ b/frontend/libs/tauri-api/src/lib/settings/UpdateAppPreference.ts @@ -1,14 +1,16 @@ import { invoke } from "@tauri-apps/api/core"; -import { PreferredDownloadDirectory } from "./GetAppPreferences"; +import { PreferredDownloadDirectory, PreferredDownloadFilename } from "./GetAppPreferences"; export interface UpdateAppPreferencesRequest { preference: PreferenceName, - value: undefined | string | boolean | PreferredDownloadDirectory, + value: undefined | string | boolean | PreferredDownloadDirectory | PreferredDownloadFilename, } export enum PreferenceName { PreferredDownloadDirectory = "preferred_download_directory", + PreferredDownloadFilename = "preferred_download_filename", + AutoDownload = "auto_download", PlaySounds = "play_sounds", DeleteFileSound = "delete_file_sound", EnqueueSuccessSound = "enqueue_success_sound",