cleanup

crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs

@@ -1,6 +1,6 @@
 use crate::core::commands::enqueue::generate_error::GenerateError;
 use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
-use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft_via_omni_endpoint::handle_artcraft_via_omni_endpoint;
+use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft;
 use crate::core::commands::generate::generate_image::tauri_generate_image_request::{
   TauriGenerateImageErrorType, TauriGenerateImageRequest, TauriGenerateImageResponse,
 };
@@ -44,55 +44,49 @@ pub async fn generate_image_command(
     });
   }
 
-  let result = handle_artcraft_via_omni_endpoint(
+  let result = handle_artcraft(
     &request,
     &app_env_configs,
     &storyteller_creds_manager,
   ).await;
 
-  match result {
-    Ok(success) => {
-      // Insert into task database
-      let db_result = success
-        .insert_into_task_database_with_frontend_payload(
-          &task_database,
-          request.frontend_caller,
-          request.frontend_subscriber_id.as_deref(),
-          request.frontend_subscriber_payload.as_deref(),
-        )
-        .await;
+  let success = match result {
+    Ok(success) => success,
+    Err(err) => return map_error_to_response(err),
+  };
 
-      if let Err(err) = db_result {
-        error!("Failed to create task in database: {:?}", err);
-      }
+  // Insert into task database
+  let db_result = success
+      .insert_into_task_database_with_frontend_payload(
+        &task_database,
+        request.frontend_caller,
+        request.frontend_subscriber_id.as_deref(),
+        request.frontend_subscriber_payload.as_deref(),
+      )
+      .await;
 
-      map_success_to_response(success, &app)
-    }
-    Err(err) => map_error_to_response(err),
+  if let Err(err) = db_result {
+    error!("Failed to create task in database: {:?}", err);
   }
-}
 
-// ── Result mapping ──
-
-fn map_success_to_response(
-  success: TaskEnqueueSuccess,
-  app: &AppHandle,
-) -> Response<TauriGenerateImageResponse, TauriGenerateImageErrorType, ()> {
   let event = GenerationEnqueueSuccessEvent {
     action: success.to_frontend_event_action(),
     service: success.to_frontend_event_service(),
     model: success.model,
   };
 
-  if let Err(err) = event.send(app) {
+  if let Err(err) = event.send(&app) {
     error!("Failed to emit event: {:?}", err);
   }
 
-  CreditsBalanceChangedEvent{}.send_infallible(app);
+  CreditsBalanceChangedEvent{}.send_infallible(&app);
 
   Ok(TauriGenerateImageResponse {}.into())
 }
 
+// ── Result mapping ──
+
+
 fn map_error_to_response(
   err: GenerateError,
 ) -> Response<TauriGenerateImageResponse, TauriGenerateImageErrorType, ()> {

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft.rs

@@ -0,0 +1,40 @@
+use log::info;
+
+use crate::core::commands::enqueue::generate_error::GenerateError;
+use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
+use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft_via_legacy::handle_artcraft_via_legacy;
+use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft_via_omni_endpoint::handle_artcraft_via_omni_endpoint;
+use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+
+/// Dispatch an image generation request to Artcraft.
+///
+/// Tries the omni endpoint first. If the model isn't supported there,
+/// falls back to the legacy artcraft_router path.
+pub async fn handle_artcraft(
+  request: &TauriGenerateImageRequest,
+  app_env_configs: &AppEnvConfigs,
+  storyteller_creds_manager: &StorytellerCredentialManager,
+) -> Result<TaskEnqueueSuccess, GenerateError> {
+  // Try the omni endpoint first.
+  let omni_result = handle_artcraft_via_omni_endpoint(
+    request,
+    app_env_configs,
+    storyteller_creds_manager,
+  ).await;
+
+  match omni_result {
+    Ok(success) => Ok(success),
+    Err(GenerateError::NotYetImplemented(_)) => {
+      // Model not supported by omni endpoint — fall back to legacy.
+      info!("Model not supported by omni endpoint, falling back to legacy artcraft_router path.");
+      handle_artcraft_via_legacy(
+        request,
+        app_env_configs,
+        storyteller_creds_manager,
+      ).await
+    }
+    Err(err) => Err(err),
+  }
+}

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_legacy.rs

@@ -0,0 +1,95 @@
+use artcraft_router::api::common_image_model::CommonImageModel;
+use artcraft_router::api::image_list_ref::ImageListRef;
+use artcraft_router::api::provider::Provider;
+use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
+use artcraft_router::client::router_artcraft_client::RouterArtcraftClient;
+use artcraft_router::client::router_client::RouterClient;
+use artcraft_router::generate::generate_image::generate_image_request::GenerateImageRequest;
+use enums::common::generation_provider::GenerationProvider;
+use enums::tauri::tasks::task_type::TaskType;
+use log::{error, info};
+
+use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason};
+use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
+use crate::core::commands::generate::generate_image::providers::artcraft::model_mapping::{
+  map_to_generation_model, map_to_router_image_model,
+};
+use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+
+/// Handle image generation via the legacy artcraft_router path.
+///
+/// This is used for models not yet supported by the omni endpoint
+/// (eg. Recraft, FluxProKontextMax, edit/inpaint models).
+pub async fn handle_artcraft_via_legacy(
+  request: &TauriGenerateImageRequest,
+  app_env_configs: &AppEnvConfigs,
+  storyteller_creds_manager: &StorytellerCredentialManager,
+) -> Result<TaskEnqueueSuccess, GenerateError> {
+  let model = request.model.ok_or(GenerateError::no_model_specified())?;
+
+  let router_model = map_to_router_image_model(model)
+    .ok_or(GenerateError::NotYetImplemented(
+      format!("Model {:?} is not supported via the legacy router path", model),
+    ))?;
+
+  let generation_model = map_to_generation_model(model);
+
+  let creds = match storyteller_creds_manager.get_credentials()? {
+    Some(creds) => creds,
+    None => return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsStorytellerCredentials)),
+  };
+
+  let client = RouterClient::Artcraft(RouterArtcraftClient::new(
+    app_env_configs.storyteller_host.clone(),
+    creds,
+  ));
+
+  // Build image inputs: prepend canvas image if present.
+  let mut image_tokens = request.image_media_tokens.clone().unwrap_or_default();
+  if let Some(canvas_token) = &request.canvas_image_media_token {
+    image_tokens.insert(0, canvas_token.clone());
+  }
+  let image_inputs = if image_tokens.is_empty() { None } else { Some(ImageListRef::MediaFileTokens(image_tokens)) };
+
+  let router_request = GenerateImageRequest {
+    model: router_model,
+    provider: Provider::Artcraft,
+    prompt: request.prompt.clone(),
+    image_inputs,
+    resolution: None, // TODO: Convert from request.resolution
+    aspect_ratio: None, // TODO: Convert from request.aspect_ratio
+    quality: None, // TODO: Convert from request.quality
+    image_batch_count: request.batch_size.map(|n| n as u16),
+    request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade,
+    generation_mode_mismatch_strategy: None,
+    idempotency_token: None,
+    horizontal_angle: request.adjust_horizontal_angle,
+    vertical_angle: request.adjust_vertical_angle,
+    zoom: request.adjust_zoom,
+  };
+
+  let plan = router_request.build()?;
+
+  info!("Legacy image generation plan: {:?}", plan);
+
+  let response = plan.generate_image(&client).await.map_err(|err| {
+    error!("Legacy image generation failed: {:?}", err);
+    GenerateError::from(err)
+  })?;
+
+  let job_id = response
+    .get_artcraft_payload()
+    .map(|p| p.inference_job_token.to_string())
+    .ok_or(GenerateError::ResponseHadNoJobTokens)?;
+
+  info!("Legacy image generation succeeded: job_id={}", job_id);
+
+  Ok(TaskEnqueueSuccess {
+    task_type: TaskType::ImageGeneration,
+    model: Some(generation_model),
+    provider: GenerationProvider::Artcraft,
+    provider_job_id: Some(job_id),
+  })
+}

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs

