fix sora downloads

Cargo.lock

@@ -694,6 +694,7 @@ dependencies = [
  "tokio",
  "tokio-util",
  "url",
+ "url_utils",
  "web_base64",
  "workspace-build-acceleration",
 ]
@@ -12457,6 +12458,14 @@ dependencies = [
  "workspace-build-acceleration",
 ]
 
+[[package]]
+name = "url_utils"
+version = "0.0.1"
+dependencies = [
+ "url",
+ "workspace-build-acceleration",
+]
+
 [[package]]
 name = "urlencoding"
 version = "2.1.3"

Cargo.toml

@@ -84,6 +84,7 @@ members = [
   "crates/lib/testing",
   "crates/lib/tts_common",
   "crates/lib/url_config",
+  "crates/lib/url_utils",
   "crates/lib/user_input_common",
   "crates/lib/web_scrapers",
   #"crates/lib/api_clients/fakeyou_client",
@@ -210,6 +211,7 @@ subprocess_common = { path = "crates/lib/subprocess_common" }
 test_data = { path = "crates/lib/test_data" }
 testing = { path = "crates/lib/testing" }
 tokens = { path = "crates/schema/public/tokens" }
+url_utils = { path = "crates/lib/url_utils" }
 web_base64 = { path = "crates/lib/data/web_base64" }
 
 # External libs

crates/desktop/artcraft/Cargo.toml

@@ -42,6 +42,7 @@ primitives.workspace = true
 sqlite_tasks.workspace = true
 storyteller_client.workspace = true
 tokens.workspace = true
+url_utils.workspace = true
 web_base64.workspace = true
 
 # External

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/download_extension.rs

@@ -0,0 +1,15 @@
+
+#[derive(Clone, Copy)]
+pub enum DownloadExtension {
+  Png,
+  Mp4,
+}
+
+impl DownloadExtension {
+  pub fn as_extension_without_period(&self) -> &'static str {
+    match self {
+      DownloadExtension::Png => "png",
+      DownloadExtension::Mp4 => "mp4",
+    }
+  }
+}

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs

@@ -9,6 +9,7 @@ use crate::core::state::task_database::TaskDatabase;
 use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
 use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
 use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
+use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension;
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
 use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
 use enums::common::generation_provider::GenerationProvider;
@@ -35,6 +36,7 @@ use storyteller_client::endpoints::media_files::upload_video_media_file_from_fil
 use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
