logging

crates/api_clients/storyteller_client/src/media_files/upload_image_media_file_from_file.rs

@@ -47,6 +47,11 @@ pub struct UploadImageFromFileArgs<'a, P: AsRef<Path>> {
   // /// If provided, this is the service provider that created the image.
   // /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
   // pub maybe_generation_provider: Option<GenerationProvider>,
+  
+  /// If provided, this groups the file into a batch
+  /// TODO: This shouldn't be set clientside without the backend generating the token 
+  ///  and cryptographically securing it. But we need to go fast here.
+  pub maybe_batch_token: Option<&'a BatchGenerationToken>,
 }
 
 
@@ -81,6 +86,10 @@ pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
   if let Some(prompt_token) = &args.maybe_prompt_token {
     form = form.text("maybe_prompt_token", prompt_token.to_string());
   }
+  
+  if let Some(batch_token) = &args.maybe_batch_token {
+    form = form.text("maybe_batch_token", batch_token.to_string());
+  }
 
   let mut request_builder = client.post(url)
       .header("User-Agent", USER_AGENT)

crates/api_clients/storyteller_client/src/recipes/upload_media_file_from_file.rs

@@ -57,6 +57,7 @@ pub async fn upload_media_file_from_file<P: AsRef<Path>>(
         path,
         is_intermediate_system_file: false,
         maybe_prompt_token,
+        maybe_batch_token: None,
       }).await;
       match result {
         Ok(result) => Ok(UploadMediaFileSuccessResponse {

crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/generic/handle_generic_bg_removal_artcraft.rs

@@ -128,6 +128,7 @@ async fn upload_image_from_base64_bytes(
         path: temp_file,
         is_intermediate_system_file: true, // NB: Probably not essential to keep this.
         maybe_prompt_token: None, // NB: Not used for bg removal.
+        maybe_batch_token: None, // NB: Not used for bg removal.
       }).await
           .map_err(|err| {
             error!("Failed to upload image media file: {:?}", err);

crates/desktop/artcraft/src/services/fal/commands/fal_background_removal_command.rs

@@ -262,6 +262,7 @@ pub async fn remove_background(
     path: result_filename,
     is_intermediate_system_file: true, // NB: Probably not essential to keep this.
     maybe_prompt_token: None, // NB: Not used for bg removal.
+    maybe_batch_token: None, // NB: Not used for bg removal.
   }).await?;
 
   // TODO: Don't re-request to simply build MediaLinks (or CDN URL). Get those from the upload API in one turn.

crates/desktop/artcraft/src/services/midjourney/threads/events/maybe_handle_text_to_image_complete_event.rs

@@ -0,0 +1,83 @@
+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::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use anyhow::anyhow;
+use artcraft_api_defs::jobs::list_session_jobs::{ListSessionJobsItem, ListSessionResultDetailsResponse};
+use enums::tauri::tasks::task_type::TaskType;
+use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
+use errors::AnyhowResult;
+use log::{error, warn};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
+use storyteller_client::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
+use tauri::AppHandle;
+use tokens::tokens::batch_generations::BatchGenerationToken;
+use tokens::tokens::media_files::MediaFileToken;
+
+pub async fn maybe_handle_text_to_image_complete_event(
+  app: &AppHandle,
+  app_env_configs: &AppEnvConfigs,
+  maybe_creds: Option<&StorytellerCredentialSet>,
+  task: &Task,
+  batch_token: &BatchGenerationToken,
+) -> AnyhowResult<()> {
+
+  match task.task_type {
+    TaskType::ImageGeneration => {} // NB: Fall-through
+    _ => return Ok(()),
+  }
+
+  match task.frontend_caller {
+    Some(TauriCommandCaller::TextToImage) => {} // NB: Fall-through
+    _ => return Ok(()),
+  }
+
+  let event = handle_batch(
+    app,
+    app_env_configs,
+    maybe_creds,
+    task,
+    batch_token,
+  ).await?;
+
+
+  if let Err(err) = event.send(&app) {
+    error!("Failed to send TextToImageGenerationCompleteEvent: {:?}", err); // Fail open
+  }
+
+  Ok(())
+}
+
+async fn handle_batch(
+  app: &AppHandle,
+  app_env_configs: &AppEnvConfigs,
+  maybe_creds: Option<&StorytellerCredentialSet>,
+  task: &Task,
+  batch_token: &BatchGenerationToken,
+) -> AnyhowResult<TextToImageGenerationCompleteEvent> {
+
+  let result = list_batch_generated_redux_media_files(
+    &app_env_configs.storyteller_host,
+    maybe_creds,
+    batch_token,
+  ).await?;
+
+  if result.media_files.is_empty() {
+    return Err(anyhow!("No media files found for batch token: {}", batch_token));
+  }
+
+  let media_files = result.media_files
+      .into_iter()
+      .map(|file| GeneratedImage {
+        media_token: file.token,
+        cdn_url: file.media_links.cdn_url,
+        maybe_thumbnail_template: file.media_links.maybe_thumbnail_template,
+      })
+      .collect();
+
+  Ok(TextToImageGenerationCompleteEvent {
+    generated_images: media_files,
+    maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(),
+    maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(),
+  })
+}

crates/desktop/artcraft/src/services/midjourney/threads/events/mod.rs

@@ -0,0 +1 @@
+pub (super) mod maybe_handle_text_to_image_complete_event;

crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs

@@ -6,6 +6,7 @@ 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::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
+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 crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
 use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
@@ -35,6 +36,7 @@ use storyteller_client::error::storyteller_error::StorytellerError;
 use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
 use storyteller_client::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
+use tokens::tokens::batch_generations::BatchGenerationToken;
 use url::Url;
 
 static PENDING_STATUSES : Lazy<HashSet<TaskStatus>> = Lazy::new(|| {
@@ -263,6 +265,12 @@ async fn upload_midjourney_batch(
 
   info!("Created prompt: {:?}", &prompt_response.prompt_token);
 
+  // TODO: Move this from clientside to the backend.
+  //  The first upload should produce a batch token that we can reuse.
+  let batch_token = BatchGenerationToken::generate();
+
+  info!("Using synthetic batch token: {:?}", &batch_token);
+
   for index in 0..4 {
     info!("Downloading generated Midjourney file...");
 
@@ -297,6 +305,7 @@ async fn upload_midjourney_batch(
         path: &download_path,
         is_intermediate_system_file: false,
         maybe_prompt_token: Some(&prompt_response.prompt_token),
+        maybe_batch_token: Some(&batch_token),
       }).await;
 
       match result {
@@ -345,6 +354,18 @@ async fn upload_midjourney_batch(
     error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
   }
 
+  let result = maybe_handle_text_to_image_complete_event(
+    app_handle,
+    app_env_configs,
+    Some(storyteller_creds),
+    local_task,
+    &batch_token,
+  ).await;
+
+  if let Err(err) = result {
+    error!("Failed to send text-to-image complete event: {:?}", err);
+  }
+
   Ok(())
 }
 

crates/desktop/artcraft/src/services/midjourney/threads/mod.rs

@@ -1 +1,2 @@
-pub mod midjourney_long_polling_thread;
+pub mod midjourney_long_polling_thread;
+mod events;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling_thread.rs

@@ -141,6 +141,7 @@ async fn polling_loop(
           path: download_path,
           is_intermediate_system_file: false,
           maybe_prompt_token: None, // TODO: this should be added soon.
+          maybe_batch_token: None, // TODO: This should be added soon.
         }).await?;
 
         info!("Uploaded to API backend: {:?}", result.media_file_token);

crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_text_to_image_complete_event.rs

@@ -69,7 +69,7 @@ pub async fn maybe_handle_text_to_image_complete_event(
   Ok(())
 }
 
-pub async fn handle_single(
+async fn handle_single(
   app: &AppHandle,
   task: &Task,
   job: &ListSessionJobsItem,
@@ -86,7 +86,7 @@ pub async fn handle_single(
   })
 }
 
-pub async fn handle_batch(
+async fn handle_batch(
   app: &AppHandle,
   app_env_configs: &AppEnvConfigs,
   maybe_creds: Option<&StorytellerCredentialSet>,

crates/desktop/artcraft/src/services/storyteller/threads/events/mod.rs

@@ -1,3 +1,3 @@
-pub mod maybe_handle_inpainting_complete_event;
-pub mod maybe_handle_text_to_image_complete_event;
-pub mod maybe_send_background_removal_complete_event;
+pub (super) mod maybe_handle_inpainting_complete_event;
+pub (super) mod maybe_handle_text_to_image_complete_event;
+pub (super) mod maybe_send_background_removal_complete_event;

crates/schema/database/mysql_queries/src/queries/media_files/create/specialized_insert/insert_media_file_from_file_upload.rs

@@ -12,6 +12,7 @@ use enums::by_table::media_files::media_file_type::MediaFileType;
 use enums::common::visibility::Visibility;
 use errors::AnyhowResult;
 use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
+use tokens::tokens::batch_generations::BatchGenerationToken;
 use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::prompts::PromptToken;
 use tokens::tokens::users::UserToken;
@@ -45,6 +46,7 @@ pub struct InsertMediaFileFromUploadArgs<'a> {
   pub sha256_checksum: &'a str,
 
   pub maybe_prompt_token: Option<&'a PromptToken>,
+  pub maybe_batch_token: Option<&'a BatchGenerationToken>,
 
   pub maybe_title: Option<&'a str>,
   pub maybe_duration_millis: Option<u64>,
@@ -116,6 +118,7 @@ SET
   checksum_sha2 = ?,
 
   maybe_prompt_token = ?,
+  maybe_batch_token = ?,
 
   maybe_title = ?,
   maybe_duration_millis = ?,
@@ -161,6 +164,7 @@ SET
       args.sha256_checksum,
 
       args.maybe_prompt_token.map(|t| t.as_str()),
+      args.maybe_batch_token.map(|t| t.as_str()),
 
       args.maybe_title,
       args.maybe_duration_millis,

crates/service/web/storyteller_web/src/http_server/deprecated_endpoints/engine/create_scene_handler.rs

@@ -206,6 +206,7 @@ pub async fn create_scene_handler(
     maybe_animation_type: None,
     maybe_mime_type: Some(mime_type),
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     file_size_bytes: file_size_bytes as u64,
     maybe_duration_millis: None,
     sha256_checksum: &hash,

crates/service/web/storyteller_web/src/http_server/deprecated_endpoints/engine/process_upload_media_file.rs

@@ -333,6 +333,7 @@ pub async fn process_upload_media_file(
     maybe_engine_category: None,
     maybe_animation_type: None,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_mime_type: Some(mime_type),
     file_size_bytes: file_size_bytes as u64,
     maybe_duration_millis: None,

crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/upload/upload_snapshot_media_file_handler.rs

@@ -247,6 +247,7 @@ pub async fn upload_snapshot_media_file_handler(
     maybe_engine_category: None,
     maybe_animation_type: None,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_mime_type: Some(&mimetype),
     file_size_bytes: file_size_bytes as u64,
     maybe_duration_millis: None,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_audio_media_file_handler.rs

@@ -325,6 +325,7 @@ pub async fn upload_audio_media_file_handler(
     maybe_engine_category: None,
     maybe_animation_type: None,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_mime_type: Some(&mimetype),
     file_size_bytes: file_size_bytes as u64,
     maybe_duration_millis,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_engine_asset/upload_engine_asset_media_file_handler.rs

@@ -210,6 +210,7 @@ pub async fn upload_engine_asset_media_file_handler(
     creator_ip_address: &ip_address,
     creator_set_visibility,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     upload_type: UploadType::Filesystem,
     maybe_engine_category: upload_media_request.maybe_engine_category,
     maybe_animation_type: upload_media_request.maybe_animation_type,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_generic/process_upload_media_file.rs

@@ -365,6 +365,7 @@ pub async fn process_upload_media_file(
     upload_type,
     maybe_engine_category: None,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_animation_type: None,
     maybe_mime_type: Some(mime_type),
     file_size_bytes: file_size_bytes as u64,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_image_media_file_handler.rs

@@ -26,6 +26,7 @@ use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes;
 use mimetypes::mimetype_to_extension::mimetype_to_extension;
 use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
 use mysql_queries::queries::media_files::create::specialized_insert::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType};
+use tokens::tokens::batch_generations::BatchGenerationToken;
 use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::prompts::PromptToken;
 
@@ -62,7 +63,13 @@ pub struct UploadImageMediaFileForm {
   #[schema(value_type = Option<String>, format = Binary)]
   maybe_prompt_token: Option<Text<PromptToken>>,
 
-  /// Optional: Whether this is a system file (eg. cover files we should hide)
+  /// Optional: Batch associated with this image
+  /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
+  #[multipart(limit = "2 KiB")]
+  #[schema(value_type = Option<String>, format = Binary)]
+  maybe_batch_token: Option<Text<BatchGenerationToken>>,
+
+  /// Optional: Whether this is a system file (e.g. cover files we should hide)
   #[multipart(limit = "2 KiB")]
   #[schema(value_type = Option<bool>, format = Binary)]
   is_intermediate_system_file: Option<Text<bool>>,
@@ -271,6 +278,10 @@ pub async fn upload_image_media_file_handler(
   let maybe_prompt_token = form.maybe_prompt_token
       .as_ref()
       .map(|token| token.0.clone());
+  
+  let maybe_batch_token = form.maybe_batch_token
+      .as_ref()
+      .map(|token| token.0.clone());
 
   let (token, record_id) = insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs {
     maybe_media_class: Some(MediaFileClass::Image),
@@ -284,6 +295,7 @@ pub async fn upload_image_media_file_handler(
     maybe_animation_type: None,
     maybe_mime_type: Some(&mimetype),
     maybe_prompt_token: maybe_prompt_token.as_ref(),
+    maybe_batch_token: maybe_batch_token.as_ref(),
     file_size_bytes: file_size_bytes as u64,
     maybe_duration_millis: None,
     sha256_checksum: &hash,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_new_engine_asset_media_file_handler.rs

@@ -207,6 +207,7 @@ pub async fn upload_new_engine_asset_media_file_handler(
     creator_set_visibility: file_info.creator_set_visibility,
     upload_type: UploadType::Filesystem,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_engine_category: Some(file_info.engine_category),
     maybe_animation_type: file_info.maybe_animation_type,
     maybe_mime_type: Some(file_info.mimetype),

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_new_scene_media_file_handler.rs

@@ -231,6 +231,7 @@ pub async fn upload_new_scene_media_file_handler(
     creator_ip_address: &ip_address,
     creator_set_visibility,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     upload_type: UploadType::Filesystem,
     maybe_engine_category: Some(MediaFileEngineCategory::Scene),
     maybe_animation_type: None,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_pmx/upload_pmx_media_file_handler.rs

@@ -265,6 +265,7 @@ pub async fn upload_pmx_media_file_handler(
     maybe_engine_category: Some(engine_category),
     maybe_animation_type,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_mime_type: Some("application/octet-stream"),
     file_size_bytes: pmx_details.file_size_bytes,
     maybe_duration_millis,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_scene_snapshot_media_file_handler.rs

@@ -225,6 +225,7 @@ pub async fn upload_scene_snapshot_media_file_handler(
     creator_set_visibility: Visibility::Hidden,
     upload_type: UploadType::Filesystem,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_engine_category: Some(MediaFileEngineCategory::Scene),
     maybe_animation_type: None,
     maybe_mime_type: Some(MIMETYPE),

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_video_new/upload_new_video_media_file_handler.rs

@@ -372,6 +372,7 @@ pub async fn upload_new_video_media_file_handler(
     maybe_engine_category: None,
     maybe_animation_type: None,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_mime_type: Some(&mimetype),
     file_size_bytes: file_size_bytes as u64,
     maybe_duration_millis,

crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/upload_video_old/upload_video_media_file_handler.rs

@@ -228,6 +228,7 @@ pub async fn upload_video_media_file_handler(
     upload_type,
     maybe_engine_category: None,
     maybe_prompt_token: None,
+    maybe_batch_token: None,
     maybe_animation_type: None,
     maybe_mime_type: Some(mimetype),
     file_size_bytes: file_size_bytes as u64,