@@ -1,15 +1,15 @@
 use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest;
 use artcraft_client::endpoints::omni_gen::generate::image::omni_gen_image::omni_gen_image_generate;
-use enums::common::generation::common_image_model::CommonImageModel;
 use enums::common::generation_provider::GenerationProvider;
 use enums::tauri::tasks::task_type::TaskType;
 use log::{error, info};
 
 use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason};
 use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
+use crate::core::commands::generate::generate_image::providers::artcraft::model_mapping::{
+  map_to_generation_model, map_to_omni_image_model,
+};
 use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
-use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel;
-use crate::core::events::generation_events::common::GenerationModel;
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
 
@@ -20,7 +20,7 @@ pub async fn handle_artcraft_via_omni_endpoint(
 ) -> Result<TaskEnqueueSuccess, GenerateError> {
   let model = request.model.ok_or(GenerateError::no_model_specified())?;
 
-  let omni_model = map_to_common_image_model(model)
+  let omni_model = map_to_omni_image_model(model)
     .ok_or(GenerateError::NotYetImplemented(
       format!("Model {:?} is not supported via the omni endpoint", model),
     ))?;
@@ -32,14 +32,11 @@ pub async fn handle_artcraft_via_omni_endpoint(
     None => return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsStorytellerCredentials)),
   };
 