+use url_utils::extension::extract_download_extension_from_url::{extract_download_extension_from_url, extract_download_extension_from_url_str};
 
 pub struct SuccessfulGeneration {
   pub prompt: Option<String>,
@@ -61,6 +63,7 @@ pub async fn handle_classic_successful_generations(
   storyteller_creds: &StorytellerCredentialSet,
   succeeded_tasks_by_id: &HashMap<TaskId, SuccessfulGeneration>,
   sqlite_tasks_by_sora_task_id: &HashMap<String, Task>,
+  recommended_download_extension: DownloadExtension,
 ) -> AnyhowResult<()> {
 
   for (task_id, generation) in succeeded_tasks_by_id.iter() {
@@ -98,7 +101,7 @@ pub async fn handle_classic_successful_generations(
 
     for (_i, item) in generation.items.iter().enumerate() {
       info!("Downloading generated file...");
-      let download_path = download_generation_item(item, &app_data_root).await?;
+      let download_path = download_generation_item(item, &app_data_root, recommended_download_extension).await?;
 
       info!("Uploading to backend...");
 
@@ -158,18 +161,26 @@ pub async fn handle_classic_successful_generations(
 }
 
 
-async fn download_generation_item(generation: &GenerationItem, app_data_root: &AppDataRoot) -> AnyhowResult<PathBuf> {
-  let url = Url::parse(&generation.url)?;
+async fn download_generation_item(
+  generation: &GenerationItem,
+  app_data_root: &AppDataRoot,
+  recommended_download_extension: DownloadExtension
+) -> AnyhowResult<PathBuf> {
+  info!("Downloading generation item from URL: {}", generation.url.as_str());
 
   let response = reqwest::get(&generation.url).await?;
   let image_bytes = response.bytes().await?;
 
-  let ext = url.path().split(".").last().unwrap_or("png");
-
+  let extension = extract_download_extension_from_url_str(&generation.url)
+      .map(|ext| ext.as_extension_without_period())
+      .unwrap_or_else(|| recommended_download_extension.as_extension_without_period());
+  
   let tempdir = app_data_root.temp_dir().path();
-  let download_filename = format!("{}.{}", generation.item_id, ext);
+  let download_filename = format!("{}.{}", generation.item_id, extension);
   let download_path = tempdir.join(download_filename);
 
+  info!("Writing to path: {:?}", download_path);
+
   let mut file = File::create(&download_path)?;
   file.write_all(&image_bytes)?;
 

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/mod.rs

@@ -2,3 +2,4 @@ pub (super) mod handle_successful_generations;
 pub (super) mod handle_failed_generations;
 pub (super) mod poll_classic_sora_tasks;
 pub (super) mod poll_sora_2_tasks;
+mod download_extension;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_classic_sora_tasks.rs

@@ -36,6 +36,7 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil
 use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
+use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension;
 
 pub async fn poll_classic_sora_tasks(
   app_handle: &AppHandle,
@@ -117,6 +118,7 @@ pub async fn poll_classic_sora_tasks(
     &storyteller_creds,
     &sora_succeeded_tasks_by_id,
     &local_sqlite_tasks_by_sora_task_id,
+    DownloadExtension::Png,
   ).await?;
 
   Ok(())

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_sora_2_tasks.rs

@@ -38,6 +38,7 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil
 use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
+use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension;
 
 pub async fn poll_sora_2_tasks(
   app_handle: &AppHandle,
@@ -110,6 +111,7 @@ pub async fn poll_sora_2_tasks(
     &storyteller_creds,
     &sora_succeeded_drafts_by_id,
     &local_sqlite_tasks_by_sora_task_id,
+    DownloadExtension::Mp4,
   ).await?;
 
   Ok(())

crates/lib/url_utils/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "url_utils"
+edition = "2024"
+version = "0.0.1"
+publish = false
+
+[package.metadata.cargo-udeps.ignore]
+normal = ["workspace-build-acceleration"]
+
+[lib]
+name = "url_utils"
+path = "src/lib.rs"
+
+[dependencies]
+
+# Workspace hack (faster compile times)
+workspace-build-acceleration.workspace = true
+
+# Internal
+# (None)
+
+# External
+url.workspace = true
+
+[dev-dependencies]
+# (None)

crates/lib/url_utils/src/extension/download_extension.rs

@@ -0,0 +1,71 @@
+
+/// When downloading a remote file, it makes sense to try to match the same file extension (with reasonable safeguards).
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum DownloadExtension {
+  // Image
+  Jpg,
+  Png,
+  // Video
+  Mp4,
+  Webm,
+  // Audio
+  Mp3,
+  Wav,
+  // Text
+  Json,
+  Txt,
+}
+
+impl DownloadExtension {
+  pub fn as_extension_without_period(&self) -> &'static str {
+    match self {
+      // Image
+      Self::Jpg => "jpg",
+      Self::Png => "png",
+      // Video
+      Self::Mp4 => "mp4",
+      Self::Webm => "webm",
+      // Audio
+      Self::Mp3 => "mp3",
+      Self::Wav => "wav",
+      // Text
+      Self::Json => "json",
+      Self::Txt => "txt",
+    }
+  }
+
+  pub fn as_extension_with_period(&self) -> &'static str {
+    match self {
+      // Image
+      Self::Jpg => ".jpg",
+      Self::Png => ".png",
+      // Video
+      Self::Mp4 => ".mp4",
+      Self::Webm => ".webm",
+      // Audio
+      Self::Mp3 => ".mp3",
+      Self::Wav => ".wav",
+      // Text
+      Self::Json => ".json",
+      Self::Txt => ".txt",
+    }
+  }
+
+  pub fn try_from_str(extension: &str) -> Option<Self> {
+    match extension.to_lowercase().as_str() {
+      // Image
+      "jpg" | "jpeg" => Some(Self::Jpg),
+      "png" => Some(Self::Png),
+      // Video
+      "mp4" => Some(Self::Mp4),
+      "webm" => Some(Self::Webm),
+      // Audio
+      "mp3" => Some(Self::Mp3),
+      "wav" => Some(Self::Wav),
+      // Text
+      "json" => Some(Self::Json),
+      "txt" => Some(Self::Txt),
+      _ => None,
+    }
+  }
+}

crates/lib/url_utils/src/extension/extract_download_extension_from_url.rs

@@ -0,0 +1,37 @@
+use crate::extension::download_extension::DownloadExtension;
+use url::Url;
+
+/// This should be used for naming files on the local file system based on their URL.
+pub fn extract_download_extension_from_url_str(url: &str) -> Option<DownloadExtension> {
+  let parsed_url = Url::parse(url).ok()?;
+  extract_download_extension_from_url(&parsed_url)
+}
+
+/// This should be used for naming files on the local file system based on their URL.
+pub fn extract_download_extension_from_url(url: &Url) -> Option<DownloadExtension> {
+  let path = url.path();
+  let extension = std::path::Path::new(path)
+      .extension()
+      .and_then(|ext| ext.to_str())?;
+  DownloadExtension::try_from_str(extension)
+}
+
+#[cfg(test)]
+mod tests {
+  use crate::extension::download_extension::DownloadExtension;
+  use crate::extension::extract_download_extension_from_url::extract_download_extension_from_url_str;
+
+  #[test]
+  fn success_case() {
+    let url = "https://storyteller.ai/example.jpg";
+    let maybe_extension= extract_download_extension_from_url_str(url);
+    assert_eq!(maybe_extension, Some(DownloadExtension::Jpg));
+  }
+
+  #[test]
+  fn difficult_url() {
+    let url = "https://videos.youtube.com/foo/bar/files/00000000-f19a-1275-a092-caa60b38f8d2%2Fraw?se=2025-11-21T22%3A26%3A53Z&sp=r&sv=2024-08-09&sr=b&skoid=8cfff87a-01fc-27c1-9091-32999dcd1389&sktid=aa8ccaf6-e6da-484e-ab11-9c7496c21c53&skt=2025-11-17T21%3A09%3A08Z&ske=2025-10-23T21%3A14%3A08Z&sks=b&skv=2024-08-04&sig=T2U6%2Brw9WJdJagBkCL4/GlFr72KhZGP3QPlobIO7JlU%3D&ac=oaisdsorprnorthcentralus";
+    let maybe_extension= extract_download_extension_from_url_str(url);
+    assert!(maybe_extension.is_none());
+  }
+}

crates/lib/url_utils/src/extension/mod.rs

@@ -0,0 +1,2 @@
+pub mod download_extension;
+pub mod extract_download_extension_from_url;

crates/lib/url_utils/src/lib.rs

@@ -0,0 +1 @@
+pub mod extension;