diff --git a/.sqlx/query-01a0bf33e8331c19dc73f064c54eaa863bec158ef7713450a06c75e246f5f08c.json b/.sqlx/query-01a0bf33e8331c19dc73f064c54eaa863bec158ef7713450a06c75e246f5f08c.json deleted file mode 100644 index c88c3496df..0000000000 --- a/.sqlx/query-01a0bf33e8331c19dc73f064c54eaa863bec158ef7713450a06c75e246f5f08c.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "db_name": "SQLite", - "query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)\n ", - "describe": { - "columns": [], - "parameters": { - "Right": 9 - }, - "nullable": [] - }, - "hash": "01a0bf33e8331c19dc73f064c54eaa863bec158ef7713450a06c75e246f5f08c" -} diff --git a/.sqlx/query-15a2f8712c80b72c86d9e77eb590d74ea21dd818a424267bcf5b88f9c993c300.json b/.sqlx/query-15a2f8712c80b72c86d9e77eb590d74ea21dd818a424267bcf5b88f9c993c300.json new file mode 100644 index 0000000000..e7583986d1 --- /dev/null +++ b/.sqlx/query-15a2f8712c80b72c86d9e77eb590d74ea21dd818a424267bcf5b88f9c993c300.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n queue_status_url,\n queue_response_url,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 11 + }, + "nullable": [] + }, + "hash": "15a2f8712c80b72c86d9e77eb590d74ea21dd818a424267bcf5b88f9c993c300" +} diff --git a/.sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json b/.sqlx/query-68f9fed45fc55aa2ec22dc1ed22795d0c99882917abb9b52b78735b8194c2624.json similarity index 70% rename from .sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json rename to .sqlx/query-68f9fed45fc55aa2ec22dc1ed22795d0c99882917abb9b52b78735b8194c2624.json index b400c3bf00..6e48486d74 100644 --- a/.sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json +++ b/.sqlx/query-68f9fed45fc55aa2ec22dc1ed22795d0c99882917abb9b52b78735b8194c2624.json @@ -1,6 +1,6 @@ { "db_name": "SQLite", - "query": "\n SELECT\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n FROM tasks\n WHERE\n provider = ?\n AND provider_job_id = ?\n ", + "query": "\n SELECT\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n queue_status_url,\n queue_response_url,\n frontend_subscriber_id,\n frontend_subscriber_payload\n FROM tasks\n WHERE\n provider = ?\n AND provider_job_id = ?\n ", "describe": { "columns": [ { @@ -39,14 +39,24 @@ "type_info": "Text" }, { - "name": "frontend_subscriber_id", + "name": "queue_status_url", "ordinal": 7, "type_info": "Text" }, { - "name": "frontend_subscriber_payload", + "name": "queue_response_url", "ordinal": 8, "type_info": "Text" + }, + { + "name": "frontend_subscriber_id", + "ordinal": 9, + "type_info": "Text" + }, + { + "name": "frontend_subscriber_payload", + "ordinal": 10, + "type_info": "Text" } ], "parameters": { @@ -61,8 +71,10 @@ true, true, true, + true, + true, true ] }, - "hash": "38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069" + "hash": "68f9fed45fc55aa2ec22dc1ed22795d0c99882917abb9b52b78735b8194c2624" } diff --git a/Cargo.lock b/Cargo.lock index 0e1fc85c2c..67542d1388 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -679,6 +679,7 @@ dependencies = [ "discord-rich-presence", "enums", "errors", + "fal_client", "filesys", "futures 0.3.32", "grok_client", diff --git a/_database/sql/artcraft_migrations/20250708000000_create_tasks_table.sql b/_database/sql/artcraft_migrations/20250708000000_create_tasks_table.sql index 033885a869..3ab5c8e843 100644 --- a/_database/sql/artcraft_migrations/20250708000000_create_tasks_table.sql +++ b/_database/sql/artcraft_migrations/20250708000000_create_tasks_table.sql @@ -19,6 +19,9 @@ -- tasks_v6.sqlite - added two fields: -- on_failure_type -- on_failure_message +-- tasks_v7.sqlite - added two fields (for fal): +-- queue_status_url +-- queue_response_url CREATE TABLE tasks ( -- Task auto-incrementing primary key. @@ -63,6 +66,16 @@ CREATE TABLE tasks ( -- Whether the user has dismissed the task from view. is_dismissed_by_user INTEGER NOT NULL DEFAULT 0, + -- OPTIONAL. + -- For some third party systems (eg. FAL), the URL to check the status of the + -- job in the queue. + queue_status_url TEXT, + + -- OPTIONAL. + -- For some third party systems (eg. FAL), the URL to check the + -- results (when the job completes) + queue_response_url TEXT, + -- OPTIONAL. -- When the generation completes, the batch token (if any). -- This will point to a collection of multiple files. diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_client.rs b/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_client.rs index 2bd63ae9de..77a834e876 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_client.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_client.rs @@ -9,4 +9,11 @@ impl RouterFalClient { pub fn new(api_key: FalApiKey, webhook_url: String) -> Self { Self { api_key, webhook_url } } + + pub fn new_from_raw_key(api_key: &str, webhook_url: String) -> Self { + Self { + api_key: FalApiKey::from_str(api_key), + webhook_url, + } + } } diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_webhook_optional_client.rs b/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_webhook_optional_client.rs index 6e07a5566b..17a970224e 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_webhook_optional_client.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/client/router_fal_webhook_optional_client.rs @@ -9,6 +9,9 @@ impl RouterFalWebhookOptionalClient { pub fn new(api_key: FalApiKey) -> Self { Self { api_key, webhook_url: None } } + pub fn from_str(api_key: &str) -> Self { + Self::new(FalApiKey::from_str(api_key)) + } pub fn new_with_webhook(api_key: FalApiKey, webhook_url: String) -> Self { Self { diff --git a/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs b/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs index 0cf8db90e7..39821e9dce 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs @@ -1,8 +1,13 @@ use crate::errors::client_error::ClientError; use crate::errors::download_error::DownloadError; use crate::errors::provider_error::ProviderError; +use fal_client::error::fal_error_plus::FalErrorPlus; +use seedance2pro_client::error::seedance2pro_error::Seedance2ProError; +use seedance2pro_client::error::seedance2pro_specific_api_error::Seedance2ProSpecificApiError; use std::error::Error; use std::fmt::{Display, Formatter}; +use artcraft_client::error::api_error::ApiError; +use artcraft_client::error::storyteller_error::StorytellerError; #[derive(Debug)] pub enum ArtcraftRouterError { @@ -23,6 +28,9 @@ pub enum ArtcraftRouterError { /// An error from an underlying provider. Provider(ProviderError), + + /// A billing error from an underlying provider. + ProviderBillingError(ProviderError), } impl Error for ArtcraftRouterError {} @@ -36,6 +44,7 @@ impl Display for ArtcraftRouterError { Self::UnsupportedProviderAndModelForNewApi(msg) => write!(f, "Unsupported provider/model (for new API during migration): {}", msg), Self::InvalidInput(msg) => write!(f, "Invalid input: {}", msg), Self::Provider(e) => write!(f, "Provider error: {}", e), + Self::ProviderBillingError(e) => write!(f, "Provider billing error: {}", e), } } } @@ -54,6 +63,25 @@ impl From for ArtcraftRouterError { impl From for ArtcraftRouterError { fn from(error: ProviderError) -> Self { - Self::Provider(error) + let is_billing_error = match &error { + ProviderError::Fal(FalErrorPlus::FalBillingError(_)) => true, + ProviderError::Seedance2Pro(Seedance2ProError::ApiSpecific(Seedance2ProSpecificApiError::BillingError { .. })) => true, + ProviderError::Storyteller(StorytellerError::Api(ApiError::PaymentRequired(_))) => true, + _ => false, + }; + if is_billing_error { + Self::ProviderBillingError(error) + } else { + Self::Provider(error) + } + } +} + +impl From for ArtcraftRouterError { + fn from(error: FalErrorPlus) -> Self { + match &error { + FalErrorPlus::FalBillingError(e) => ArtcraftRouterError::ProviderBillingError(ProviderError::Fal(error)), + _ => ArtcraftRouterError::Provider(ProviderError::Fal(error)), + } } } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_dev.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_dev.rs index 92241c5822..23d8383dc5 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_dev.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_dev.rs @@ -58,6 +58,8 @@ pub async fn execute_fal_flux_1_dev( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_schnell.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_schnell.rs index 735df6f1dc..17873cbc1c 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_schnell.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_1_schnell.rs @@ -58,6 +58,8 @@ pub async fn execute_fal_flux_1_schnell( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1.rs index c871e942e2..f0c3a10e7d 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1.rs @@ -35,6 +35,8 @@ pub async fn execute_fal_flux_pro_1p1( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1_ultra.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1_ultra.rs index be689f43a7..b48084b8a7 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1_ultra.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_flux_pro_1p1_ultra.rs @@ -35,6 +35,8 @@ pub async fn execute_fal_flux_pro_1p1_ultra( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1.rs index f0d734137b..a84359c50d 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1.rs @@ -67,6 +67,8 @@ pub async fn execute_fal_gpt_image_1( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1p5.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1p5.rs index 46361e36c8..6439a0b070 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1p5.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_1p5.rs @@ -67,6 +67,8 @@ pub async fn execute_fal_gpt_image_1p5( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_2.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_2.rs index 58211fdbae..6b462ddc18 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_2.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_gpt_image_2.rs @@ -64,6 +64,8 @@ pub async fn execute_fal_gpt_image_2( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana.rs index 03cf930657..85d4f56f70 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana.rs @@ -63,6 +63,8 @@ pub async fn execute_fal_nano_banana( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_2.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_2.rs index 94dc5df554..702268403b 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_2.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_2.rs @@ -67,6 +67,8 @@ pub async fn execute_fal_nano_banana_2( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_pro.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_pro.rs index 2d9711a9c0..933f53a859 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_pro.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_nano_banana_pro.rs @@ -66,6 +66,8 @@ pub async fn execute_fal_nano_banana_pro( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4.rs index d45086f673..6cc1e32f99 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4.rs @@ -61,6 +61,8 @@ pub async fn execute_fal_seedream_4( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4p5.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4p5.rs index b87cd06573..3b5287b896 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4p5.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_4p5.rs @@ -61,6 +61,8 @@ pub async fn execute_fal_seedream_4p5( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_5_lite.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_5_lite.rs index a9d435bde2..d44d565ef4 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_5_lite.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/execute/fal/generate_image_fal_seedream_5_lite.rs @@ -61,6 +61,8 @@ pub async fn execute_fal_seedream_5_lite( Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: webhook_response.request_id, gateway_request_id: webhook_response.gateway_request_id, + maybe_status_url: None, + maybe_response_url: None, maybe_outbound_request: Some(outbound_request), })) } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs index b238126474..a330c7c5a6 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs @@ -13,6 +13,12 @@ pub struct FalImageResponsePayload { pub request_id: Option, pub gateway_request_id: Option, + /// The queue status URL (for polling job progress). + pub maybe_status_url: Option, + + /// The queue response URL (for fetching completed results). + pub maybe_response_url: Option, + /// The outbound request that was sent to Fal. /// Stored as a trait object so any Request type can be captured. /// Use `format!("{:?}", ...)` or `format!("{:#?}", ...)` to print. diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_dev/request.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_dev/request.rs index a421aacd02..d410a3971c 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_dev/request.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_dev/request.rs @@ -27,6 +27,8 @@ impl FalFlux1DevRequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -36,6 +38,8 @@ impl FalFlux1DevRequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -48,6 +52,8 @@ impl FalFlux1DevRequestState { struct FalResponseIds { request_id: Option, gateway_request_id: Option, + status_url: Option, + response_url: Option, } async fn send_fal_request( @@ -57,20 +63,22 @@ async fn send_fal_request( if let Some(webhook_url) = &client.webhook_url { let response = request .send_webhook_request(&client.api_key, webhook_url) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: response.request_id, gateway_request_id: response.gateway_request_id, + status_url: None, + response_url: None, }) } else { let response = request .send_queue_request(&client.api_key) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: Some(response.request_id), gateway_request_id: None, + status_url: Some(response.status_url), + response_url: Some(response.response_url), }) } } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_schnell/request.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_schnell/request.rs index f1891f9d1b..7b5ac10fd7 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_schnell/request.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/flux_1_schnell/request.rs @@ -27,6 +27,8 @@ impl FalFlux1SchnellRequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -36,6 +38,8 @@ impl FalFlux1SchnellRequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -48,6 +52,8 @@ impl FalFlux1SchnellRequestState { struct FalResponseIds { request_id: Option, gateway_request_id: Option, + status_url: Option, + response_url: Option, } async fn send_fal_request( @@ -57,20 +63,22 @@ async fn send_fal_request( if let Some(webhook_url) = &client.webhook_url { let response = request .send_webhook_request(&client.api_key, webhook_url) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: response.request_id, gateway_request_id: response.gateway_request_id, + status_url: None, + response_url: None, }) } else { let response = request .send_queue_request(&client.api_key) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: Some(response.request_id), gateway_request_id: None, + status_url: Some(response.status_url), + response_url: Some(response.response_url), }) } } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_2/request.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_2/request.rs index 92bd5d916f..faa09a7e0b 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_2/request.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_2/request.rs @@ -27,6 +27,8 @@ impl FalNanoBanana2RequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -36,6 +38,8 @@ impl FalNanoBanana2RequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -48,6 +52,8 @@ impl FalNanoBanana2RequestState { struct FalResponseIds { request_id: Option, gateway_request_id: Option, + status_url: Option, + response_url: Option, } async fn send_fal_request( @@ -57,20 +63,22 @@ async fn send_fal_request( if let Some(webhook_url) = &client.webhook_url { let response = request .send_webhook_request(&client.api_key, webhook_url) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: response.request_id, gateway_request_id: response.gateway_request_id, + status_url: None, + response_url: None, }) } else { let response = request .send_queue_request(&client.api_key) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: Some(response.request_id), gateway_request_id: None, + status_url: Some(response.status_url), + response_url: Some(response.response_url), }) } } diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_pro/request.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_pro/request.rs index 5b7f785da0..e1cb749451 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_pro/request.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image_v2/providers/fal/nano_banana_pro/request.rs @@ -27,6 +27,8 @@ impl FalNanoBananaProRequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -36,6 +38,8 @@ impl FalNanoBananaProRequestState { Ok(GenerateImageResponse::Fal(FalImageResponsePayload { request_id: payload.request_id, gateway_request_id: payload.gateway_request_id, + maybe_status_url: payload.status_url, + maybe_response_url: payload.response_url, maybe_outbound_request: Some(outbound), })) } @@ -48,6 +52,8 @@ impl FalNanoBananaProRequestState { struct FalResponseIds { request_id: Option, gateway_request_id: Option, + status_url: Option, + response_url: Option, } /// Send a FAL request via webhook (if URL present) or queue (if not). @@ -58,20 +64,22 @@ async fn send_fal_request( if let Some(webhook_url) = &client.webhook_url { let response = request .send_webhook_request(&client.api_key, webhook_url) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: response.request_id, gateway_request_id: response.gateway_request_id, + status_url: None, + response_url: None, }) } else { let response = request .send_queue_request(&client.api_key) - .await - .map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?; + .await?; Ok(FalResponseIds { request_id: Some(response.request_id), gateway_request_id: None, + status_url: Some(response.status_url), + response_url: Some(response.response_url), }) } } diff --git a/crates/desktop/artcraft/Cargo.toml b/crates/desktop/artcraft/Cargo.toml index cec696fca5..d969c47f6b 100644 --- a/crates/desktop/artcraft/Cargo.toml +++ b/crates/desktop/artcraft/Cargo.toml @@ -33,6 +33,7 @@ build_metadata.workspace = true cookie_store.workspace = true enums.workspace = true errors.workspace = true +fal_client.workspace = true filesys.workspace = true grok_client.workspace = true uuid_utils.workspace = true diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_flux_kontext_edit.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_flux_kontext_edit.rs index 7d36dcd03f..39bfc958b5 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_flux_kontext_edit.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_flux_kontext_edit.rs @@ -124,5 +124,7 @@ pub async fn handle_artcraft_flux_kontext_edit( model: Some(GenerationModel::FluxProKontextMax), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_gpt_image_1_edit.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_gpt_image_1_edit.rs index 0a310651bd..2499208ca3 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_gpt_image_1_edit.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_artcraft_gpt_image_1_edit.rs @@ -122,5 +122,7 @@ pub async fn handle_artcraft_gpt_image_1_edit( model: Some(GenerationModel::GptImage1), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs index 076c05b4f7..5641be4206 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/artcraft/handle_image_edit_artcraft_via_router.rs @@ -109,6 +109,8 @@ pub(super) async fn handle_image_edit_artcraft_via_router( model: Some(generation_model), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_id), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/sora/handle_sora_gpt_image_1_edit.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/sora/handle_sora_gpt_image_1_edit.rs index 3ae7404283..af5994dc90 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/sora/handle_sora_gpt_image_1_edit.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_edit/sora/handle_sora_gpt_image_1_edit.rs @@ -176,5 +176,7 @@ pub async fn handle_sora_gpt_image_1_edit( model: Some(GenerationModel::GptImage1), provider: GenerationProvider::Sora, provider_job_id: Some(response.task_id.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_dev_juggernaut_inpaint.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_dev_juggernaut_inpaint.rs index d077bd1ccc..a54c354332 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_dev_juggernaut_inpaint.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_dev_juggernaut_inpaint.rs @@ -113,6 +113,8 @@ pub async fn handle_artcraft_flux_dev_juggernaut_inpaint( model: Some(GenerationModel::FluxDevJuggernaut), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageInpaintEdit, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_1_inpaint.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_1_inpaint.rs index d631d543ba..d1955c15db 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_1_inpaint.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_1_inpaint.rs @@ -112,6 +112,8 @@ pub async fn handle_artcraft_flux_pro_1_inpaint( model: Some(GenerationModel::FluxPro1), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageInpaintEdit, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_kontext_inpaint.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_kontext_inpaint.rs index 1ebbb9c2c7..cb586cc6ee 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_kontext_inpaint.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_flux_pro_kontext_inpaint.rs @@ -125,5 +125,7 @@ pub async fn handle_artcraft_flux_pro_kontext_inpaint( model: Some(GenerationModel::FluxProKontextMax), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageInpaintEdit, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_nano_banana_inpaint.rs b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_nano_banana_inpaint.rs index 3084332af3..a53d6587ad 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_nano_banana_inpaint.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/image_inpaint/artcraft/handle_artcraft_nano_banana_inpaint.rs @@ -110,5 +110,7 @@ pub async fn handle_artcraft_nano_banana_inpaint( model: Some(GenerationModel::Gemini25Flash), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageInpaintEdit, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_artcraft_gpt_image_1_text_to_image.rs b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_artcraft_gpt_image_1_text_to_image.rs index 70462ebd39..63df433fa5 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_artcraft_gpt_image_1_text_to_image.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_artcraft_gpt_image_1_text_to_image.rs @@ -108,5 +108,7 @@ pub async fn handle_artcraft_gpt_image_1_text_to_image( model: Some(GenerationModel::GptImage1), provider_job_id: Some(job_id.to_string()), task_type: TaskType::ImageGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs index df2f130589..e4d95675b2 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/artcraft/handle_text_to_image_artcraft_via_router.rs @@ -87,6 +87,8 @@ pub(super) async fn handle_text_to_image_artcraft_via_router( model: Some(generation_model), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_id), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/grok/handle_grok.rs b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/grok/handle_grok.rs index 84271c2a57..48fee7d5c7 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/grok/handle_grok.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/grok/handle_grok.rs @@ -68,6 +68,8 @@ pub async fn handle_grok( model: Some(GenerationModel::GrokImage), provider_job_id: Some(job_id), task_type: TaskType::ImageGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/midjourney/handle_midjourney.rs b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/midjourney/handle_midjourney.rs index 4eceda6d30..088837d6ef 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/midjourney/handle_midjourney.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/midjourney/handle_midjourney.rs @@ -111,6 +111,8 @@ pub async fn handle_midjourney( model: Some(GenerationModel::Midjourney), provider_job_id: Some(job_id), task_type: TaskType::ImageGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/sora/handle_gpt_image_1_sora_text_to_image.rs b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/sora/handle_gpt_image_1_sora_text_to_image.rs index 04001a9d4f..cf4f48a3b6 100644 --- a/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/sora/handle_gpt_image_1_sora_text_to_image.rs +++ b/crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/sora/handle_gpt_image_1_sora_text_to_image.rs @@ -111,6 +111,8 @@ pub async fn handle_gpt_image_1_sora_text_to_image( task_type: TaskType::ImageGeneration, model: Some(GenerationModel::GptImage1), provider_job_id: Some(response.task_id.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/common/maybe_notify_frontend_of_billing_errors.rs b/crates/desktop/artcraft/src/core/commands/enqueue/common/maybe_notify_frontend_of_billing_errors.rs index 9612551df2..f74b47f509 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/common/maybe_notify_frontend_of_billing_errors.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/common/maybe_notify_frontend_of_billing_errors.rs @@ -28,10 +28,12 @@ fn billing_error( reason: &BillingIssueReason, ) { let provider = match reason.provider { + BillingProvider::Artcraft => GenerationProvider::Artcraft, BillingProvider::Fal => GenerationProvider::Fal, + BillingProvider::Kinovi => GenerationProvider::Artcraft, // NB: We don't support Kinovi yet. BillingProvider::Midjourney => GenerationProvider::Midjourney, + BillingProvider::Muapi=> GenerationProvider::Artcraft, // NB: We don't support Muapi yet. BillingProvider::Sora => GenerationProvider::Sora, - BillingProvider::Storyteller => GenerationProvider::Artcraft, }; warn!("Billing issue with: {:?}", provider); ShowProviderBillingModalEvent::send_for_provider(provider, app); diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs index b0c164253b..2a651b864f 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs @@ -93,10 +93,12 @@ pub struct BillingIssueReason { #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] pub enum BillingProvider { + Artcraft, Fal, + Kinovi, Midjourney, + Muapi, Sora, - Storyteller, } #[derive(Debug)] @@ -230,6 +232,15 @@ impl From for GenerateError { ArtcraftRouterError::Client(e) => Self::BadInput(BadInputReason::WrongImageArguments(e.to_string())), ArtcraftRouterError::Download(e) => Self::ArtcraftRouterDownloadError(e), ArtcraftRouterError::InvalidInput(msg) => Self::BadInput(BadInputReason::WrongImageArguments(msg)), + ArtcraftRouterError::ProviderBillingError(err) => { + let provider = match err { + ProviderError::Fal(_) => BillingProvider::Fal, + ProviderError::Muapi(_) => BillingProvider::Muapi, + ProviderError::Seedance2Pro(_) => BillingProvider::Kinovi, + ProviderError::Storyteller(_) => BillingProvider::Artcraft, + }; + Self::BillingIssue(BillingIssueReason { provider }) + }, ArtcraftRouterError::Provider(ProviderError::Storyteller(e)) => Self::ProviderFailure(ProviderFailureReason::StorytellerError(e)), ArtcraftRouterError::Provider(ProviderError::Fal(_)) => Self::FalNoLongerSupported, ArtcraftRouterError::Provider(ProviderError::Muapi(_)) => Self::ArtcraftRouterNotYetSupportedProvider("muapi"), diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/artcraft/handle_generic_bg_removal_artcraft.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/artcraft/handle_generic_bg_removal_artcraft.rs index aa27501fe0..9dc96b388c 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/artcraft/handle_generic_bg_removal_artcraft.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_bg_removal/artcraft/handle_generic_bg_removal_artcraft.rs @@ -79,6 +79,8 @@ pub async fn handle_generic_bg_removal_artcraft( task_type: TaskType::BackgroundRemoval, provider_job_id: Some(success_result.inference_job_token.to_string()), model: None, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/artcraft/handle_artcraft_splat_via_router.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/artcraft/handle_artcraft_splat_via_router.rs index e3a4cbedbf..97640f380f 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/artcraft/handle_artcraft_splat_via_router.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/artcraft/handle_artcraft_splat_via_router.rs @@ -58,5 +58,7 @@ pub(super) async fn handle_artcraft_splat_via_router( model: Some(generation_model), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_id), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/handle_worldlabs_marble.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/handle_worldlabs_marble.rs index aa6a859940..7e479ac7a0 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/handle_worldlabs_marble.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/handle_worldlabs_marble.rs @@ -102,6 +102,8 @@ pub async fn handle_worldlabs_marble( model: Some(GenerationModel::WorldlabsMarble), provider_job_id: Some(world_id.0), task_type: TaskType::ImageGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/artcraft/handle_artcraft_object.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/artcraft/handle_artcraft_object.rs index 7944a7fb38..5f19e7ed15 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/artcraft/handle_artcraft_object.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_object/artcraft/handle_artcraft_object.rs @@ -152,5 +152,7 @@ pub async fn handle_artcraft_object( model: used_model, provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs index bd8f07ba45..b45cafdfa4 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs @@ -14,7 +14,8 @@ pub struct TaskEnqueueSuccess { pub model: Option, pub provider: GenerationProvider, pub provider_job_id: Option, - // TODO: We may want to change the `model` type - this has weird ownership and semantics + pub maybe_queue_status_url: Option, + pub maybe_queue_response_url: Option, } impl TaskEnqueueSuccess{ @@ -121,6 +122,8 @@ impl TaskEnqueueSuccess{ model_type, provider: self.provider, provider_job_id: self.provider_job_id.as_deref(), + queue_status_url: self.maybe_queue_status_url.as_deref(), + queue_response_url: self.maybe_queue_response_url.as_deref(), frontend_caller, frontend_subscriber_id, frontend_subscriber_payload, diff --git a/crates/desktop/artcraft/src/core/commands/generate/common/mod.rs b/crates/desktop/artcraft/src/core/commands/generate/common/mod.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs index 1fa3e6eff7..d9deb28e6b 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs @@ -1,7 +1,8 @@ use crate::core::commands::enqueue::common::notify_frontend_of_errors::notify_frontend_of_errors; -use crate::core::commands::enqueue::generate_error::GenerateError; +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::handle_artcraft; +use crate::core::commands::generate::generate_image::providers::router::handle_router::handle_router; use crate::core::commands::generate::generate_image::tauri_generate_image_request::{ TauriGenerateImageErrorType, TauriGenerateImageRequest, TauriGenerateImageResponse, }; @@ -10,6 +11,7 @@ use crate::core::commands::response::shorthand::Response; use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; use crate::core::events::functional_events::credits_balance_changed_event::CreditsBalanceChangedEvent; use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent; +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::task_database::TaskDatabase; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; @@ -24,32 +26,37 @@ pub async fn generate_image_command( app_env_configs: State<'_, AppEnvConfigs>, task_database: State<'_, TaskDatabase>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, + credential_cache: State<'_, ProviderCredentialLoadingCache>, ) -> Response { info!("generate_image_command called, request: {:?}", request); - // Only Artcraft provider is supported via this command. - // Other providers (Grok, Sora, Midjourney) should use the legacy commands. + let provider = request.provider.unwrap_or(GenerationProvider::Artcraft); - let is_artcraft_provider = matches!( - request.provider, - None | Some(GenerationProvider::Artcraft) - ); - - if !is_artcraft_provider { - return Err(CommandErrorResponseWrapper { - status: CommandErrorStatus::BadRequest, - error_message: Some("Only Artcraft is supported.".to_string()), - error_type: Some(TauriGenerateImageErrorType::NoProviderAvailable), - error_details: None, - }); - } - - let result = handle_artcraft( - &request, - &app_env_configs, - &storyteller_creds_manager, - ).await; + let result = match provider { + GenerationProvider::Artcraft => { + handle_artcraft( + &request, + &app_env_configs, + &storyteller_creds_manager, + ).await + } + // Midjourney uses its own legacy command path, not this one. + GenerationProvider::Midjourney => { + Err(GenerateError::NotYetImplemented( + "Midjourney should use its dedicated command".to_string(), + )) + } + // All other providers go through the router. + other => { + handle_router( + &request, + other, + &app_env_configs, + &credential_cache, + ).await + } + }; match result { Ok(success) => handle_success_behavior(&app, &task_database, &request, success).await, @@ -119,6 +126,8 @@ async fn handle_error_behavior( let error_type = match &err { GenerateError::BadInput(_) => TauriGenerateImageErrorType::BadInput, + GenerateError::MissingCredentials(MissingCredentialsReason::NeedsFalApiKey) => TauriGenerateImageErrorType::NeedsFalApiKey, + GenerateError::MissingCredentials(MissingCredentialsReason::NeedsGrokCredentials) => TauriGenerateImageErrorType::NeedsGrokCredentials, GenerateError::MissingCredentials(_) => TauriGenerateImageErrorType::NeedsStorytellerCredentials, GenerateError::NoProviderAvailable => TauriGenerateImageErrorType::NoProviderAvailable, GenerateError::BillingIssue(_) => TauriGenerateImageErrorType::BillingIssue, diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs index abcb05a5dc..31c0835955 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/handle_artcraft_via_omni_endpoint.rs @@ -65,6 +65,8 @@ pub async fn handle_artcraft_via_omni_endpoint( model: Some(generation_model), provider: GenerationProvider::Artcraft, provider_job_id: Some(response.inference_job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_2_lora_angles.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_2_lora_angles.rs index a92210b7b7..e8df3d4f43 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_2_lora_angles.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_2_lora_angles.rs @@ -70,6 +70,8 @@ pub async fn handle_flux_2_lora_angles( model: Some(GenerationModel::Flux2LoraAngles), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_id), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_dev_juggernaut_inpaint.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_dev_juggernaut_inpaint.rs index b6e2f90d12..ee90cd3040 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_dev_juggernaut_inpaint.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_dev_juggernaut_inpaint.rs @@ -76,6 +76,8 @@ pub async fn handle_flux_dev_juggernaut_inpaint( model: Some(GenerationModel::FluxDevJuggernaut), provider: GenerationProvider::Artcraft, provider_job_id: Some(response.inference_job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_pro_kontext_edit.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_pro_kontext_edit.rs index 12236bcc67..566acb0733 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_pro_kontext_edit.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_flux_pro_kontext_edit.rs @@ -74,6 +74,8 @@ pub async fn handle_flux_pro_kontext_edit( model: Some(GenerationModel::FluxProKontextMax), provider: GenerationProvider::Artcraft, provider_job_id: Some(response.inference_job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_qwen_edit_2511_angles.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_qwen_edit_2511_angles.rs index 774d648c78..30037464b0 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_qwen_edit_2511_angles.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft/legacy/artcraft_qwen_edit_2511_angles.rs @@ -70,6 +70,8 @@ pub async fn handle_qwen_edit_2511_angles( model: Some(GenerationModel::QwenEdit2511Angles), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_id), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/mod.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/mod.rs index 8a4247a99f..cd026817bf 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/mod.rs @@ -1 +1,2 @@ -pub mod artcraft; \ No newline at end of file +pub mod artcraft; +pub mod router; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_api_providers.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_api_providers.rs new file mode 100644 index 0000000000..b2f6534dc8 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_api_providers.rs @@ -0,0 +1,159 @@ +use artcraft_client::utils::api_host::ApiHost; +use artcraft_router::api::image_list_ref::ImageListRef; +use artcraft_router::api::provider::Provider; +use artcraft_router::client::generation_mode_mismatch_strategy::GenerationModeMismatchStrategy; +use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy; +use artcraft_router::client::router_client::RouterClient; +use artcraft_router::client::router_fal_client::RouterFalClient; +use artcraft_router::client::router_fal_webhook_optional_client::RouterFalWebhookOptionalClient; +use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; +use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; +use artcraft_router::generate::generate_image_v2::image_generation_draft_or_request::ImageGenerationDraftOrRequest; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_type::TaskType; +use log::{info, warn}; +use tokens::tokens::media_files::MediaFileToken; + +use crate::core::api_adapters::models::image::tauri_image_model_to_generation_model::tauri_image_model_to_generation_model; +use crate::core::api_adapters::models::image::tauri_image_model_to_router_model::tauri_image_model_to_router_model; +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::router::utils::convert_enums_to_router::{convert_aspect_ratio, convert_quality, convert_resolution}; +use crate::core::commands::generate::generate_image::providers::router::utils::map_media_files_to_urls::map_media_file_tokens_to_cdn_urls; +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::state::app_env_configs::app_env_configs::AppEnvConfigs; + +/// Handle image generation for providers that authenticate via API key. +pub async fn handle_api_key_provider( + request: &TauriGenerateImageRequest, + provider: GenerationProvider, + api_key: &str, + app_env_configs: &AppEnvConfigs, +) -> Result { + match provider { + GenerationProvider::Fal => { + handle_fal(request, api_key, &app_env_configs.storyteller_host).await + } + _ => { + Err(GenerateError::NotYetImplemented( + format!("API key provider {:?} is not yet supported via the router path", provider), + )) + } + } +} + +// ── FAL ── + +async fn handle_fal( + request: &TauriGenerateImageRequest, + api_key: &str, + api_host: &ApiHost, +) -> Result { + let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; + + let router_model = tauri_image_model_to_router_model(tauri_model) + .ok_or(GenerateError::NotYetImplemented( + format!("Model {:?} is not supported via the FAL router path", tauri_model), + ))?; + + // Collect all media file tokens that need resolving. + let image_inputs = resolve_image_inputs(request, api_host).await?; + + let router_request = GenerateImageRequestBuilder { + model: router_model, + provider: Provider::Fal, + prompt: request.prompt.clone(), + image_inputs, + resolution: request.resolution.map(convert_resolution), + aspect_ratio: request.aspect_ratio.map(convert_aspect_ratio), + quality: request.quality.map(convert_quality), + image_batch_count: request.batch_size.map(|n| n as u16), + horizontal_angle: request.adjust_horizontal_angle, + vertical_angle: request.adjust_vertical_angle, + zoom: request.adjust_zoom, + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, + generation_mode_mismatch_strategy: Some(GenerationModeMismatchStrategy::GenerateAnyway), + idempotency_token: None, + }; + + let fal_client = RouterFalWebhookOptionalClient::from_str(api_key); + let client = RouterClient::FalWebhookOptional(fal_client); + + info!("Building FAL image generation plan: model={:?}", router_model); + + let request = match router_request.build2() { + Ok(ImageGenerationDraftOrRequest::Request(request)) => request, + Ok(ImageGenerationDraftOrRequest::Draft(draft)) => { + warn!("Fal is trying to send draft request: {:?}", draft); + return Err(GenerateError::NotYetImplemented("Fal should not be sending draft requests".to_string())); + }, + Err(err) => { + warn!("Could not use FAL: {:?}", err); + return Err(GenerateError::NotYetImplemented("Error Message: TODO".to_string())); + } + }; + + info!("Executing FAL image generation. Request: {:?}", request); + + match request.send_request(&client).await { + Ok(response) => { + build_task_enqueue_success(tauri_model, response) + }, + Err(err) => { + warn!("Fal image generation failed: {:?}", err); + Err(err.into()) + } + } +} + +// ── Helpers ── + +async fn resolve_image_inputs( + request: &TauriGenerateImageRequest, + api_host: &ApiHost, +) -> Result, GenerateError> { + let mut tokens: Vec = Vec::new(); + + if let Some(canvas_token) = &request.canvas_image_media_token { + tokens.push(canvas_token.clone()); + } + + if let Some(scene_token) = &request.scene_image_media_token { + tokens.push(scene_token.clone()); + } + + if let Some(media_tokens) = &request.image_media_tokens { + tokens.extend(media_tokens.clone()); + } + + if tokens.is_empty() { + return Ok(None); + } + + let urls = map_media_file_tokens_to_cdn_urls(&tokens, api_host).await?; + Ok(Some(ImageListRef::Urls(urls))) +} + +fn build_task_enqueue_success( + tauri_model: TauriImageModel, + response: GenerateImageResponse, +) -> Result { + let fal_payload = response.get_fal_payload() + .ok_or(GenerateError::ResponseHadNoJobTokens)?; + + let provider_job_id = fal_payload.request_id + .or(fal_payload.gateway_request_id) + .ok_or(GenerateError::ResponseHadNoJobTokens)?; + + let generation_model = tauri_image_model_to_generation_model(tauri_model); + + Ok(TaskEnqueueSuccess { + task_type: TaskType::ImageGeneration, + model: Some(generation_model), + provider: GenerationProvider::Fal, + provider_job_id: Some(provider_job_id), + maybe_queue_status_url: fal_payload.maybe_status_url, + maybe_queue_response_url: fal_payload.maybe_response_url, + }) +} diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_router.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_router.rs new file mode 100644 index 0000000000..e0222d335c --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_router.rs @@ -0,0 +1,66 @@ +use enums::common::generation_provider::GenerationProvider; +use log::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::router::handle_api_providers::handle_api_key_provider; +use crate::core::commands::generate::generate_image::providers::router::handle_web_login_providers::handle_web_login_provider; +use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest; +use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload; +use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey; +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; + +/// Dispatch an image generation request to a third-party provider via the router. +/// +/// This handles providers that authenticate with the user's own credentials +/// (API key or web login) rather than going through the Artcraft backend. +pub async fn handle_router( + request: &TauriGenerateImageRequest, + provider: GenerationProvider, + app_env_configs: &AppEnvConfigs, + credential_cache: &ProviderCredentialLoadingCache, +) -> Result { + let credential_key = map_provider_to_credential_key(provider)?; + + info!("handle_router: provider={:?}, credential_key={:?}", provider, credential_key); + + let payload = credential_cache.get_credentials(credential_key) + .map_err(|err| { + GenerateError::AnyhowError(anyhow::anyhow!("Failed to load credentials: {:?}", err)) + })? + .ok_or_else(|| map_provider_to_missing_credentials_error(provider))?; + + match payload { + ProviderCredentialPayload::ApiKey(api_key_data) => { + handle_api_key_provider(request, provider, api_key_data.as_str(), app_env_configs).await + } + ProviderCredentialPayload::WebLogin(web_login_data) => { + handle_web_login_provider(request, provider, &web_login_data).await + } + } +} + +// ── Helpers ── + +fn map_provider_to_credential_key( + provider: GenerationProvider, +) -> Result { + match provider { + GenerationProvider::Fal => Ok(ProviderCredentialKey::FalApiKey), + _ => Err(GenerateError::NotYetImplemented( + format!("Provider {:?} does not have a mapped credential key yet", provider), + )), + } +} + +fn map_provider_to_missing_credentials_error(provider: GenerationProvider) -> GenerateError { + match provider { + GenerationProvider::Fal => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsFalApiKey), + GenerationProvider::Grok => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsGrokCredentials), + GenerationProvider::Midjourney => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyCredentials), + GenerationProvider::Sora => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsSoraCredentials), + GenerationProvider::WorldLabs => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsWorldLabsCredentials), + _ => GenerateError::NoProviderAvailable, + } +} diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_web_login_providers.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_web_login_providers.rs new file mode 100644 index 0000000000..d9359f325f --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/handle_web_login_providers.rs @@ -0,0 +1,17 @@ +use enums::common::generation_provider::GenerationProvider; + +use crate::core::commands::enqueue::generate_error::GenerateError; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest; +use crate::core::providers::credentials::payload::web_login::WebLoginData; + +/// Handle image generation for providers that authenticate via web login. +/// +/// Not yet implemented — placeholder for Grok, Runway, Higgsfield, etc. +pub async fn handle_web_login_provider( + _request: &TauriGenerateImageRequest, + _provider: GenerationProvider, + _web_login: &WebLoginData, +) -> Result { + unimplemented!("Web login provider image generation is not yet supported via the router path") +} diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/mod.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/mod.rs new file mode 100644 index 0000000000..ea74c39d9a --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/mod.rs @@ -0,0 +1,4 @@ +pub mod handle_router; +mod handle_api_providers; +mod handle_web_login_providers; +mod utils; diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/convert_enums_to_router.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/convert_enums_to_router.rs new file mode 100644 index 0000000000..f81bbd5779 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/convert_enums_to_router.rs @@ -0,0 +1,50 @@ +use artcraft_router::api::common_aspect_ratio::CommonAspectRatio as RouterAspectRatio; +use artcraft_router::api::common_quality::CommonQuality as RouterQuality; +use artcraft_router::api::common_resolution::CommonResolution as RouterResolution; +use enums::common::generation::common_aspect_ratio::CommonAspectRatio as EnumsAspectRatio; +use enums::common::generation::common_quality::CommonQuality as EnumsQuality; +use enums::common::generation::common_resolution::CommonResolution as EnumsResolution; + +pub fn convert_aspect_ratio(value: EnumsAspectRatio) -> RouterAspectRatio { + match value { + EnumsAspectRatio::Auto => RouterAspectRatio::Auto, + EnumsAspectRatio::Square => RouterAspectRatio::Square, + EnumsAspectRatio::WideThreeByTwo => RouterAspectRatio::WideThreeByTwo, + EnumsAspectRatio::WideFourByThree => RouterAspectRatio::WideFourByThree, + EnumsAspectRatio::WideFiveByFour => RouterAspectRatio::WideFiveByFour, + EnumsAspectRatio::WideSixteenByNine => RouterAspectRatio::WideSixteenByNine, + EnumsAspectRatio::WideTwentyOneByNine => RouterAspectRatio::WideTwentyOneByNine, + EnumsAspectRatio::TallTwoByThree => RouterAspectRatio::TallTwoByThree, + EnumsAspectRatio::TallThreeByFour => RouterAspectRatio::TallThreeByFour, + EnumsAspectRatio::TallFourByFive => RouterAspectRatio::TallFourByFive, + EnumsAspectRatio::TallNineBySixteen => RouterAspectRatio::TallNineBySixteen, + EnumsAspectRatio::TallNineByTwentyOne => RouterAspectRatio::TallNineByTwentyOne, + EnumsAspectRatio::Wide => RouterAspectRatio::Wide, + EnumsAspectRatio::Tall => RouterAspectRatio::Tall, + EnumsAspectRatio::Auto2k => RouterAspectRatio::Auto2k, + EnumsAspectRatio::Auto3k => RouterAspectRatio::Auto3k, + EnumsAspectRatio::Auto4k => RouterAspectRatio::Auto4k, + EnumsAspectRatio::SquareHd => RouterAspectRatio::SquareHd, + } +} + +pub fn convert_quality(value: EnumsQuality) -> RouterQuality { + match value { + EnumsQuality::High => RouterQuality::High, + EnumsQuality::Medium => RouterQuality::Medium, + EnumsQuality::Low => RouterQuality::Low, + } +} + +pub fn convert_resolution(value: EnumsResolution) -> RouterResolution { + match value { + EnumsResolution::OneK => RouterResolution::OneK, + EnumsResolution::TwoK => RouterResolution::TwoK, + EnumsResolution::ThreeK => RouterResolution::ThreeK, + EnumsResolution::FourK => RouterResolution::FourK, + EnumsResolution::HalfK => RouterResolution::HalfK, + EnumsResolution::FourEightyP => RouterResolution::FourEightyP, + EnumsResolution::SevenTwentyP => RouterResolution::SevenTwentyP, + EnumsResolution::TenEightyP => RouterResolution::TenEightyP, + } +} diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/map_media_files_to_urls.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/map_media_files_to_urls.rs new file mode 100644 index 0000000000..81f47670e7 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/map_media_files_to_urls.rs @@ -0,0 +1,27 @@ +use artcraft_client::endpoints::media_files::get_media_file::get_media_file; +use artcraft_client::utils::api_host::ApiHost; +use log::info; +use tokens::tokens::media_files::MediaFileToken; + +use crate::core::commands::enqueue::generate_error::GenerateError; + +/// Resolve a list of media file tokens to their CDN URLs. +/// +/// FAL (and other third-party providers) can fetch images by URL directly, +/// so we just need the CDN link — no need to download the bytes locally. +pub async fn map_media_file_tokens_to_cdn_urls( + tokens: &[MediaFileToken], + api_host: &ApiHost, +) -> Result, GenerateError> { + let mut urls = Vec::with_capacity(tokens.len()); + + for token in tokens { + info!("Resolving media file token {} to CDN URL...", token.as_str()); + let response = get_media_file(api_host, token).await?; + let cdn_url = response.media_file.media_links.cdn_url.to_string(); + info!("Resolved {} -> {}", token.as_str(), &cdn_url); + urls.push(cdn_url); + } + + Ok(urls) +} diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/mod.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/mod.rs new file mode 100644 index 0000000000..d9e43ac9c4 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/router/utils/mod.rs @@ -0,0 +1,2 @@ +pub mod convert_enums_to_router; +pub mod map_media_files_to_urls; diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs index 1a4efafaea..1e37568ff6 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs @@ -116,6 +116,8 @@ pub enum TauriGenerateImageErrorType { ServerError, /// Needs to be logged into Artcraft NeedsStorytellerCredentials, + /// Needs a FAL API key + NeedsFalApiKey, /// Needs Grok credentials NeedsGrokCredentials, /// Billing issue diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_1p6_pro.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_1p6_pro.rs index ca84b1f8a5..b1215731d6 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_1p6_pro.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_1p6_pro.rs @@ -71,5 +71,7 @@ pub (super) async fn handle_artcraft_kling_1p6_pro( model: Some(GenerationModel::Kling1_6), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_master.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_master.rs index eb443f03fc..8d08027003 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_master.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_master.rs @@ -70,5 +70,7 @@ pub (super) async fn handle_artcraft_kling_2p1_master( model: Some(GenerationModel::Kling21Master), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_pro.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_pro.rs index c085324560..dc504a2e4d 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_pro.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p1_pro.rs @@ -70,5 +70,7 @@ pub (super) async fn handle_artcraft_kling_2p1_pro( model: Some(GenerationModel::Kling21Pro), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p5_turbo_pro.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p5_turbo_pro.rs index 032489b319..254fed394f 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p5_turbo_pro.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p5_turbo_pro.rs @@ -70,5 +70,7 @@ pub (super) async fn handle_artcraft_kling_2p5_turbo_pro( model: Some(GenerationModel::Kling2p5TurboPro), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p6_pro.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p6_pro.rs index 783fd9e48d..37fddcd4a6 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p6_pro.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_kling_2p6_pro.rs @@ -72,5 +72,7 @@ pub (super) async fn handle_artcraft_kling_2p6_pro( model: Some(GenerationModel::Kling2p6Pro), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_seedance_1_lite.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_seedance_1_lite.rs index ac5cc62f49..58b071a275 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_seedance_1_lite.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_seedance_1_lite.rs @@ -66,5 +66,7 @@ pub (super) async fn handle_artcraft_seedance_1_lite( model: Some(GenerationModel::Kling1_6), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2.rs index f19ed8e319..b686b294b3 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2.rs @@ -72,5 +72,7 @@ pub (super) async fn handle_artcraft_sora2( model: Some(GenerationModel::Sora2), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2_pro.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2_pro.rs index da71f50a58..344db5c873 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2_pro.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_sora2_pro.rs @@ -71,5 +71,7 @@ pub (super) async fn handle_artcraft_sora2_pro( model: Some(GenerationModel::Sora2Pro), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo2.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo2.rs index 9dd9b1944f..2994130fe7 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo2.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo2.rs @@ -69,5 +69,7 @@ pub (super) async fn handle_artcraft_veo2( model: Some(GenerationModel::Veo2), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3.rs index ff63be0f92..632c8e2fb3 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3.rs @@ -71,5 +71,7 @@ pub (super) async fn handle_artcraft_veo3( model: Some(GenerationModel::Veo3), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3_fast.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3_fast.rs index 9fe34fc7ee..5558295787 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3_fast.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3_fast.rs @@ -66,5 +66,7 @@ pub (super) async fn handle_artcraft_veo3_fast( model: Some(GenerationModel::Veo3Fast), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1.rs index 2636c02f84..df7683d225 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1.rs @@ -74,5 +74,7 @@ pub (super) async fn handle_artcraft_veo3p1( model: Some(GenerationModel::Veo3p1), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1_fast.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1_fast.rs index ab7dfe7591..ca912d649a 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1_fast.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_veo3p1_fast.rs @@ -74,5 +74,7 @@ pub (super) async fn handle_artcraft_veo3p1_fast( model: Some(GenerationModel::Veo3p1Fast), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_token.to_string()), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs index 2fed730e18..ca4c0a6a08 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs @@ -79,6 +79,8 @@ pub(super) async fn handle_artcraft_video_via_router( model: Some(generation_model), provider: GenerationProvider::Artcraft, provider_job_id: Some(job_id), + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/grok/handle_grok_video.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/grok/handle_grok_video.rs index 6656f0625f..c38543eb98 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/grok/handle_grok_video.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/grok/handle_grok_video.rs @@ -96,6 +96,8 @@ pub async fn handle_grok_video( model: Some(GenerationModel::GrokVideo), provider_job_id: Some(post_id.to_string()), task_type: TaskType::VideoGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/generate_video/sora2/handle_sora_sora2.rs b/crates/desktop/artcraft/src/core/commands/generate/generate_video/sora2/handle_sora_sora2.rs index 8520b4dfa7..a351dacdb0 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/generate_video/sora2/handle_sora_sora2.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/generate_video/sora2/handle_sora_sora2.rs @@ -98,5 +98,7 @@ pub async fn handle_sora_sora2( model: Some(GenerationModel::Sora2), provider_job_id: Some(job_id.to_string()), task_type: TaskType::VideoGeneration, + maybe_queue_status_url: None, + maybe_queue_response_url: None, }) } diff --git a/crates/desktop/artcraft/src/core/commands/generate/mod.rs b/crates/desktop/artcraft/src/core/commands/generate/mod.rs index 295de01e41..b59be414cd 100644 --- a/crates/desktop/artcraft/src/core/commands/generate/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/generate/mod.rs @@ -1,2 +1,3 @@ pub mod generate_image; pub mod generate_video; +pub mod common; diff --git a/crates/desktop/artcraft/src/core/commands/providers/get_provider_order_command.rs b/crates/desktop/artcraft/src/core/commands/providers/deprecated/get_provider_order_command.rs similarity index 100% rename from crates/desktop/artcraft/src/core/commands/providers/get_provider_order_command.rs rename to crates/desktop/artcraft/src/core/commands/providers/deprecated/get_provider_order_command.rs diff --git a/crates/desktop/artcraft/src/core/commands/providers/deprecated/mod.rs b/crates/desktop/artcraft/src/core/commands/providers/deprecated/mod.rs new file mode 100644 index 0000000000..83c74a6411 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/providers/deprecated/mod.rs @@ -0,0 +1,2 @@ +pub mod get_provider_order_command; +pub mod set_provider_order_command; diff --git a/crates/desktop/artcraft/src/core/commands/providers/set_provider_order_command.rs b/crates/desktop/artcraft/src/core/commands/providers/deprecated/set_provider_order_command.rs similarity index 100% rename from crates/desktop/artcraft/src/core/commands/providers/set_provider_order_command.rs rename to crates/desktop/artcraft/src/core/commands/providers/deprecated/set_provider_order_command.rs diff --git a/crates/desktop/artcraft/src/core/commands/providers/mod.rs b/crates/desktop/artcraft/src/core/commands/providers/mod.rs index e3a06fefa8..74b32e1655 100644 --- a/crates/desktop/artcraft/src/core/commands/providers/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/providers/mod.rs @@ -1,2 +1,4 @@ -pub mod get_provider_order_command; -pub mod set_provider_order_command; \ No newline at end of file +pub mod deprecated; +pub mod provider_list_command; +pub mod provider_set_api_key_command; +pub mod provider_clear_command; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/providers/provider_clear_command.rs b/crates/desktop/artcraft/src/core/commands/providers/provider_clear_command.rs new file mode 100644 index 0000000000..150a7a2dda --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/providers/provider_clear_command.rs @@ -0,0 +1,33 @@ +use crate::core::commands::response::shorthand::ResponseOrErrorMessage; +use crate::core::commands::response::success_response_wrapper::SerializeMarker; +use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey; +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; +use log::{error, info}; +use serde_derive::{Deserialize, Serialize}; +use tauri::State; + +#[derive(Deserialize)] +pub struct ProviderClearRequest { + pub provider_credential: ProviderCredentialKey, +} + +#[derive(Serialize)] +pub struct ProviderClearResponse {} + +impl SerializeMarker for ProviderClearResponse {} + +#[tauri::command] +pub async fn provider_clear_command( + request: ProviderClearRequest, + credential_cache: State<'_, ProviderCredentialLoadingCache>, +) -> ResponseOrErrorMessage { + info!("provider_clear_command called for key: {:?}", request.provider_credential); + + credential_cache.delete_credentials(request.provider_credential) + .map_err(|err| { + error!("Failed to clear provider credential: {:?}", err); + "Failed to clear provider credential" + })?; + + Ok(ProviderClearResponse {}.into()) +} diff --git a/crates/desktop/artcraft/src/core/commands/providers/provider_list_command.rs b/crates/desktop/artcraft/src/core/commands/providers/provider_list_command.rs new file mode 100644 index 0000000000..55eee65b61 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/providers/provider_list_command.rs @@ -0,0 +1,106 @@ +use crate::core::commands::response::shorthand::ResponseOrErrorMessage; +use crate::core::commands::response::success_response_wrapper::SerializeMarker; +use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload; +use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey; +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; +use crate::core::providers::credentials::provider_credential_type::ProviderCredentialType; +use log::{info, warn}; +use serde_derive::Serialize; +use tauri::State; + +/// All known credential keys. Add new ones here as providers are added. +const ALL_KEYS: &[ProviderCredentialKey] = &[ + ProviderCredentialKey::FalApiKey, + ProviderCredentialKey::ReplicateApiKey, + ProviderCredentialKey::GrokWebLogin, + ProviderCredentialKey::HiggsfieldWebLogin, + ProviderCredentialKey::MidjourneyLogin, + ProviderCredentialKey::RunwayWebLogin, +]; + +const REDACTED_KEY_VISIBLE_CHARS: usize = 6; + +#[derive(Serialize)] +pub struct ProviderListResponse { + pub providers: Vec, +} + +impl SerializeMarker for ProviderListResponse {} + +#[derive(Serialize)] +pub struct ProviderListEntry { + pub provider_credential: ProviderCredentialKey, + pub credential_type: ProviderCredentialType, + pub has_credentials: bool, + pub maybe_details: Option, +} + +#[derive(Serialize)] +pub struct ProviderCredentialDetails { + /// For API keys: the first few characters followed by asterisks. + pub maybe_key_start: Option, + /// For API keys: the full key value. + pub maybe_full_key: Option, + /// For web logins: the email address if available. + pub maybe_email_address: Option, + /// For web logins: the username if available. + pub maybe_username: Option, +} + +#[tauri::command] +pub async fn provider_list_command( + credential_cache: State<'_, ProviderCredentialLoadingCache>, +) -> ResponseOrErrorMessage { + info!("provider_list_command called"); + + let mut providers = Vec::with_capacity(ALL_KEYS.len()); + + for &key in ALL_KEYS { + let maybe_payload = match credential_cache.get_credentials(key) { + Ok(payload) => payload, + Err(err) => { + warn!("Error checking credential for {:?}: {:?}", key, err); + None + } + }; + + let has_credentials = maybe_payload.is_some(); + + let maybe_details = maybe_payload.map(|payload| match payload { + ProviderCredentialPayload::ApiKey(data) => { + ProviderCredentialDetails { + maybe_key_start: Some(redact_key(data.as_str())), + maybe_full_key: Some(data.as_str().to_string()), + maybe_email_address: None, + maybe_username: None, + } + } + ProviderCredentialPayload::WebLogin(data) => { + ProviderCredentialDetails { + maybe_key_start: None, + maybe_full_key: None, + maybe_email_address: data.email_address, + maybe_username: data.username, + } + } + }); + + providers.push(ProviderListEntry { + provider_credential: key, + credential_type: key.get_type(), + has_credentials, + maybe_details, + }); + } + + Ok(ProviderListResponse { providers }.into()) +} + +fn redact_key(key: &str) -> String { + if key.len() <= REDACTED_KEY_VISIBLE_CHARS { + "********".to_string() + } else { + let visible = &key[..REDACTED_KEY_VISIBLE_CHARS]; + format!("{}********", visible) + } +} diff --git a/crates/desktop/artcraft/src/core/commands/providers/provider_set_api_key_command.rs b/crates/desktop/artcraft/src/core/commands/providers/provider_set_api_key_command.rs new file mode 100644 index 0000000000..fc08ce0f2d --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/providers/provider_set_api_key_command.rs @@ -0,0 +1,39 @@ +use crate::core::commands::response::shorthand::ResponseOrErrorMessage; +use crate::core::commands::response::success_response_wrapper::SerializeMarker; +use crate::core::providers::credentials::payload::api_key::ApiKeyData; +use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload; +use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey; +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; +use log::{error, info}; +use serde_derive::{Deserialize, Serialize}; +use tauri::State; + +#[derive(Deserialize)] +pub struct ProviderSetApiKeyRequest { + pub provider_credential: ProviderCredentialKey, + pub api_key: String, +} + +#[derive(Serialize)] +pub struct ProviderSetApiKeyResponse {} + +impl SerializeMarker for ProviderSetApiKeyResponse {} + +#[tauri::command] +pub async fn provider_set_api_key_command( + request: ProviderSetApiKeyRequest, + credential_cache: State<'_, ProviderCredentialLoadingCache>, +) -> ResponseOrErrorMessage { + info!("provider_set_api_key_command called for provider: {:?}", request.provider_credential); + + let api_key_data = ApiKeyData::from_str(&request.api_key); + let payload = ProviderCredentialPayload::ApiKey(api_key_data); + + credential_cache.save_credentials(request.provider_credential, payload) + .map_err(|err| { + error!("Failed to save API key: {:?}", err); + "Failed to save API key" + })?; + + Ok(ProviderSetApiKeyResponse {}.into()) +} diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs index 0ec319b02f..033ea2236e 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs @@ -21,6 +21,8 @@ use crate::services::sora::state::sora_task_queue::SoraTaskQueue; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge; use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; +use crate::core::threads::third_party_task_polling_thread::third_party_task_polling_thread::third_party_task_polling_thread; use crate::services::worldlabs::threads::worldlabs_marble_task_polling::worldlabs_marble_task_polling; use errors::AnyhowResult; use tauri::{AppHandle, Manager}; @@ -39,6 +41,7 @@ pub async fn handle_tauri_startup( grok_image_prompt_queue: GrokImagePromptQueue, worldlabs_bearer_bridge: WorldlabsBearerBridge, worldlabs_creds_manager: WorldlabsCredentialManager, + credential_cache: ProviderCredentialLoadingCache, ) -> AnyhowResult<()> { set_app_log_level( @@ -116,6 +119,15 @@ pub async fn handle_tauri_startup( storyteller_creds_manager.clone(), )); + tauri::async_runtime::spawn(third_party_task_polling_thread( + app.clone(), + app_env_configs.clone(), + root.clone(), + task_database.clone(), + storyteller_creds_manager.clone(), + credential_cache, + )); + spawn_discord_presence_thread()?; initially_size_and_position_windows(&app, &root); diff --git a/crates/desktop/artcraft/src/core/mod.rs b/crates/desktop/artcraft/src/core/mod.rs index 810a551526..cbdc9bd9fb 100644 --- a/crates/desktop/artcraft/src/core/mod.rs +++ b/crates/desktop/artcraft/src/core/mod.rs @@ -3,6 +3,7 @@ pub mod artcraft_error; pub mod commands; pub mod events; pub mod lifecycle; +pub mod providers; pub mod state; pub mod threads; pub mod utils; diff --git a/crates/desktop/artcraft/src/core/providers/credentials/mod.rs b/crates/desktop/artcraft/src/core/providers/credentials/mod.rs new file mode 100644 index 0000000000..1af6c8b918 --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/mod.rs @@ -0,0 +1,4 @@ +pub mod payload; +pub mod provider_credential_key; +pub mod provider_credential_loading_cache; +pub mod provider_credential_type; diff --git a/crates/desktop/artcraft/src/core/providers/credentials/payload/api_key.rs b/crates/desktop/artcraft/src/core/providers/credentials/payload/api_key.rs new file mode 100644 index 0000000000..c00bfedcde --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/payload/api_key.rs @@ -0,0 +1,100 @@ +use std::error::Error; +use std::fmt::Display; +use std::path::Path; + +#[derive(Clone)] +pub struct ApiKeyData(String); + +impl ApiKeyData { + pub fn from_str(s: &str) -> Self { + Self(s.trim().to_string()) + } + + pub fn as_str(&self) -> &str { + self.0.as_str() + } + + pub fn load_from_file>(file_path: P) -> Result { + let contents = std::fs::read_to_string(file_path) + .map_err(ApiKeyDataError::IoError)?; + Ok(Self(contents.trim().to_string())) + } + + pub fn save_to_file>(&self, file_path: P) -> Result<(), ApiKeyDataError> { + std::fs::write(file_path, self.0.trim()) + .map_err(ApiKeyDataError::IoError)?; + Ok(()) + } +} + +#[derive(Debug)] +pub enum ApiKeyDataError { + IoError(std::io::Error), +} + +impl Error for ApiKeyDataError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + ApiKeyDataError::IoError(e) => Some(e), + } + } +} + +impl Display for ApiKeyDataError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ApiKeyDataError::IoError(e) => write!(f, "IO error: {}", e), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::NamedTempFile; + use std::io::Write; + + #[test] + fn load_trims_whitespace() { + let mut file = NamedTempFile::new().unwrap(); + write!(file, " sk-abc123 \n").unwrap(); + + let data = ApiKeyData::load_from_file(file.path()).unwrap(); + assert_eq!(data.as_str(), "sk-abc123"); + } + + #[test] + fn load_trims_newlines() { + let mut file = NamedTempFile::new().unwrap(); + write!(file, "\nmy_secret_key\n\n").unwrap(); + + let data = ApiKeyData::load_from_file(file.path()).unwrap(); + assert_eq!(data.as_str(), "my_secret_key"); + } + + #[test] + fn save_trims_whitespace() { + let file = NamedTempFile::new().unwrap(); + let data = ApiKeyData::from_str(" sk-padded "); + data.save_to_file(file.path()).unwrap(); + + let contents = std::fs::read_to_string(file.path()).unwrap(); + assert_eq!(contents, "sk-padded"); + } + + #[test] + fn round_trip() { + let file = NamedTempFile::new().unwrap(); + let original = ApiKeyData::from_str("sk-round-trip-test-key"); + original.save_to_file(file.path()).unwrap(); + + let loaded = ApiKeyData::load_from_file(file.path()).unwrap(); + assert_eq!(loaded.as_str(), "sk-round-trip-test-key"); + } + + #[test] + fn load_nonexistent_file_returns_io_error() { + let result = ApiKeyData::load_from_file("/tmp/nonexistent_api_key_file_12345.txt"); + assert!(result.is_err()); + } +} diff --git a/crates/desktop/artcraft/src/core/providers/credentials/payload/mod.rs b/crates/desktop/artcraft/src/core/providers/credentials/payload/mod.rs new file mode 100644 index 0000000000..8d5083bee3 --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/payload/mod.rs @@ -0,0 +1,3 @@ +pub mod provider_credential_payload; +pub mod api_key; +pub mod web_login; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/providers/credentials/payload/provider_credential_payload.rs b/crates/desktop/artcraft/src/core/providers/credentials/payload/provider_credential_payload.rs new file mode 100644 index 0000000000..c2b8ab562c --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/payload/provider_credential_payload.rs @@ -0,0 +1,10 @@ +use crate::core::providers::credentials::payload::api_key::ApiKeyData; +use crate::core::providers::credentials::payload::web_login::WebLoginData; + +#[derive(Clone)] +pub enum ProviderCredentialPayload { + ApiKey(ApiKeyData), + + // TODO: There might be logins in the future that use weird header states, etc. + WebLogin(WebLoginData), +} diff --git a/crates/desktop/artcraft/src/core/providers/credentials/payload/web_login.rs b/crates/desktop/artcraft/src/core/providers/credentials/payload/web_login.rs new file mode 100644 index 0000000000..4f9a41766e --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/payload/web_login.rs @@ -0,0 +1,251 @@ +use chrono::{DateTime, Utc}; +use serde_derive::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::error::Error; +use std::fmt::Display; +use std::path::Path; + +const CURRENT_VERSION: u8 = 1; + +#[derive(Clone)] +pub struct WebLoginData { + pub cookies_header: Option, + pub additional_headers: Option>, + + pub username: Option, + pub email_address: Option, + + pub created_at: Option>, + pub updated_at: Option>, +} + +impl WebLoginData { + pub fn new() -> Self { + Self { + cookies_header: None, + additional_headers: None, + username: None, + email_address: None, + created_at: None, + updated_at: None, + } + } + + pub fn load_from_file>(file_path: P) -> Result { + let contents = std::fs::read_to_string(file_path) + .map_err(WebLoginDataError::IoError)?; + + let serializable: WebLoginDataSerializable = serde_json::from_str(&contents) + .map_err(WebLoginDataError::DeserializeError)?; + + Ok(Self { + cookies_header: serializable.cookies_header, + additional_headers: serializable.additional_headers, + username: serializable.username, + email_address: serializable.email_address, + created_at: serializable.created_at, + updated_at: serializable.updated_at, + }) + } + + pub fn save_to_file>(&self, file_path: P) -> Result<(), WebLoginDataError> { + let serializable = WebLoginDataSerializable { + version: CURRENT_VERSION, + cookies_header: self.cookies_header.clone(), + additional_headers: self.additional_headers.clone(), + username: self.username.clone(), + email_address: self.email_address.clone(), + created_at: self.created_at, + updated_at: self.updated_at, + }; + + let contents = serde_json::to_string_pretty(&serializable) + .map_err(WebLoginDataError::SerializeError)?; + + std::fs::write(file_path, contents) + .map_err(WebLoginDataError::IoError)?; + + Ok(()) + } +} + +#[derive(Debug)] +pub enum WebLoginDataError { + IoError(std::io::Error), + DeserializeError(serde_json::Error), + SerializeError(serde_json::Error), +} + +impl Error for WebLoginDataError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + WebLoginDataError::IoError(e) => Some(e), + WebLoginDataError::DeserializeError(e) => Some(e), + WebLoginDataError::SerializeError(e) => Some(e), + _ => None, + } + } +} + +impl Display for WebLoginDataError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + WebLoginDataError::IoError(e) => write!(f, "IO error: {}", e), + WebLoginDataError::DeserializeError(e) => write!(f, "Deserialization error: {}", e), + WebLoginDataError::SerializeError(e) => write!(f, "Serialization error: {}", e), + } + } +} + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +struct WebLoginDataSerializable { + version: u8, + + #[serde(skip_serializing_if = "Option::is_none")] + cookies_header: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + additional_headers: Option>, + + #[serde(skip_serializing_if = "Option::is_none")] + username: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + email_address: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + created_at: Option>, + + #[serde(skip_serializing_if = "Option::is_none")] + updated_at: Option>, +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::NamedTempFile; + use std::io::Write; + + #[test] + fn save_writes_version() { + let file = NamedTempFile::new().unwrap(); + let data = WebLoginData::new(); + data.save_to_file(file.path()).unwrap(); + + let contents = std::fs::read_to_string(file.path()).unwrap(); + assert!(contents.contains("\"version\": 1")); + } + + #[test] + fn load_reads_all_fields() { + let mut file = NamedTempFile::new().unwrap(); + write!(file, r#"{{ + "version": 1, + "cookies_header": "session=abc123; visitor=xyz", + "additional_headers": {{"Authorization": "Bearer tok"}}, + "username": "testuser", + "email_address": "test@example.com", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-05-10T12:00:00Z" + }}"#).unwrap(); + + let data = WebLoginData::load_from_file(file.path()).unwrap(); + assert_eq!(data.cookies_header.as_deref(), Some("session=abc123; visitor=xyz")); + assert_eq!(data.username.as_deref(), Some("testuser")); + assert_eq!(data.email_address.as_deref(), Some("test@example.com")); + assert!(data.additional_headers.is_some()); + let headers = data.additional_headers.unwrap(); + assert_eq!(headers.get("Authorization").unwrap(), "Bearer tok"); + assert!(data.created_at.is_some()); + assert!(data.updated_at.is_some()); + } + + #[test] + fn load_tolerates_absent_optional_fields() { + let mut file = NamedTempFile::new().unwrap(); + write!(file, r#"{{"version": 1}}"#).unwrap(); + + let data = WebLoginData::load_from_file(file.path()).unwrap(); + assert!(data.cookies_header.is_none()); + assert!(data.additional_headers.is_none()); + assert!(data.username.is_none()); + assert!(data.email_address.is_none()); + assert!(data.created_at.is_none()); + assert!(data.updated_at.is_none()); + } + + #[test] + fn load_tolerates_null_fields() { + let mut file = NamedTempFile::new().unwrap(); + write!(file, r#"{{ + "version": 1, + "cookies_header": null, + "username": null, + "email_address": null + }}"#).unwrap(); + + let data = WebLoginData::load_from_file(file.path()).unwrap(); + assert!(data.cookies_header.is_none()); + assert!(data.username.is_none()); + assert!(data.email_address.is_none()); + } + + #[test] + fn load_invalid_json_returns_error() { + let mut file = NamedTempFile::new().unwrap(); + write!(file, "not valid json at all").unwrap(); + + let result = WebLoginData::load_from_file(file.path()); + assert!(result.is_err()); + } + + #[test] + fn load_nonexistent_file_returns_io_error() { + let result = WebLoginData::load_from_file("/tmp/nonexistent_web_login_12345.json"); + assert!(result.is_err()); + } + + #[test] + fn round_trip_full() { + let file = NamedTempFile::new().unwrap(); + + let mut headers = HashMap::new(); + headers.insert("X-Custom".to_string(), "value".to_string()); + + let now = Utc::now(); + + let original = WebLoginData { + cookies_header: Some("session=roundtrip".to_string()), + additional_headers: Some(headers), + username: Some("round_trip_user".to_string()), + email_address: Some("rt@example.com".to_string()), + created_at: Some(now), + updated_at: Some(now), + }; + + original.save_to_file(file.path()).unwrap(); + let loaded = WebLoginData::load_from_file(file.path()).unwrap(); + + assert_eq!(loaded.cookies_header.as_deref(), Some("session=roundtrip")); + assert_eq!(loaded.username.as_deref(), Some("round_trip_user")); + assert_eq!(loaded.email_address.as_deref(), Some("rt@example.com")); + assert_eq!(loaded.additional_headers.as_ref().unwrap().get("X-Custom").unwrap(), "value"); + assert!(loaded.created_at.is_some()); + assert!(loaded.updated_at.is_some()); + } + + #[test] + fn round_trip_minimal() { + let file = NamedTempFile::new().unwrap(); + + let original = WebLoginData::new(); + original.save_to_file(file.path()).unwrap(); + + let loaded = WebLoginData::load_from_file(file.path()).unwrap(); + assert!(loaded.cookies_header.is_none()); + assert!(loaded.additional_headers.is_none()); + assert!(loaded.username.is_none()); + assert!(loaded.email_address.is_none()); + } +} diff --git a/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_key.rs b/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_key.rs new file mode 100644 index 0000000000..4ba3f3ce3b --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_key.rs @@ -0,0 +1,50 @@ +use serde_derive::{Deserialize, Serialize}; +use crate::core::providers::credentials::provider_credential_type::ProviderCredentialType; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ProviderCredentialKey { + + // ========== API KEYS========== + + FalApiKey, + ReplicateApiKey, + + // ========== WEB LOGINS ========== + + GrokWebLogin, + HiggsfieldWebLogin, + MidjourneyLogin, + RunwayWebLogin, +} + + +impl ProviderCredentialKey { + pub fn get_type(&self) -> ProviderCredentialType { + match self { + // Api keys + Self::FalApiKey => ProviderCredentialType::ApiKey, + Self::ReplicateApiKey => ProviderCredentialType::ApiKey, + // Web logins + Self::GrokWebLogin => ProviderCredentialType::WebLogin, + Self::HiggsfieldWebLogin => ProviderCredentialType::WebLogin, + Self::MidjourneyLogin => ProviderCredentialType::WebLogin, + Self::RunwayWebLogin => ProviderCredentialType::WebLogin, + } + } + + pub fn get_filename(&self) -> &'static str { + match self { + // Api keys + ProviderCredentialKey::FalApiKey => "fal.api_key.txt", + ProviderCredentialKey::ReplicateApiKey => "replicate.api_key.txt", + // Web logins + ProviderCredentialKey::GrokWebLogin => "grok.web_login.toml", + ProviderCredentialKey::HiggsfieldWebLogin => "higgsfield.web_login.toml", + ProviderCredentialKey::MidjourneyLogin => "midjourney.web_login.toml", + ProviderCredentialKey::RunwayWebLogin => "runway.web_login.toml", + } + } +} + + diff --git a/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_loading_cache.rs b/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_loading_cache.rs new file mode 100644 index 0000000000..deed0e0636 --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_loading_cache.rs @@ -0,0 +1,187 @@ +use std::collections::HashMap; +use std::error::Error; +use std::fmt::Display; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; + +use log::{info, warn}; + +use crate::core::providers::credentials::payload::api_key::{ApiKeyData, ApiKeyDataError}; +use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload; +use crate::core::providers::credentials::payload::web_login::{WebLoginData, WebLoginDataError}; +use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey; +use crate::core::providers::credentials::provider_credential_type::ProviderCredentialType; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; + +struct CacheEntry { + payload: ProviderCredentialPayload, + loaded_at: Instant, +} + +#[derive(Clone)] +pub struct ProviderCredentialLoadingCache { + inner: Arc>>, + app_data_root: AppDataRoot, + ttl: Duration, +} + +impl ProviderCredentialLoadingCache { + pub fn new(app_data_root: AppDataRoot) -> Self { + Self { + inner: Arc::new(RwLock::new(HashMap::new())), + app_data_root, + ttl: Duration::from_secs(300), // 5 minutes + } + } + + pub fn new_with_ttl(app_data_root: AppDataRoot, ttl: Duration) -> Self { + Self { + inner: Arc::new(RwLock::new(HashMap::new())), + app_data_root, + ttl, + } + } + + pub fn get_credentials( + &self, + key: ProviderCredentialKey, + ) -> Result, ProviderCredentialLoadingCacheError> { + // Check cache first. + { + let cache = self.inner.read() + .map_err(|_| ProviderCredentialLoadingCacheError::LockPoisoned)?; + + if let Some(entry) = cache.get(&key) { + if entry.loaded_at.elapsed() < self.ttl { + return Ok(Some(entry.payload.clone())); + } + } + } + + // Cache miss or expired — load from disk. + let file_path = self.app_data_root.credentials_dir().path().join(key.get_filename()); + + if !file_path.exists() { + return Ok(None); + } + + info!("Loading credential from disk: {:?}", file_path); + + let payload = match key.get_type() { + ProviderCredentialType::ApiKey => { + let data = ApiKeyData::load_from_file(&file_path) + .map_err(ProviderCredentialLoadingCacheError::ApiKeyError)?; + ProviderCredentialPayload::ApiKey(data) + } + ProviderCredentialType::WebLogin => { + let data = WebLoginData::load_from_file(&file_path) + .map_err(ProviderCredentialLoadingCacheError::WebLoginError)?; + ProviderCredentialPayload::WebLogin(data) + } + }; + + // Store in cache. + { + let mut cache = self.inner.write() + .map_err(|_| ProviderCredentialLoadingCacheError::LockPoisoned)?; + + cache.insert(key, CacheEntry { + payload: payload.clone(), + loaded_at: Instant::now(), + }); + } + + Ok(Some(payload)) + } + + pub fn save_credentials( + &self, + key: ProviderCredentialKey, + payload: ProviderCredentialPayload, + ) -> Result<(), ProviderCredentialLoadingCacheError> { + let file_path = self.app_data_root.credentials_dir().path().join(key.get_filename()); + + info!("Saving credential to disk: {:?}", file_path); + + match (&payload, key.get_type()) { + (ProviderCredentialPayload::ApiKey(data), ProviderCredentialType::ApiKey) => { + data.save_to_file(&file_path) + .map_err(ProviderCredentialLoadingCacheError::ApiKeyError)?; + } + (ProviderCredentialPayload::WebLogin(data), ProviderCredentialType::WebLogin) => { + data.save_to_file(&file_path) + .map_err(ProviderCredentialLoadingCacheError::WebLoginError)?; + } + _ => { + warn!("Credential type mismatch for key {:?}", key); + return Err(ProviderCredentialLoadingCacheError::TypeMismatch); + } + } + + // Update cache. + { + let mut cache = self.inner.write() + .map_err(|_| ProviderCredentialLoadingCacheError::LockPoisoned)?; + + cache.insert(key, CacheEntry { + payload, + loaded_at: Instant::now(), + }); + } + + Ok(()) + } + + /// Invalidate a cached entry, forcing a re-read from disk on next access. + pub fn invalidate(&self, key: ProviderCredentialKey) { + if let Ok(mut cache) = self.inner.write() { + cache.remove(&key); + } + } + + pub fn delete_credentials(&self, key: ProviderCredentialKey) -> Result<(), ProviderCredentialLoadingCacheError> { + let file_path = self.app_data_root.credentials_dir().path().join(key.get_filename()); + + if file_path.exists() { + std::fs::remove_file(file_path) + .map_err(|e| ProviderCredentialLoadingCacheError::IoError(e))?; + } + + self.invalidate(key); + + Ok(()) + } +} + +#[derive(Debug)] +pub enum ProviderCredentialLoadingCacheError { + IoError(std::io::Error), + ApiKeyError(ApiKeyDataError), + WebLoginError(WebLoginDataError), + TypeMismatch, + LockPoisoned, +} + +impl Error for ProviderCredentialLoadingCacheError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + Self::IoError(err) => Some(err), + Self::ApiKeyError(err) => Some(err), + Self::WebLoginError(err) => Some(err), + _ => None, + } + } +} + +impl Display for ProviderCredentialLoadingCacheError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::IoError(e) => write!(f, "IO error: {}", e), + Self::ApiKeyError(e) => write!(f, "API key error: {}", e), + Self::WebLoginError(e) => write!(f, "Web login error: {}", e), + Self::TypeMismatch => write!(f, "Credential type does not match key type"), + Self::LockPoisoned => write!(f, "Internal lock poisoned"), + } + } +} diff --git a/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_type.rs b/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_type.rs new file mode 100644 index 0000000000..fd113fcec8 --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/credentials/provider_credential_type.rs @@ -0,0 +1,9 @@ +use serde_derive::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ProviderCredentialType { + ApiKey, + WebLogin, +} + diff --git a/crates/desktop/artcraft/src/core/providers/mod.rs b/crates/desktop/artcraft/src/core/providers/mod.rs new file mode 100644 index 0000000000..48c828e752 --- /dev/null +++ b/crates/desktop/artcraft/src/core/providers/mod.rs @@ -0,0 +1 @@ +pub mod credentials; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/state/task_database_version.rs b/crates/desktop/artcraft/src/core/state/task_database_version.rs index 2766c30434..d4a7d8c32b 100644 --- a/crates/desktop/artcraft/src/core/state/task_database_version.rs +++ b/crates/desktop/artcraft/src/core/state/task_database_version.rs @@ -7,4 +7,4 @@ /// We increment this whenever we make a change to the database schema. /// This prevents deadlock on startup if the schema has changed from the /// previous version. -pub const TASK_DATABASE_VERSION: u32 = 6; +pub const TASK_DATABASE_VERSION: u32 = 7; diff --git a/crates/desktop/artcraft/src/core/threads/mod.rs b/crates/desktop/artcraft/src/core/threads/mod.rs index 80525f98bb..85518c0d82 100644 --- a/crates/desktop/artcraft/src/core/threads/mod.rs +++ b/crates/desktop/artcraft/src/core/threads/mod.rs @@ -1,2 +1,3 @@ pub mod discord_presence_thread; -pub mod main_window_thread; \ No newline at end of file +pub mod main_window_thread; +pub mod third_party_task_polling_thread; \ No newline at end of file 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 new file mode 100644 index 0000000000..afa04b3cb4 --- /dev/null +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs @@ -0,0 +1,339 @@ +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::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::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; +use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; +use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use artcraft_client::error::api_error::ApiError; +use artcraft_client::error::storyteller_error::StorytellerError; +use artcraft_client::endpoints::media_files::get_media_file::get_media_file; +use artcraft_client::endpoints::media_files::legacy_upload_media_file_from_file::{ + legacy_upload_media_file_from_file, LegacyUploadMediaFileFromFileArgs, +}; +use artcraft_client::endpoints::media_files::upload_image_media_file_from_file::{ + upload_image_media_file_from_file, UploadImageFromFileArgs, +}; +use artcraft_client::endpoints::media_files::upload_video_media_file_from_file::{ + upload_video_media_file_from_file, UploadVideoFromFileArgs, +}; +use artcraft_client::endpoints::prompts::create_prompt::create_prompt; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass; +use enums::tauri::tasks::task_type::TaskType; +use fal_client::polling::poll_job_response::poll_job_response::PollJobResponse; +use fal_client::polling::poll_job_response::success_case_extractors::PollResponseExtractedContents; +use log::{error, info, warn}; +use sqlite_tasks::queries::task::Task; +use sqlite_tasks::queries::update_successful_task_status_with_metadata::{ + update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs, +}; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use std::time::Duration; +use tauri::AppHandle; +use tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::media_files::MediaFileToken; +use uuid_utils::uuid::generate_random_uuid; + +pub async fn handle_fal_complete( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + task: &Task, + job_response: PollJobResponse, +) { + info!("[FalComplete] Handling completed task {}", task.id.as_str()); + + let result = handle_fal_complete_inner( + app_handle, + app_env_configs, + app_data_root, + task_database, + storyteller_creds_manager, + task, + job_response, + ).await; + + if let Err(err) = result { + error!("[FalComplete] Failed to handle task {}: {:?}", task.id.as_str(), err); + } +} + +async fn handle_fal_complete_inner( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + task: &Task, + job_response: PollJobResponse, +) -> Result<(), Box> { + let creds = storyteller_creds_manager.get_credentials()? + .ok_or("No Storyteller credentials available")?; + + let extracted = job_response.extracted_contents; + + // Determine what kind of media we got and collect download URLs. + let (urls, media_class, generation_action) = collect_media_urls(task, &extracted)?; + + if urls.is_empty() { + warn!("[FalComplete] Task {} completed but no downloadable media found in response", task.id.as_str()); + //warn!("[FalComplete] Raw body: {}", &job_response.raw_body[..job_response.raw_body.len().min(500)]); + return Ok(()); + } + + // Create a prompt record + let prompt_response = create_prompt( + &app_env_configs.storyteller_host, + Some(&creds), + CreatePromptRequest { + uuid_idempotency_token: generate_random_uuid(), + positive_prompt: None, + negative_prompt: None, + model_type: None, + generation_provider: Some(GenerationProvider::Fal), + maybe_generation_mode: None, + maybe_aspect_ratio: None, + maybe_resolution: None, + maybe_batch_count: None, + maybe_generate_audio: None, + maybe_duration_seconds: None, + }, + ).await?; + + info!("[FalComplete] Created prompt: {:?}", prompt_response.prompt_token); + + let maybe_batch_token = if urls.len() > 1 { + let token = BatchGenerationToken::generate(); + info!("[FalComplete] Using batch token for {} files: {:?}", urls.len(), token); + Some(token) + } else { + None + }; + + let mut maybe_primary_media_file_token: Option = None; + + for (i, url) in urls.iter().enumerate() { + info!("[FalComplete] Downloading result {} from: {}", i, url); + + let download_path = download_file(url, app_data_root, i).await?; + + info!("[FalComplete] Uploading result {} to backend...", i); + + let media_token = upload_to_backend( + &creds, + app_env_configs, + &download_path, + &prompt_response.prompt_token, + maybe_batch_token.as_ref(), + media_class, + ).await?; + + info!("[FalComplete] Uploaded as media file: {:?}", media_token); + + if maybe_primary_media_file_token.is_none() { + maybe_primary_media_file_token = Some(media_token); + } + } + + // Look up CDN/thumbnail URLs for the primary media file + let mut maybe_cdn_url = None; + let mut maybe_thumbnail_url_template = None; + + if let Some(media_file_token) = maybe_primary_media_file_token.as_ref() { + match get_media_file(&app_env_configs.storyteller_host, media_file_token).await { + Ok(response) => { + maybe_cdn_url = Some(response.media_file.media_links.cdn_url.to_string()); + maybe_thumbnail_url_template = media_links_to_thumbnail_template(&response.media_file.media_links) + .map(|s| s.to_string()); + } + Err(err) => { + error!("[FalComplete] Failed to look up media file after upload: {:?} (failing open)", err); + } + } + } + + // Mark the task as completed + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { + db: task_database.get_connection(), + task_id: &task.id, + maybe_batch_token: maybe_batch_token.as_ref(), + maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), + maybe_primary_media_file_class: Some(media_class), + maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref(), + maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(), + }).await?; + + if updated { + let event = GenerationCompleteEvent { + action: Some(generation_action), + service: GenerationServiceProvider::Fal, + model: None, + }; + event.send_infallible(app_handle); + } + + info!("[FalComplete] Task {} fully handled", task.id.as_str()); + + Ok(()) +} + +// ── Helpers ── + +/// Collect downloadable media URLs from the extracted response contents. +fn collect_media_urls( + task: &Task, + extracted: &Option, +) -> Result<(Vec, TaskMediaFileClass, GenerationAction), Box> { + let extracted = match extracted { + Some(e) => e, + None => return Ok((vec![], TaskMediaFileClass::Image, GenerationAction::GenerateImage)), + }; + + // Images (batch) + if let Some(images) = &extracted.images { + let urls: Vec = images.iter() + .filter_map(|img| img.url.clone()) + .collect(); + if !urls.is_empty() { + return Ok((urls, TaskMediaFileClass::Image, GenerationAction::GenerateImage)); + } + } + + // Single image (e.g. background removal) + if let Some(image) = &extracted.image { + if let Some(url) = &image.url { + return Ok((vec![url.clone()], TaskMediaFileClass::Image, GenerationAction::GenerateImage)); + } + } + + // Video + if let Some(video) = &extracted.video { + if let Some(url) = &video.url { + return Ok((vec![url.clone()], TaskMediaFileClass::Video, GenerationAction::GenerateVideo)); + } + } + + // 3D model (GLB) + if let Some(glb) = &extracted.model_glb { + if let Some(url) = &glb.url { + return Ok((vec![url.clone()], TaskMediaFileClass::Dimensional, GenerationAction::ImageTo3d)); + } + } + + Ok((vec![], TaskMediaFileClass::Image, GenerationAction::GenerateImage)) +} + +async fn download_file( + url: &str, + app_data_root: &AppDataRoot, + index: usize, +) -> Result> { + let response = reqwest::get(url).await?; + let bytes = response.bytes().await?; + + let extension = url_utils::download_extension::extract_download_extension_from_url::extract_download_extension_from_url_str(url) + .map(|ext| ext.as_extension_without_period()) + .unwrap_or("bin"); + + let tempdir = app_data_root.temp_dir().path(); + let filename = format!("fal_{}_{}.{}", generate_random_uuid(), index, extension); + let download_path = tempdir.join(filename); + + let mut file = File::create(&download_path)?; + file.write_all(&bytes)?; + + Ok(download_path) +} + +const MAX_UPLOAD_RETRIES: u32 = 5; +const INITIAL_RETRY_DELAY_SECS: u64 = 10; + +async fn upload_to_backend( + creds: &StorytellerCredentialSet, + app_env_configs: &AppEnvConfigs, + download_path: &PathBuf, + prompt_token: &tokens::tokens::prompts::PromptToken, + maybe_batch_token: Option<&BatchGenerationToken>, + media_class: TaskMediaFileClass, +) -> Result> { + let mut retry_delay_secs = INITIAL_RETRY_DELAY_SECS; + + for attempt in 0..MAX_UPLOAD_RETRIES { + let result = try_upload(creds, app_env_configs, download_path, prompt_token, maybe_batch_token, media_class).await; + + match result { + Ok(token) => return Ok(token), + Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => { + if attempt + 1 < MAX_UPLOAD_RETRIES { + warn!( + "[FalComplete] Upload rate-limited (429), retrying in {}s (attempt {}/{})", + retry_delay_secs, + attempt + 1, + MAX_UPLOAD_RETRIES, + ); + tokio::time::sleep(Duration::from_secs(retry_delay_secs)).await; + retry_delay_secs = (retry_delay_secs * 2).min(60); + } else { + error!("[FalComplete] Upload rate-limited after {} attempts, giving up", MAX_UPLOAD_RETRIES); + return Err(Box::new(StorytellerError::Api(ApiError::TooManyRequests( + "Max retries exceeded".to_string(), + )))); + } + } + Err(err) => return Err(Box::new(err)), + } + } + + unreachable!() +} + +async fn try_upload( + creds: &StorytellerCredentialSet, + app_env_configs: &AppEnvConfigs, + download_path: &PathBuf, + prompt_token: &tokens::tokens::prompts::PromptToken, + maybe_batch_token: Option<&BatchGenerationToken>, + media_class: TaskMediaFileClass, +) -> Result { + let media_token = match media_class { + TaskMediaFileClass::Video => { + let result = upload_video_media_file_from_file(UploadVideoFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(creds), + path: download_path, + maybe_prompt_token: Some(prompt_token), + }).await?; + result.media_file_token + } + TaskMediaFileClass::Dimensional => { + let result = legacy_upload_media_file_from_file(LegacyUploadMediaFileFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(creds), + path: download_path, + }).await?; + result.media_file_token + } + _ => { + let result = upload_image_media_file_from_file(UploadImageFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(creds), + path: download_path, + is_intermediate_system_file: false, + maybe_prompt_token: Some(prompt_token), + maybe_batch_token, + }).await?; + result.media_file_token + } + }; + + Ok(media_token) +} diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_failure.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_failure.rs new file mode 100644 index 0000000000..2f754eb9eb --- /dev/null +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_failure.rs @@ -0,0 +1,51 @@ +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_failed_event::GenerationFailedEvent; +use crate::core::state::task_database::TaskDatabase; +use enums::tauri::tasks::task_status::TaskStatus; +use enums::tauri::tasks::task_type::TaskType; +use log::{error, info}; +use sqlite_tasks::queries::task::Task; +use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; +use tauri::AppHandle; + +pub async fn handle_fal_failure( + app_handle: &AppHandle, + task_database: &TaskDatabase, + task: &Task, + reason: &str, +) { + info!("[FalPolling] Marking task {} as failed: {}", task.id.as_str(), reason); + + let update_result = update_task_status(UpdateTaskArgs { + db: task_database.get_connection(), + task_id: &task.id, + status: TaskStatus::CompleteFailure, + }).await; + + if let Err(err) = update_result { + error!("[FalPolling] Failed to update task status for {}: {:?}", task.id.as_str(), err); + } + + let action = task_type_to_generation_action(task.task_type); + + let event = GenerationFailedEvent { + action, + service: GenerationServiceProvider::Fal, + model: None, + reason: Some(reason.to_string()), + }; + + event.send_infallible(app_handle); +} + +fn task_type_to_generation_action(task_type: TaskType) -> GenerationAction { + match task_type { + TaskType::ImageGeneration => GenerationAction::GenerateImage, + TaskType::VideoGeneration => GenerationAction::GenerateVideo, + TaskType::BackgroundRemoval => GenerationAction::RemoveBackground, + TaskType::ObjectGeneration => GenerationAction::ImageTo3d, + TaskType::GaussianGeneration => GenerationAction::GenerateGaussian, + TaskType::ImageInpaintEdit => GenerationAction::ImageInpaintEdit, + } +} diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/mod.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/mod.rs new file mode 100644 index 0000000000..95e4589169 --- /dev/null +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/mod.rs @@ -0,0 +1,3 @@ +pub mod handle_fal_complete; +pub mod handle_fal_failure; +pub mod poll_fal_tasks; diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/poll_fal_tasks.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/poll_fal_tasks.rs new file mode 100644 index 0000000000..81f9f585ac --- /dev/null +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/poll_fal_tasks.rs @@ -0,0 +1,183 @@ +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; +use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload; +use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey; +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; +use crate::core::threads::third_party_task_polling_thread::handlers::fal::handle_fal_complete::handle_fal_complete; +use crate::core::threads::third_party_task_polling_thread::handlers::fal::handle_fal_failure::handle_fal_failure; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use enums::tauri::tasks::task_status::TaskStatus; +use fal_client::creds::fal_api_key::FalApiKey; +use fal_client::polling::poll_job_response::poll_job_response::{poll_job_response, PollJobResponseArgs}; +use fal_client::polling::poll_job_status::poll_job_status::{poll_job_status, FalJobStatus, PollJobStatusArgs}; +use log::{error, info, warn}; +use sqlite_tasks::queries::task::Task; +use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; +use tauri::AppHandle; + +pub async fn poll_fal_tasks( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + credential_cache: &ProviderCredentialLoadingCache, + fal_tasks: &[&Task], +) { + let api_key = match load_fal_api_key(credential_cache) { + Some(key) => key, + None => { + warn!("[FalPolling] No FAL API key configured, skipping poll"); + return; + } + }; + + for task in fal_tasks { + let result = poll_single_fal_task( + app_handle, + app_env_configs, + app_data_root, + task_database, + storyteller_creds_manager, + &api_key, + task, + ).await; + + if let Err(err) = result { + error!( + "[FalPolling] Error processing task {}: {:?}", + task.id.as_str(), + err, + ); + } + } +} + +async fn poll_single_fal_task( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + api_key: &FalApiKey, + task: &Task, +) -> Result<(), Box> { + let status_url = match &task.queue_status_url { + Some(url) => url.as_str(), + None => { + warn!("[FalPolling] Task {} has no queue_status_url, marking as failed", task.id.as_str()); + handle_fal_failure( + app_handle, + task_database, + task, + "No queue status URL to check", + ).await; + return Ok(()); + } + }; + + // Step 1: Check job status + info!("[FalPolling] Checking status for task {}: {}", task.id.as_str(), status_url); + + let status_response = poll_job_status(PollJobStatusArgs { + status_url, + api_key, + }).await?; + + match status_response.status { + FalJobStatus::InQueue | FalJobStatus::InProgress => { + info!( + "[FalPolling] Task {} still {:?} (queue_position={:?})", + task.id.as_str(), + status_response.status, + status_response.maybe_queue_position, + ); + return Ok(()); + } + FalJobStatus::Unknown(ref status) => { + warn!("[FalPolling] Task {} has unknown status: {}", task.id.as_str(), status); + return Ok(()); + } + FalJobStatus::Completed => { + info!("[FalPolling] Task {} completed, fetching results...", task.id.as_str()); + } + } + + // Step 2: Fetch the response + let response_url = pick_response_url(task, &status_response.response_url); + + let response_url = match response_url { + Some(url) => url, + None => { + warn!("[FalPolling] Task {} completed but has no response URL, marking as failed", task.id.as_str()); + handle_fal_failure( + app_handle, + task_database, + task, + "Job completed but no response URL available", + ).await; + return Ok(()); + } + }; + + let job_response = match poll_job_response(PollJobResponseArgs { + response_url: &response_url, + api_key, + }).await { + Ok(response) => response, + Err(err) => { + error!( + "[FalPolling] Failed to fetch response for task {}: {}. Skipping for now.", + task.id.as_str(), + err, + ); + return Ok(()); + } + }; + + // Step 3: Handle completed job + handle_fal_complete( + app_handle, + app_env_configs, + app_data_root, + task_database, + storyteller_creds_manager, + task, + job_response, + ).await; + + Ok(()) +} + +// ── Helpers ── + +fn load_fal_api_key(credential_cache: &ProviderCredentialLoadingCache) -> Option { + match credential_cache.get_credentials(ProviderCredentialKey::FalApiKey) { + Ok(Some(ProviderCredentialPayload::ApiKey(data))) => { + Some(FalApiKey::from_str(data.as_str())) + } + _ => None, + } +} + +/// Use the task's stored response URL first; fall back to the one from the status poll. +/// Log if they differ. +fn pick_response_url(task: &Task, polled_response_url: &Option) -> Option { + match (&task.queue_response_url, polled_response_url) { + (Some(db_url), Some(polled_url)) => { + if db_url != polled_url { + warn!( + "[FalPolling] Task {} response URLs differ: db={}, polled={}. Using db URL.", + task.id.as_str(), + db_url, + polled_url, + ); + } + Some(db_url.clone()) + } + (Some(db_url), None) => Some(db_url.clone()), + (None, Some(polled_url)) => Some(polled_url.clone()), + (None, None) => None, + } +} diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/mod.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/mod.rs new file mode 100644 index 0000000000..e02f03de99 --- /dev/null +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/mod.rs @@ -0,0 +1 @@ +pub mod fal; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/mod.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/mod.rs new file mode 100644 index 0000000000..55faddd942 --- /dev/null +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/mod.rs @@ -0,0 +1,2 @@ +pub mod third_party_task_polling_thread; +pub mod handlers; diff --git a/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/third_party_task_polling_thread.rs b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/third_party_task_polling_thread.rs new file mode 100644 index 0000000000..552633c4cd --- /dev/null +++ b/crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/third_party_task_polling_thread.rs @@ -0,0 +1,138 @@ +use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache; +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; +use crate::core::threads::third_party_task_polling_thread::handlers::fal::poll_fal_tasks::poll_fal_tasks; +use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use enums::common::generation_provider::GenerationProvider; +use log::{error, info, warn}; +use sqlite_tasks::queries::list_non_artcraft_pending_tasks::{ + list_non_artcraft_pending_tasks, ListNonArtcraftPendingTasksArgs, +}; +use sqlite_tasks::queries::task::Task; +use std::time::Duration; +use tauri::AppHandle; + +const SLEEP_NO_THIRD_PARTY_JOBS_SEEN: Duration = Duration::from_secs(10); +const SLEEP_THIRD_PARTY_JOBS_SEEN: Duration = Duration::from_secs(2); +const SLEEP_BETWEEN_FAL_POLLS: Duration = Duration::from_secs(1); +const SLEEP_ON_ERROR: Duration = Duration::from_secs(30); + +pub async fn third_party_task_polling_thread( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, + task_database: TaskDatabase, + storyteller_creds_manager: StorytellerCredentialManager, + credential_cache: ProviderCredentialLoadingCache, +) -> ! { + let mut has_ever_seen_third_party_jobs = false; + + loop { + let result = poll_iteration( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &storyteller_creds_manager, + &credential_cache, + &mut has_ever_seen_third_party_jobs, + ).await; + + if let Err(err) = result { + error!("[ThirdPartyPolling] Error in polling loop: {:?}", err); + tokio::time::sleep(SLEEP_ON_ERROR).await; + } + } +} + +async fn poll_iteration( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + credential_cache: &ProviderCredentialLoadingCache, + has_ever_seen_third_party_jobs: &mut bool, +) -> Result<(), PollError> { + let task_list = list_non_artcraft_pending_tasks(ListNonArtcraftPendingTasksArgs { + db: task_database.get_connection(), + task_statuses: &TASK_DATABASE_PENDING_STATUSES, + }).await?; + + let tasks = task_list.tasks; + + if tasks.is_empty() { + let sleep_duration = if *has_ever_seen_third_party_jobs { + SLEEP_THIRD_PARTY_JOBS_SEEN + } else { + SLEEP_NO_THIRD_PARTY_JOBS_SEEN + }; + tokio::time::sleep(sleep_duration).await; + return Ok(()); + } + + *has_ever_seen_third_party_jobs = true; + + let fal_tasks: Vec<&Task> = tasks.iter() + .filter(|t| t.provider == GenerationProvider::Fal) + .collect(); + + let non_fal_tasks: Vec<&Task> = tasks.iter() + .filter(|t| t.provider != GenerationProvider::Fal) + .collect(); + + if !non_fal_tasks.is_empty() { + for task in &non_fal_tasks { + warn!( + "[ThirdPartyPolling] Skipping non-FAL task: id={}, provider={:?}, type={:?}", + task.id.as_str(), + task.provider, + task.task_type, + ); + } + } + + if fal_tasks.is_empty() { + tokio::time::sleep(SLEEP_THIRD_PARTY_JOBS_SEEN).await; + return Ok(()); + } + + info!("[ThirdPartyPolling] {} FAL job(s) ready to check", fal_tasks.len()); + + poll_fal_tasks( + app_handle, + app_env_configs, + app_data_root, + task_database, + storyteller_creds_manager, + credential_cache, + &fal_tasks, + ).await; + + tokio::time::sleep(SLEEP_BETWEEN_FAL_POLLS).await; + + Ok(()) +} + +// ── Error ── + +#[derive(Debug)] +enum PollError { + SqliteTasksError(sqlite_tasks::error::SqliteTasksError), +} + +impl From for PollError { + fn from(err: sqlite_tasks::error::SqliteTasksError) -> Self { + Self::SqliteTasksError(err) + } +} + +impl std::fmt::Display for PollError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::SqliteTasksError(err) => write!(f, "SQLite error: {:?}", err), + } + } +} diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index eb057e1ff4..c602406bf1 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -22,8 +22,11 @@ use crate::core::commands::get_app_info_command::get_app_info_command; use crate::core::commands::load_without_cors_command::load_without_cors_command; use crate::core::commands::media_files::media_file_delete_command::media_file_delete_command; use crate::core::commands::platform_info_command::platform_info_command; -use crate::core::commands::providers::get_provider_order_command::get_provider_order_command; -use crate::core::commands::providers::set_provider_order_command::set_provider_order_command; +use crate::core::commands::providers::deprecated::get_provider_order_command::get_provider_order_command; +use crate::core::commands::providers::deprecated::set_provider_order_command::set_provider_order_command; +use crate::core::commands::providers::provider_clear_command::provider_clear_command; +use crate::core::commands::providers::provider_list_command::provider_list_command; +use crate::core::commands::providers::provider_set_api_key_command::provider_set_api_key_command; use crate::core::commands::task_queue::get_task_queue_command::get_task_queue_command; use crate::core::commands::task_queue::mark_task_as_dismissed_command::mark_task_as_dismissed_command; use crate::core::commands::task_queue::tasks_nuke_all_command::tasks_nuke_all_command; @@ -93,6 +96,9 @@ pub fn run() { println!("Loading app preferences..."); let app_preferences = load_app_preferences_or_default(&app_data_root); + let provider_credential_cache = crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache::new(app_data_root.clone()); + let provider_credential_cache_2 = provider_credential_cache.clone(); + // NB: tauri-plugin-http stores the credentials on disk, so we can defer to that for now. // println!("Attempting to read existing artcraft credentials..."); // let storyteller_creds_manager = StorytellerCredentialManager::initialize_from_disk_infallible(&app_data_root); @@ -174,6 +180,7 @@ pub fn run() { grok_prompt_queue_2, worldlabs_bearer_bridge_2, worldlabs_creds_manager_2, + provider_credential_cache_2, ).await; if let Err(err) = result { @@ -196,6 +203,7 @@ pub fn run() { .manage(sora_task_queue) .manage(storyteller_creds_manager_3) .manage(worldlabs_bearer_bridge) + .manage(provider_credential_cache) .manage(worldlabs_creds_manager); // TODO: Break this out into another module, because RustRover/IntelliJ lags with these macros. @@ -218,6 +226,9 @@ pub fn run() { get_app_preferences_command, get_provider_order_command, get_task_queue_command, + provider_clear_command, + provider_list_command, + provider_set_api_key_command, grok_clear_credentials_command, grok_get_credential_info_command, grok_open_login_command, 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 557bc23997..64bf2042c8 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 @@ -94,7 +94,8 @@ pub async fn handle_classic_successful_generations( maybe_resolution: None, maybe_batch_count: None, maybe_generate_audio: None, - maybe_duration_seconds: None, }; + maybe_duration_seconds: None, + }; let prompt_response = create_prompt( &app_env_configs.storyteller_host, diff --git a/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_credits_command.rs b/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_credits_command.rs index 974ea5dc3b..be02b1cd53 100644 --- a/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_credits_command.rs +++ b/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_credits_command.rs @@ -15,7 +15,7 @@ use log::{error, info}; use serde_derive::{Deserialize, Serialize}; use artcraft_client::endpoints::credits::get_session_credits::get_session_credits; use tauri::{AppHandle, State}; -use crate::core::commands::providers::get_provider_order_command::GetProviderOrderResponse; +use crate::core::commands::providers::deprecated::get_provider_order_command::GetProviderOrderResponse; use crate::core::commands::response::success_response_wrapper::SerializeMarker; // TODO(bt,2025-09-12): Cache credits and clear cache on generate calls. diff --git a/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_subscription_command.rs b/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_subscription_command.rs index a622368982..98cdf3e429 100644 --- a/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_subscription_command.rs +++ b/crates/desktop/artcraft/src/services/storyteller/commands/storyteller_get_subscription_command.rs @@ -1,6 +1,6 @@ use chrono::{DateTime, Utc}; use crate::core::commands::deprecated::image_edit::enqueue_edit_image_command::{ImageEditModel, EditImageQuality, EditImageSize}; -use crate::core::commands::providers::get_provider_order_command::GetProviderOrderResponse; +use crate::core::commands::providers::deprecated::get_provider_order_command::GetProviderOrderResponse; use crate::core::commands::response::shorthand::ResponseOrErrorMessage; use crate::core::commands::response::success_response_wrapper::SerializeMarker; use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; diff --git a/crates/schema/database/sqlite_tasks/src/queries/create_task.rs b/crates/schema/database/sqlite_tasks/src/queries/create_task.rs index 1d7e4385db..86edbf5bf3 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/create_task.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/create_task.rs @@ -14,6 +14,8 @@ pub struct CreateTaskArgs<'a> { pub model_type: Option, pub provider: GenerationProvider, pub provider_job_id: Option<&'a str>, + pub queue_status_url: Option<&'a str>, + pub queue_response_url: Option<&'a str>, pub frontend_caller: Option, pub frontend_subscriber_id: Option<&'a str>, pub frontend_subscriber_payload: Option<&'a str>, @@ -40,11 +42,13 @@ pub async fn create_task( model_type, provider, provider_job_id, + queue_status_url, + queue_response_url, frontend_caller, frontend_subscriber_id, frontend_subscriber_payload ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "#, task_id_temp, status_temp, @@ -52,6 +56,8 @@ pub async fn create_task( model_type_temp, provider_temp, args.provider_job_id, + args.queue_status_url, + args.queue_response_url, frontend_caller_temp, args.frontend_subscriber_id, args.frontend_subscriber_payload diff --git a/crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs b/crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs index 3d1e3aaf80..ac259ac24e 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs @@ -33,6 +33,8 @@ pub async fn get_task_by_provider_and_provider_job_id( provider, provider_job_id, frontend_caller, + queue_status_url, + queue_response_url, frontend_subscriber_id, frontend_subscriber_payload FROM tasks @@ -69,6 +71,8 @@ pub async fn get_task_by_provider_and_provider_job_id( .transpose()?, provider: GenerationProvider::from_str(provider)?, provider_job_id: record.provider_job_id, + queue_status_url: record.queue_status_url, + queue_response_url: record.queue_response_url, frontend_caller: record.frontend_caller .map(|caller| TauriCommandCaller::from_str(&caller)) .transpose()?, diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_non_artcraft_pending_tasks.rs b/crates/schema/database/sqlite_tasks/src/queries/list_non_artcraft_pending_tasks.rs new file mode 100644 index 0000000000..0385003d01 --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/list_non_artcraft_pending_tasks.rs @@ -0,0 +1,82 @@ +use crate::connection::TaskDbConnection; +use crate::error::SqliteTasksError; +use crate::queries::task::{RawTask, Task}; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_model_type::TaskModelType; +use enums::tauri::tasks::task_status::TaskStatus; +use enums::tauri::tasks::task_type::TaskType; +use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; +use sqlx::{QueryBuilder, Sqlite}; +use std::collections::HashSet; +use tokens::tokens::sqlite::tasks::TaskId; + +pub struct ListNonArtcraftPendingTasksArgs<'a> { + pub db: &'a TaskDbConnection, + pub task_statuses: &'a HashSet, +} + +pub struct NonArtcraftTaskList { + pub tasks: Vec, +} + +/// List all tasks that are NOT from the Artcraft provider and match the given statuses. +pub async fn list_non_artcraft_pending_tasks( + args: ListNonArtcraftPendingTasksArgs<'_>, +) -> Result { + let artcraft_provider = GenerationProvider::Artcraft.to_string(); + + let mut query_builder: QueryBuilder = QueryBuilder::new(r#" + SELECT + id, + task_status, + task_type, + model_type, + provider, + provider_job_id, + queue_status_url, + queue_response_url, + frontend_caller, + frontend_subscriber_id, + frontend_subscriber_payload + FROM tasks + WHERE provider != + "#); + + query_builder.push_bind(artcraft_provider); + + if !args.task_statuses.is_empty() { + query_builder.push(" AND task_status IN ("); + let mut separated = query_builder.separated(", "); + for task_status in args.task_statuses.into_iter() { + separated.push_bind(task_status.to_str()); + } + separated.push_unseparated(") "); + } + + let query = query_builder.build_query_as::(); + let results = query.fetch_all(args.db.get_pool()).await?; + + let mut tasks: Vec = Vec::new(); + + for task in results { + tasks.push(Task { + id: TaskId::new_from_str(&task.id), + status: TaskStatus::from_str(&task.task_status)?, + task_type: TaskType::from_str(&task.task_type)?, + model_type: task.model_type + .map(|model| TaskModelType::from_str(&model)) + .transpose()?, + provider: GenerationProvider::from_str(&task.provider)?, + provider_job_id: task.provider_job_id, + queue_status_url: task.queue_status_url, + queue_response_url: task.queue_response_url, + frontend_caller: task.frontend_caller + .map(|caller| TauriCommandCaller::from_str(&caller)) + .transpose()?, + frontend_subscriber_id: task.frontend_subscriber_id, + frontend_subscriber_payload: task.frontend_subscriber_payload, + }); + } + + Ok(NonArtcraftTaskList { tasks }) +} diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs index 20e5f2003b..bc9af859ae 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs @@ -32,6 +32,8 @@ pub async fn list_tasks_by_provider_and_status( model_type, provider, provider_job_id, + queue_status_url, + queue_response_url, frontend_caller, frontend_subscriber_id, frontend_subscriber_payload @@ -73,6 +75,8 @@ pub async fn list_tasks_by_provider_and_status( .transpose()?, provider: GenerationProvider::from_str(&task.provider)?, provider_job_id: task.provider_job_id, + queue_status_url: task.queue_status_url, + queue_response_url: task.queue_response_url, frontend_caller: task.frontend_caller .map(|caller| TauriCommandCaller::from_str(&caller)) .transpose()?, diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs index 1aff07d89c..41bda79832 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs @@ -32,6 +32,8 @@ pub async fn list_tasks_by_provider_and_tokens( provider, provider_job_id, frontend_caller, + queue_status_url, + queue_response_url, frontend_subscriber_id, frontend_subscriber_payload FROM tasks @@ -72,6 +74,8 @@ pub async fn list_tasks_by_provider_and_tokens( .transpose()?, provider: GenerationProvider::from_str(&task.provider)?, provider_job_id: task.provider_job_id, + queue_status_url: task.queue_status_url, + queue_response_url: task.queue_response_url, frontend_caller: task.frontend_caller .map(|caller| TauriCommandCaller::from_str(&caller)) .transpose()?, diff --git a/crates/schema/database/sqlite_tasks/src/queries/mod.rs b/crates/schema/database/sqlite_tasks/src/queries/mod.rs index 8628ed1884..86b20be6e8 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/mod.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/mod.rs @@ -1,5 +1,6 @@ pub mod create_task; pub mod dismiss_all_tasks; +pub mod list_non_artcraft_pending_tasks; pub mod get_task_by_provider_and_provider_job_id; pub mod list_tasks_by_provider_and_status; pub mod list_tasks_by_provider_and_tokens; diff --git a/crates/schema/database/sqlite_tasks/src/queries/task.rs b/crates/schema/database/sqlite_tasks/src/queries/task.rs index 27ebe63db1..cb266caa1f 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/task.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/task.rs @@ -13,6 +13,8 @@ pub struct Task { pub model_type: Option, pub provider: GenerationProvider, pub provider_job_id: Option, + pub queue_status_url: Option, + pub queue_response_url: Option, pub frontend_caller: Option, pub frontend_subscriber_id: Option, pub frontend_subscriber_payload: Option, @@ -27,6 +29,8 @@ pub (crate) struct RawTask { pub (crate) model_type: Option, pub (crate) provider: String, pub (crate) provider_job_id: Option, + pub (crate) queue_status_url: Option, + pub (crate) queue_response_url: Option, pub (crate) frontend_caller: Option, pub (crate) frontend_subscriber_id: Option, pub (crate) frontend_subscriber_payload: Option, diff --git a/frontend/libs/api-enums/src/lib/common/generation_provider.ts b/frontend/libs/api-enums/src/lib/common/generation_provider.ts index cf41aabe62..bea5098e0c 100644 --- a/frontend/libs/api-enums/src/lib/common/generation_provider.ts +++ b/frontend/libs/api-enums/src/lib/common/generation_provider.ts @@ -8,6 +8,5 @@ export enum GenerationProvider { Midjourney = "midjourney", Sora = "sora", WorldLabs = "world_labs", - // NB: We should build our own open source Fal - //Fal = "fal", + Fal = "fal", } diff --git a/frontend/libs/components/model-selector/src/lib/provider-icons.tsx b/frontend/libs/components/model-selector/src/lib/provider-icons.tsx index b23df3a060..625dd32eb5 100644 --- a/frontend/libs/components/model-selector/src/lib/provider-icons.tsx +++ b/frontend/libs/components/model-selector/src/lib/provider-icons.tsx @@ -9,6 +9,7 @@ const GENERATION_PROVIDER_TO_CREATOR: Partial switch (provider) { case GenerationProvider.Artcraft: return "ArtCraft"; + case GenerationProvider.Fal: + return "FAL"; case GenerationProvider.Grok: return "Grok"; case GenerationProvider.Midjourney: diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx index 9fb6ddd084..b1c326636d 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx +++ b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx @@ -68,7 +68,7 @@ export const PromptBoxImage = ({ }: PromptBoxImageProps) => { useSignals(); - console.debug("Selected model:", selectedModel); + console.debug("Selected model and provider:", selectedModel, selectedProvider); useEffect(() => { if (imageMediaId && url) { @@ -335,6 +335,8 @@ export const PromptBoxImage = ({ timestamp: Date.now(), }); + console.debug("Image Generation Request", request); + const generateResponse = await GenerateImage(request); console.debug("PromptBoxImage - generateResponse", generateResponse); diff --git a/frontend/libs/components/settings-modal/src/lib/panes/AccountSettings/ProviderCredentialsGroupBlock.tsx b/frontend/libs/components/settings-modal/src/lib/panes/AccountSettings/ProviderCredentialsGroupBlock.tsx new file mode 100644 index 0000000000..fbb2a36587 --- /dev/null +++ b/frontend/libs/components/settings-modal/src/lib/panes/AccountSettings/ProviderCredentialsGroupBlock.tsx @@ -0,0 +1,184 @@ +import { invoke } from "@tauri-apps/api/core"; +import { useEffect, useState } from "react"; + +// ── Types matching the Rust ProviderCredentialKey enum ── + +type ProviderCredentialKey = + | "fal_api_key" + | "replicate_api_key" + | "grok_web_login" + | "higgsfield_web_login" + | "midjourney_login" + | "runway_web_login"; + +interface ProviderCredentialDetails { + maybe_key_start: string | null; + maybe_full_key: string | null; + maybe_email_address: string | null; + maybe_username: string | null; +} + +interface ProviderListEntry { + provider_credential: ProviderCredentialKey; + credential_type: string; + has_credentials: boolean; + maybe_details: ProviderCredentialDetails | null; +} + +interface ProviderListResponse { + providers: ProviderListEntry[]; +} + +// ── API key input row ── + +interface ApiKeyRowProps { + label: string; + credentialKey: ProviderCredentialKey; + initialRedactedValue: string; + initialFullKey: string; +} + +const ApiKeyRow = ({ + label, + credentialKey, + initialRedactedValue, + initialFullKey, +}: ApiKeyRowProps) => { + // The actual key value (editable). + const [fullKey, setFullKey] = useState(initialFullKey); + // The redacted display value. + const [redacted, setRedacted] = useState(initialRedactedValue); + // Whether the input is focused (show real value vs redacted). + const [isFocused, setIsFocused] = useState(false); + // Track if user changed something since last save. + const [isDirty, setIsDirty] = useState(false); + + const handleChange = (e: React.ChangeEvent) => { + const newValue = e.target.value; + setFullKey(newValue); + setIsDirty(true); + }; + + const handleFocus = () => { + setIsFocused(true); + }; + + const handleBlur = () => { + setIsFocused(false); + + if (!isDirty) return; + setIsDirty(false); + + const trimmed = fullKey.trim(); + + if (trimmed === "") { + invoke("provider_clear_command", { + request: { provider_credential: credentialKey }, + }).then(() => { + setFullKey(""); + setRedacted(""); + console.log(`[ProviderCredentials] Cleared ${credentialKey}`); + }).catch((e) => { + console.error(`[ProviderCredentials] Error clearing ${credentialKey}:`, e); + }); + } else { + invoke("provider_set_api_key_command", { + request: { provider_credential: credentialKey, api_key: trimmed }, + }).then(() => { + setFullKey(trimmed); + setRedacted( + trimmed.length > 6 + ? trimmed.substring(0, 6) + "********" + : "********" + ); + console.log(`[ProviderCredentials] Saved ${credentialKey}`); + }).catch((e) => { + console.error(`[ProviderCredentials] Error saving ${credentialKey}:`, e); + }); + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter") { + (e.target as HTMLInputElement).blur(); + } + }; + + const hasKey = isFocused ? fullKey.length > 0 : redacted.length > 0; + + return ( +
+ + +
+ ); +}; + +// ── Main block ── + +export const ProviderCredentialsGroupBlock = () => { + const [providers, setProviders] = useState([]); + const [isLoading, setIsLoading] = useState(true); + + useEffect(() => { + const fetchProviders = async () => { + try { + const result = (await invoke("provider_list_command")) as { + payload: ProviderListResponse; + }; + setProviders(result.payload.providers); + } catch (e) { + console.error("[ProviderCredentials] Error fetching provider list:", e); + } finally { + setIsLoading(false); + } + }; + fetchProviders(); + }, []); + + const findProvider = (key: ProviderCredentialKey) => + providers.find((p) => p.provider_credential === key); + + if (isLoading) { + return null; + } + + const falProvider = findProvider("fal_api_key"); + const replicateProvider = findProvider("replicate_api_key"); + + return ( +
+

API Keys

+ + +
+ ); +}; diff --git a/frontend/libs/components/settings-modal/src/lib/panes/ApiSettingsPane.tsx b/frontend/libs/components/settings-modal/src/lib/panes/ApiSettingsPane.tsx new file mode 100644 index 0000000000..977af54823 --- /dev/null +++ b/frontend/libs/components/settings-modal/src/lib/panes/ApiSettingsPane.tsx @@ -0,0 +1,9 @@ +import { ProviderCredentialsGroupBlock } from "./AccountSettings/ProviderCredentialsGroupBlock"; + +export const ApiSettingsPane = () => { + return ( +
+ +
+ ); +}; diff --git a/frontend/libs/components/settings-modal/src/lib/settings-modal.tsx b/frontend/libs/components/settings-modal/src/lib/settings-modal.tsx index 5e05b66f4e..694c83cdb8 100644 --- a/frontend/libs/components/settings-modal/src/lib/settings-modal.tsx +++ b/frontend/libs/components/settings-modal/src/lib/settings-modal.tsx @@ -4,6 +4,7 @@ import { FontAwesomeIcon } from "@fortawesome/react-fontawesome"; import { faUser, faCog, + faKey, faVolumeHigh, faCircleInfo, faCreditCard, @@ -19,6 +20,7 @@ import { ProviderPrioritySettingsPane } from "./panes/ProviderPrioritySettingsPa import { ExperimentalSettingsPane } from "./panes/ExperimentalSettingsPane"; import { gtagEvent } from "@storyteller/google-analytics"; import { BillingSettingsPane } from "./panes/BillingSettingsPane"; +import { ApiSettingsPane } from "./panes/ApiSettingsPane"; import { AppearanceSettingsPane } from "./panes/AppearanceSettingsPane"; import { Button } from "@storyteller/ui-button"; import { useExperimentalStore } from "./experimental-store"; @@ -36,6 +38,7 @@ type SettingsSection = | "general" | "appearance" | "accounts" + | "api_keys" | "alerts" | "about" | "provider_priority" @@ -74,6 +77,7 @@ export const SettingsModal = ({ { id: "general" as const, label: "General", icon: faCog }, { id: "accounts" as const, label: "Accounts", icon: faUser }, + { id: "api_keys" as const, label: "API Keys", icon: faKey }, { id: "billing" as const, label: "Plan & Credits", icon: faCreditCard }, /* { @@ -105,6 +109,8 @@ export const SettingsModal = ({ globalAccountLogoutCallback={globalAccountLogoutCallback} /> ); + case "api_keys": + return ; case "about": return ; case "provider_priority": diff --git a/frontend/libs/model-list/src/lib/lists/ImageModels.ts b/frontend/libs/model-list/src/lib/lists/ImageModels.ts index e63fae9826..7b935943ac 100644 --- a/frontend/libs/model-list/src/lib/lists/ImageModels.ts +++ b/frontend/libs/model-list/src/lib/lists/ImageModels.ts @@ -116,6 +116,10 @@ export const IMAGE_MODELS: ImageModel[] = [ fullName: "Flux 1 Dev", category: "image", creator: ModelCreator.BlackForestLabs, + providers: [ + GenerationProvider.Artcraft, + GenerationProvider.Fal, + ], selectorName: "Flux 1 Dev", selectorDescription: "Fast, but lower quality", selectorBadges: ["10 sec."], @@ -142,6 +146,10 @@ export const IMAGE_MODELS: ImageModel[] = [ fullName: "Flux 1 Schnell", category: "image", creator: ModelCreator.BlackForestLabs, + providers: [ + GenerationProvider.Artcraft, + GenerationProvider.Fal, + ], selectorName: "Flux 1 Schnell", selectorDescription: "Fastest image gen, but lowest quality", selectorBadges: ["10 sec."], @@ -203,6 +211,10 @@ export const IMAGE_MODELS: ImageModel[] = [ fullName: "Nano Banana 2", category: "image", creator: ModelCreator.Google, + providers: [ + GenerationProvider.Artcraft, + GenerationProvider.Fal, + ], selectorName: "Nano Banana 2", selectorDescription: "Fast instructive editing", selectorBadges: ["25 sec."], @@ -243,6 +255,10 @@ export const IMAGE_MODELS: ImageModel[] = [ fullName: "Nano Banana Pro", category: "image", creator: ModelCreator.Google, + providers: [ + GenerationProvider.Artcraft, + GenerationProvider.Fal, + ], selectorName: "Nano Banana Pro", selectorDescription: "Powerful instructive editing", selectorBadges: ["30 sec."], diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 606077f5af..6322a026d2 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -606,6 +606,10 @@ "name": "@storyteller/ui-pagedraw", "version": "0.0.1" }, + "libs/components/pagescene": { + "name": "@storyteller/ui-pagescene", + "version": "0.0.1" + }, "libs/components/pagination": { "name": "@storyteller/ui-pagination", "version": "0.0.1" @@ -10307,6 +10311,10 @@ "resolved": "libs/components/pagedraw", "link": true }, + "node_modules/@storyteller/ui-pagescene": { + "resolved": "libs/components/pagescene", + "link": true + }, "node_modules/@storyteller/ui-pagination": { "resolved": "libs/components/pagination", "link": true