-  // Build the omni gen request.
+  // Build image_media_tokens: prepend canvas image if present.
   let mut image_media_tokens = request.image_media_tokens.clone().unwrap_or_default();
-
-  // Prepend canvas image if present.
   if let Some(canvas_token) = &request.canvas_image_media_token {
     image_media_tokens.insert(0, canvas_token.clone());
   }
-
   let image_media_tokens = if image_media_tokens.is_empty() { None } else { Some(image_media_tokens) };
 
   let omni_request = OmniGenImageCostAndGenerateRequest {
@@ -47,8 +44,8 @@ pub async fn handle_artcraft_via_omni_endpoint(
     model: Some(omni_model),
     prompt: request.prompt.clone(),
     image_media_tokens,
-    resolution: None, // TODO: Convert from request.resolution
-    aspect_ratio: None, // TODO: Convert from request.aspect_ratio
+    resolution: None,
+    aspect_ratio: None,
     quality: request.quality,
     image_batch_count: request.batch_size.map(|n| n as u16),
     adjust_horizontal_angle: request.adjust_horizontal_angle,
@@ -76,54 +73,3 @@ pub async fn handle_artcraft_via_omni_endpoint(
     provider_job_id: Some(response.inference_job_token.to_string()),
   })
 }
-
-// ── Model mapping helpers ──
-
-fn map_to_common_image_model(model: TauriImageModel) -> Option<CommonImageModel> {
-  match model {
-    TauriImageModel::Flux1Dev => Some(CommonImageModel::Flux1Dev),
-    TauriImageModel::Flux1Schnell => Some(CommonImageModel::Flux1Schnell),
-    TauriImageModel::FluxPro11 => Some(CommonImageModel::FluxPro11),
-    TauriImageModel::FluxPro11Ultra => Some(CommonImageModel::FluxPro11Ultra),
-    TauriImageModel::Gemini25Flash => Some(CommonImageModel::NanoBanana), // NB: Alias for Nano Banana
-    TauriImageModel::GptImage1 => Some(CommonImageModel::GptImage1),
-    TauriImageModel::GptImage1p5 => Some(CommonImageModel::GptImage1p5),
-    TauriImageModel::GptImage2 => Some(CommonImageModel::GptImage2),
-    TauriImageModel::NanoBanana => Some(CommonImageModel::NanoBanana),
-    TauriImageModel::NanoBanana2 => Some(CommonImageModel::NanoBanana2),
-    TauriImageModel::NanoBananaPro => Some(CommonImageModel::NanoBananaPro),
-    TauriImageModel::Seedream4 => Some(CommonImageModel::Seedream4),
-    TauriImageModel::Seedream4p5 => Some(CommonImageModel::Seedream4p5),
-    TauriImageModel::Seedream5Lite => Some(CommonImageModel::Seedream5Lite),
-    // Models not in CommonImageModel (Grok, Recraft, Midjourney, edit-only, inpaint-only)
-    _ => None,
-  }
-}
-
-fn map_to_generation_model(model: TauriImageModel) -> GenerationModel {
-  match model {
-    TauriImageModel::Flux1Dev => GenerationModel::Flux1Dev,
-    TauriImageModel::Flux1Schnell => GenerationModel::Flux1Schnell,
-    TauriImageModel::Flux2LoraAngles => GenerationModel::Flux2LoraAngles,
-    TauriImageModel::FluxDevJuggernaut => GenerationModel::FluxDevJuggernaut,
-    TauriImageModel::FluxPro1 => GenerationModel::FluxPro1,
-    TauriImageModel::FluxPro11 => GenerationModel::FluxPro11,
-    TauriImageModel::FluxPro11Ultra => GenerationModel::FluxPro11Ultra,
-    TauriImageModel::FluxProKontextMax => GenerationModel::FluxProKontextMax,
-    TauriImageModel::Gemini25Flash => GenerationModel::NanoBanana, // NB: Alias for Nano Banana
-    TauriImageModel::GptImage1 => GenerationModel::GptImage1,
-    TauriImageModel::GptImage1p5 => GenerationModel::GptImage1p5,
-    TauriImageModel::GptImage2 => GenerationModel::GptImage2,
-    TauriImageModel::GrokImage => GenerationModel::GrokImage,
-    TauriImageModel::NanoBanana => GenerationModel::NanoBanana,
-    TauriImageModel::NanoBanana2 => GenerationModel::NanoBanana2,
-    TauriImageModel::NanoBananaPro => GenerationModel::NanoBananaPro,
-    TauriImageModel::QwenEdit2511Angles => GenerationModel::QwenEdit2511Angles,
-    TauriImageModel::Seedream4 => GenerationModel::Seedream4,
-    TauriImageModel::Seedream4p5 => GenerationModel::Seedream4p5,
-    TauriImageModel::Seedream5Lite => GenerationModel::Seedream5Lite,
-    // TODO:
-    TauriImageModel::Midjourney => GenerationModel::Flux1Dev, // Fallback
-    TauriImageModel::Recraft3 => GenerationModel::Flux1Dev, // Fallback
-  }
-}

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/mod.rs

@@ -1 +1,6 @@
-pub mod handle_artcraft_via_omni_endpoint;
+mod handle_artcraft;
+mod handle_artcraft_via_legacy;
+mod handle_artcraft_via_omni_endpoint;
+mod model_mapping;
+
+pub use handle_artcraft::handle_artcraft;

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/model_mapping.rs

@@ -0,0 +1,79 @@
+use artcraft_router::api::common_image_model::CommonImageModel as RouterCommonImageModel;
+use enums::common::generation::common_image_model::CommonImageModel as OmniCommonImageModel;
+
+use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel;
+use crate::core::events::generation_events::common::GenerationModel;
+
+/// Map TauriImageModel to the omni endpoint's CommonImageModel.
+/// Returns None for models not supported by the omni endpoint.
+pub fn map_to_omni_image_model(model: TauriImageModel) -> Option<OmniCommonImageModel> {
+  match model {
+    TauriImageModel::Flux1Dev => Some(OmniCommonImageModel::Flux1Dev),
+    TauriImageModel::Flux1Schnell => Some(OmniCommonImageModel::Flux1Schnell),
+    TauriImageModel::FluxPro11 => Some(OmniCommonImageModel::FluxPro11),
+    TauriImageModel::FluxPro11Ultra => Some(OmniCommonImageModel::FluxPro11Ultra),
+    TauriImageModel::GptImage1 => Some(OmniCommonImageModel::GptImage1),
+    TauriImageModel::GptImage1p5 => Some(OmniCommonImageModel::GptImage1p5),
+    TauriImageModel::GptImage2 => Some(OmniCommonImageModel::GptImage2),
+    TauriImageModel::NanoBanana => Some(OmniCommonImageModel::NanoBanana),
+    TauriImageModel::NanoBanana2 => Some(OmniCommonImageModel::NanoBanana2),
+    TauriImageModel::NanoBananaPro => Some(OmniCommonImageModel::NanoBananaPro),
+    TauriImageModel::Gemini25Flash => Some(OmniCommonImageModel::NanoBanana),
+    TauriImageModel::Seedream4 => Some(OmniCommonImageModel::Seedream4),
+    TauriImageModel::Seedream4p5 => Some(OmniCommonImageModel::Seedream4p5),
+    TauriImageModel::Seedream5Lite => Some(OmniCommonImageModel::Seedream5Lite),
+    _ => None,
+  }
+}
+
+/// Map TauriImageModel to the artcraft_router's CommonImageModel.
+/// Returns None for models not supported by the router (Grok, Midjourney).
+pub fn map_to_router_image_model(model: TauriImageModel) -> Option<RouterCommonImageModel> {
+  match model {
+    TauriImageModel::Flux1Dev => Some(RouterCommonImageModel::Flux1Dev),
+    TauriImageModel::Flux1Schnell => Some(RouterCommonImageModel::Flux1Schnell),
+    TauriImageModel::FluxPro11 => Some(RouterCommonImageModel::FluxPro11),
+    TauriImageModel::FluxPro11Ultra => Some(RouterCommonImageModel::FluxPro11Ultra),
+    TauriImageModel::GptImage1 => Some(RouterCommonImageModel::GptImage1),
+    TauriImageModel::GptImage1p5 => Some(RouterCommonImageModel::GptImage1p5),
+    TauriImageModel::GptImage2 => Some(RouterCommonImageModel::GptImage2),
+    TauriImageModel::NanoBanana => Some(RouterCommonImageModel::NanoBanana),
+    TauriImageModel::NanoBanana2 => Some(RouterCommonImageModel::NanoBanana2),
+    TauriImageModel::NanoBananaPro => Some(RouterCommonImageModel::NanoBananaPro),
+    TauriImageModel::Gemini25Flash => Some(RouterCommonImageModel::NanoBanana),
+    TauriImageModel::Seedream4 => Some(RouterCommonImageModel::Seedream4),
+    TauriImageModel::Seedream4p5 => Some(RouterCommonImageModel::Seedream4p5),
+    TauriImageModel::Seedream5Lite => Some(RouterCommonImageModel::Seedream5Lite),
+    TauriImageModel::QwenEdit2511Angles => Some(RouterCommonImageModel::QwenEdit2511Angles),
+    TauriImageModel::Flux2LoraAngles => Some(RouterCommonImageModel::Flux2LoraAngles),
+    _ => None,
+  }
+}
+
+/// Map TauriImageModel to GenerationModel for frontend events.
+pub fn map_to_generation_model(model: TauriImageModel) -> GenerationModel {
+  match model {
+    TauriImageModel::Flux1Dev => GenerationModel::Flux1Dev,
+    TauriImageModel::Flux1Schnell => GenerationModel::Flux1Schnell,
+    TauriImageModel::FluxPro11 => GenerationModel::FluxPro11,
+    TauriImageModel::FluxPro11Ultra => GenerationModel::FluxPro11Ultra,
+    TauriImageModel::GrokImage => GenerationModel::GrokImage,
+    TauriImageModel::Recraft3 => GenerationModel::Flux1Dev, // Fallback
+    TauriImageModel::GptImage1 => GenerationModel::GptImage1,
+    TauriImageModel::GptImage1p5 => GenerationModel::GptImage1p5,
+    TauriImageModel::GptImage2 => GenerationModel::GptImage2,
+    TauriImageModel::Gemini25Flash => GenerationModel::NanoBanana,
+    TauriImageModel::NanoBanana => GenerationModel::NanoBanana,
+    TauriImageModel::NanoBanana2 => GenerationModel::NanoBanana2,
+    TauriImageModel::NanoBananaPro => GenerationModel::NanoBananaPro,
+    TauriImageModel::Seedream4 => GenerationModel::Seedream4,
+    TauriImageModel::Seedream4p5 => GenerationModel::Seedream4p5,
+    TauriImageModel::Seedream5Lite => GenerationModel::Seedream5Lite,
+    TauriImageModel::Midjourney => GenerationModel::Flux1Dev, // Fallback
+    TauriImageModel::FluxProKontextMax => GenerationModel::FluxProKontextMax,
+    TauriImageModel::QwenEdit2511Angles => GenerationModel::QwenEdit2511Angles,
+    TauriImageModel::Flux2LoraAngles => GenerationModel::Flux2LoraAngles,
+    TauriImageModel::FluxDevJuggernaut => GenerationModel::FluxDevJuggernaut,
+    TauriImageModel::FluxPro1 => GenerationModel::FluxPro1,
+  }
+}