From 012b7a9567f59e74af49abc00a5113e3fdea1b69 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 15 Jul 2025 22:28:03 -0400 Subject: [PATCH] artcraft task polling (#727) --- ...7ab4529eac2192475be7d08589886703ba1cb.json | 12 ++ .../src/jobs/list_session_jobs.rs | 43 +++++ .../storyteller_client/src/jobs/mod.rs | 1 + .../api_clients/storyteller_client/src/lib.rs | 1 + .../src/utils/basic_json_get_request.rs | 63 +++++++ .../storyteller_client/src/utils/mod.rs | 1 + .../startup/bootstrap_task_database.rs | 6 +- .../lifecycle/startup/handle_tauri_startup.rs | 22 ++- .../src/core/lifecycle/startup/mod.rs | 3 +- .../spawn_storytoller_task_polling_thread.rs | 23 +++ .../artcraft/src/core/state/task_database.rs | 1 + .../enum_conversion/generation_provider.rs | 11 ++ .../src/core/utils/enum_conversion/mod.rs | 2 + .../core/utils/enum_conversion/task_type.rs | 11 ++ crates/desktop/artcraft/src/core/utils/mod.rs | 1 + crates/desktop/artcraft/src/lib.rs | 6 +- .../artcraft/src/services/storyteller/mod.rs | 3 +- .../src/services/storyteller/threads/mod.rs | 1 + .../storyteller_task_polling_thread.rs | 156 ++++++++++++++++++ .../database/sqlite_tasks/src/connection.rs | 1 + .../schema/database/sqlite_tasks/src/error.rs | 14 +- .../list_tasks_by_provider_and_tokens.rs | 101 ++++++++++++ .../database/sqlite_tasks/src/queries/mod.rs | 4 +- .../src/queries/update_task_status.rs | 39 +++++ .../src/common/responses/job_details.rs | 5 +- .../src/common/responses/media_links.rs | 7 +- .../src/jobs/list_session_jobs.rs | 10 +- 27 files changed, 521 insertions(+), 27 deletions(-) create mode 100644 .sqlx/query-dc0367e9e43c3d7dda3623e77337ab4529eac2192475be7d08589886703ba1cb.json create mode 100644 crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs create mode 100644 crates/api_clients/storyteller_client/src/jobs/mod.rs create mode 100644 crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs create mode 100644 crates/desktop/artcraft/src/core/lifecycle/startup/spawn_storytoller_task_polling_thread.rs create mode 100644 crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs create mode 100644 crates/desktop/artcraft/src/core/utils/enum_conversion/mod.rs create mode 100644 crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs create mode 100644 crates/desktop/artcraft/src/services/storyteller/threads/mod.rs create mode 100644 crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs create mode 100644 crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs create mode 100644 crates/schema/database/sqlite_tasks/src/queries/update_task_status.rs diff --git a/.sqlx/query-dc0367e9e43c3d7dda3623e77337ab4529eac2192475be7d08589886703ba1cb.json b/.sqlx/query-dc0367e9e43c3d7dda3623e77337ab4529eac2192475be7d08589886703ba1cb.json new file mode 100644 index 0000000000..b4ed499639 --- /dev/null +++ b/.sqlx/query-dc0367e9e43c3d7dda3623e77337ab4529eac2192475be7d08589886703ba1cb.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "\n UPDATE tasks\n SET task_status = ?\n WHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + }, + "hash": "dc0367e9e43c3d7dda3623e77337ab4529eac2192475be7d08589886703ba1cb" +} diff --git a/crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs b/crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs new file mode 100644 index 0000000000..accf2a6043 --- /dev/null +++ b/crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs @@ -0,0 +1,43 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::basic_json_get_request::basic_json_get_request; +use artcraft_api_defs::jobs::list_session_jobs::{ListSessionJobsSuccessResponse, LIST_SESSION_JOBS_URL_PATH}; +use enums::common::job_status_plus::JobStatusPlus; +use std::collections::HashSet; + +#[derive(Clone)] +pub enum States { + All, + Include(HashSet), + Exclude(HashSet), +} + +pub async fn list_session_jobs( + api_host: &ApiHost, + maybe_creds: Option<&StorytellerCredentialSet>, + job_states: States, +) -> Result { + + let url = match job_states { + States::All => + LIST_SESSION_JOBS_URL_PATH.to_string(), + States::Include(states) => + format!("{}?include_states={}", LIST_SESSION_JOBS_URL_PATH, states_string(&states)), + States::Exclude(states) => + format!("{}?exclude_states={}", LIST_SESSION_JOBS_URL_PATH, states_string(&states)), + }; + + Ok(basic_json_get_request( + api_host, + &url, + maybe_creds, + ).await?) +} + +fn states_string(states: &HashSet) -> String { + states.iter() + .map(|state| state.to_str()) + .collect::>() + .join(",") +} \ No newline at end of file diff --git a/crates/api_clients/storyteller_client/src/jobs/mod.rs b/crates/api_clients/storyteller_client/src/jobs/mod.rs new file mode 100644 index 0000000000..164271aad7 --- /dev/null +++ b/crates/api_clients/storyteller_client/src/jobs/mod.rs @@ -0,0 +1 @@ +pub mod list_session_jobs; \ No newline at end of file diff --git a/crates/api_clients/storyteller_client/src/lib.rs b/crates/api_clients/storyteller_client/src/lib.rs index 228b2137de..e0af3d9aac 100644 --- a/crates/api_clients/storyteller_client/src/lib.rs +++ b/crates/api_clients/storyteller_client/src/lib.rs @@ -3,6 +3,7 @@ pub mod credentials; pub mod error; pub mod generate; +pub mod jobs; pub mod media_files; pub mod recipes; pub mod shared_response_types; diff --git a/crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs b/crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs new file mode 100644 index 0000000000..0adb217339 --- /dev/null +++ b/crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs @@ -0,0 +1,63 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::api_error::ApiError; +use crate::error::client_error::ClientError; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::constants::{APPLICATION_JSON, USER_AGENT}; +use crate::utils::filter_bad_response::filter_bad_response; +use artcraft_api_defs::generate::image::generate_flux_1_schnell_text_to_image::{GenerateFlux1SchnellTextToImageRequest, GenerateFlux1SchnellTextToImageResponse, GENERATE_FLUX_1_SCHNELL_TEXT_TO_IMAGE_PATH}; +use log::debug; +use reqwest::Client; +use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; + +pub async fn basic_json_get_request( + api_host: &ApiHost, + route_path: &str, + maybe_creds: Option<&StorytellerCredentialSet>, +) -> Result { + + let url = get_route(api_host, route_path); + + debug!("Requesting {:?}", &url); + + let client = Client::builder() + .gzip(true) + .build() + .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; + + let mut request_builder = client.get(url) + .header("User-Agent", USER_AGENT) + .header("Accept", APPLICATION_JSON); + + if let Some(creds) = maybe_creds { + if let Some(header) = &creds.maybe_as_cookie_header() { + request_builder = request_builder.header("Cookie", header); + } + } + + let response = request_builder + .send() + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let response = filter_bad_response(response) + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let response_body = &response.text() + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + debug!("Response body: {:?}", response_body); + + let response = serde_json::from_str(&response_body) + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + Ok(response) +} + +fn get_route(api_host: &ApiHost, route_path: &str) -> String { + let api_hostname = api_host.to_api_hostname_and_scheme(); + format!("{}{}", api_hostname, route_path) +} diff --git a/crates/api_clients/storyteller_client/src/utils/mod.rs b/crates/api_clients/storyteller_client/src/utils/mod.rs index defbbbc4c5..2211888f7f 100644 --- a/crates/api_clients/storyteller_client/src/utils/mod.rs +++ b/crates/api_clients/storyteller_client/src/utils/mod.rs @@ -1,4 +1,5 @@ pub mod api_host; +pub mod basic_json_get_request; pub mod basic_json_post_request; pub mod constants; pub mod filter_bad_response; diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/bootstrap_task_database.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/bootstrap_task_database.rs index e9575d10c9..343241fd17 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/bootstrap_task_database.rs +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/bootstrap_task_database.rs @@ -3,8 +3,8 @@ use crate::core::state::task_database::TaskDatabase; use errors::AnyhowResult; use tauri::{AppHandle, Manager}; -pub async fn bootstrap_task_database(app: &AppHandle, root: &AppDataRoot) -> AnyhowResult<()> { +pub async fn bootstrap_task_database(app: &AppHandle, root: &AppDataRoot) -> AnyhowResult { let task_database = TaskDatabase::connect(root).await?; - app.manage(task_database); - Ok(()) + app.manage(task_database.clone()); + Ok(task_database) } 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 1521b46ddc..4d75750f74 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 @@ -1,9 +1,25 @@ use crate::core::lifecycle::startup::bootstrap_task_database::bootstrap_task_database; +use crate::core::lifecycle::startup::spawn_storytoller_task_polling_thread::spawn_storyteller_task_polling_thread; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; use errors::AnyhowResult; -use tauri::AppHandle; +use tauri::{AppHandle, Manager}; + +pub async fn handle_tauri_startup( + app: AppHandle, + root: AppDataRoot, + app_env_configs: AppEnvConfigs, + storyteller_creds_manager: StorytellerCredentialManager, +) -> AnyhowResult<()> { + let task_database = bootstrap_task_database(&app, &root).await?; + + spawn_storyteller_task_polling_thread( + &app, + &app_env_configs, + &task_database, + &storyteller_creds_manager, + )?; -pub async fn handle_tauri_startup(app: AppHandle, root: AppDataRoot) -> AnyhowResult<()> { - bootstrap_task_database(&app, &root).await?; Ok(()) } diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/mod.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/mod.rs index 368fd30356..f3a123ffca 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/mod.rs +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/mod.rs @@ -1,2 +1,3 @@ pub mod bootstrap_task_database; -pub mod handle_tauri_startup; \ No newline at end of file +pub mod handle_tauri_startup; +mod spawn_storytoller_task_polling_thread; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/spawn_storytoller_task_polling_thread.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/spawn_storytoller_task_polling_thread.rs new file mode 100644 index 0000000000..609164eafd --- /dev/null +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/spawn_storytoller_task_polling_thread.rs @@ -0,0 +1,23 @@ +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; +use crate::services::storyteller::threads::storyteller_task_polling_thread::storyteller_task_polling_thread; +use errors::AnyhowResult; +use tauri::{AppHandle, Manager}; + +pub fn spawn_storyteller_task_polling_thread( + app: &AppHandle, + app_env_configs: &AppEnvConfigs, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + + tauri::async_runtime::spawn(storyteller_task_polling_thread( + app.clone(), + app_env_configs.clone(), + task_database.clone(), + storyteller_creds_manager.clone(), + )); + + Ok(()) +} diff --git a/crates/desktop/artcraft/src/core/state/task_database.rs b/crates/desktop/artcraft/src/core/state/task_database.rs index 9df07c4cd2..4467785f0d 100644 --- a/crates/desktop/artcraft/src/core/state/task_database.rs +++ b/crates/desktop/artcraft/src/core/state/task_database.rs @@ -2,6 +2,7 @@ use crate::core::state::data_dir::app_data_root::AppDataRoot; use errors::AnyhowResult; use sqlite_tasks::connection::TaskDbConnection; +#[derive(Clone)] pub struct TaskDatabase { connection: TaskDbConnection, } diff --git a/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs b/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs new file mode 100644 index 0000000000..d667d7cc53 --- /dev/null +++ b/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs @@ -0,0 +1,11 @@ +use crate::core::events::generation_events::common::GenerationServiceProvider; +use enums::common::generation_provider::GenerationProvider; + +// TODO(bt,2025-07-15): Get rid of GenerationServiceProvider +pub fn to_generation_service_provider(provider: GenerationProvider) -> GenerationServiceProvider { + match provider { + GenerationProvider::Artcraft => GenerationServiceProvider::Artcraft, + GenerationProvider::Fal => GenerationServiceProvider::Fal, + GenerationProvider::Sora => GenerationServiceProvider::Sora, + } +} diff --git a/crates/desktop/artcraft/src/core/utils/enum_conversion/mod.rs b/crates/desktop/artcraft/src/core/utils/enum_conversion/mod.rs new file mode 100644 index 0000000000..67b4f3707f --- /dev/null +++ b/crates/desktop/artcraft/src/core/utils/enum_conversion/mod.rs @@ -0,0 +1,2 @@ +pub mod generation_provider; +pub mod task_type; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs b/crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs new file mode 100644 index 0000000000..b9ac6dd479 --- /dev/null +++ b/crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs @@ -0,0 +1,11 @@ +use crate::core::events::generation_events::common::GenerationAction; +use enums::tauri::tasks::task_type::TaskType; + +pub fn 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, + } +} diff --git a/crates/desktop/artcraft/src/core/utils/mod.rs b/crates/desktop/artcraft/src/core/utils/mod.rs index d9003fc855..3b42fdd79b 100644 --- a/crates/desktop/artcraft/src/core/utils/mod.rs +++ b/crates/desktop/artcraft/src/core/utils/mod.rs @@ -2,6 +2,7 @@ pub mod best_window_size_heuristic; pub mod clear_all_webview_cookies; pub mod download_media_file_to_temp_dir; pub mod download_url_to_temp_dir; +pub mod enum_conversion; pub mod get_url_file_extension; pub mod image; pub mod large_file_transfer; diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index ea05077137..43b83bdc49 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -129,11 +129,13 @@ pub fn run() { // )?; //} let app = app.handle().clone(); - let app2 = app.clone(); + let handle = app.clone(); let root = app_data_root_2.clone(); + let env_config = app_env_configs_2.clone(); + let storyteller_creds = storyteller_creds_manager_2.clone(); tauri::async_runtime::block_on(async move { - if let Err(err) = handle_tauri_startup(app2, root).await { + if let Err(err) = handle_tauri_startup(handle, root, env_config, storyteller_creds).await { error!("Failed to handle Tauri startup: {:?}", err); panic!("Failed to handle Tauri startup: {:?}", err); } diff --git a/crates/desktop/artcraft/src/services/storyteller/mod.rs b/crates/desktop/artcraft/src/services/storyteller/mod.rs index 0d0b840411..51f42a9be7 100644 --- a/crates/desktop/artcraft/src/services/storyteller/mod.rs +++ b/crates/desktop/artcraft/src/services/storyteller/mod.rs @@ -1 +1,2 @@ -pub mod state; \ No newline at end of file +pub mod state; +pub mod threads; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/mod.rs b/crates/desktop/artcraft/src/services/storyteller/threads/mod.rs new file mode 100644 index 0000000000..09250df48d --- /dev/null +++ b/crates/desktop/artcraft/src/services/storyteller/threads/mod.rs @@ -0,0 +1 @@ +pub mod storyteller_task_polling_thread; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs new file mode 100644 index 0000000000..6411c0fcf2 --- /dev/null +++ b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs @@ -0,0 +1,156 @@ +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::events::generation_events::generation_failed_event::GenerationFailedEvent; +use crate::core::events::sendable_event_trait::SendableEvent; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; +use crate::core::state::task_database::TaskDatabase; +use crate::core::utils::enum_conversion::generation_provider::to_generation_service_provider; +use crate::core::utils::enum_conversion::task_type::to_generation_action; +use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; +use crate::services::sora::state::sora_task_queue::SoraTaskQueue; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use anyhow::anyhow; +use enums::common::generation_provider::GenerationProvider; +use enums::common::job_status_plus::JobStatusPlus; +use enums::tauri::tasks::task_status::TaskStatus; +use errors::AnyhowResult; +use log::{error, info}; +use reqwest::Url; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; +use std::collections::HashMap; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; +use storyteller_client::jobs::list_session_jobs::{list_session_jobs, States}; +use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file; +use tauri::AppHandle; +use tempdir::TempDir; + +pub async fn storyteller_task_polling_thread( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + task_database: TaskDatabase, + storyteller_creds_manager: StorytellerCredentialManager, +) -> ! { + loop { + let res = polling_loop( + &app_handle, + &app_env_configs, + &task_database, + &storyteller_creds_manager, + ).await; + if let Err(err) = res { + error!("An error occurred: {:?}", err); + } + // NB: Only sleep if an error occurs. + tokio::time::sleep(std::time::Duration::from_millis(30_000)).await; + } +} + +async fn polling_loop( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + loop { + // Wait before next request for jobs. + tokio::time::sleep(std::time::Duration::from_millis(5_000)).await; + + let creds = storyteller_creds_manager.get_credentials()?; + + let result = list_session_jobs( + &app_env_configs.storyteller_host, + creds.as_ref(), + States::All, + ).await; + + let jobs = match result { + Ok(result) => result.jobs, + Err(err) => { + match &err { + StorytellerError::Api(ApiError::TooManyRequests(message)) => { + error!("Too many requests (sleeping): {:?}", message); + tokio::time::sleep(std::time::Duration::from_millis(60_000)).await; + } + _ => {} + } + return Err(anyhow!(err)); + } + }; + + let job_ids = jobs.iter() + .map(|job| job.job_token.to_string()) + .collect::>(); + + let tasks = list_tasks_by_provider_and_tokens(ListTasksArgs { + db: task_database.get_connection(), + provider: GenerationProvider::Artcraft, + provider_job_ids: Some(job_ids), + }).await?; + + let tasks = tasks.tasks; + + let jobs_by_id = jobs.iter() + .map(|job| (job.job_token.to_string(), job)) + .collect::>(); + + let tasks_by_provider_job_id = tasks.iter() + .filter_map(|task| { + if let Some(provider_job_id) = &task.provider_job_id { + Some((provider_job_id.clone(), task.clone())) + } else { + None + } + }) + //.map(|task| (task.provider_job_id.clone(), task.clone())) + .collect::>(); + + for job in jobs.iter() { + match job.status.status { + JobStatusPlus::CompleteSuccess => {} // Fall-through. + _ => continue, + } + + let task = match tasks_by_provider_job_id.get(job.job_token.as_str()) { + Some(job) => job, + None => continue, + }; + + match task.status { + TaskStatus::CompleteSuccess => continue, + _ => {} // Fall-through. + } + + let updated = update_task_status(UpdateTaskArgs { + db: task_database.get_connection(), + task_id: &task.id, + status: TaskStatus::CompleteSuccess, + }).await?; + + if !updated { + continue; // If anything breaks with queries, don't spam events. + } + + let service = to_generation_service_provider(task.provider); + let action = to_generation_action(task.task_type); + + let event = GenerationCompleteEvent { + action: Some(action), + service, + model: None, // TODO + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open + } + } + + } +} diff --git a/crates/schema/database/sqlite_tasks/src/connection.rs b/crates/schema/database/sqlite_tasks/src/connection.rs index 50972d4670..2a46d832d0 100644 --- a/crates/schema/database/sqlite_tasks/src/connection.rs +++ b/crates/schema/database/sqlite_tasks/src/connection.rs @@ -3,6 +3,7 @@ use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode}; use sqlx::{Pool, Sqlite, SqlitePool}; use std::path::Path; +#[derive(Clone)] pub struct TaskDbConnection { pool: Pool, } diff --git a/crates/schema/database/sqlite_tasks/src/error.rs b/crates/schema/database/sqlite_tasks/src/error.rs index a04b9ec225..708c4fd196 100644 --- a/crates/schema/database/sqlite_tasks/src/error.rs +++ b/crates/schema/database/sqlite_tasks/src/error.rs @@ -3,12 +3,13 @@ use std::fmt::{Display, Formatter}; #[derive(Debug)] pub enum SqliteTasksError { - SqlxError(sqlx::Error), - //TaskNotFound, - //TaskAlreadyExists, - //InvalidTaskStatus, - //InvalidTaskType, - //InvalidGenerationProvider, + SqlxError(sqlx::Error), + TaskParseError(String), + //TaskNotFound, + //TaskAlreadyExists, + //InvalidTaskStatus, + //InvalidTaskType, + //InvalidGenerationProvider, } impl Error for SqliteTasksError {} @@ -17,6 +18,7 @@ impl Display for SqliteTasksError { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { SqliteTasksError::SqlxError(err) => write!(f, "SQLx error: {}", err), + SqliteTasksError::TaskParseError(err) => write!(f, "Task parse error: {}", err), } } } 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 new file mode 100644 index 0000000000..bb60ee90ad --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs @@ -0,0 +1,101 @@ +use crate::connection::TaskDbConnection; +use crate::error::SqliteTasksError; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_status::TaskStatus; +use enums::tauri::tasks::task_type::TaskType; +use sqlx::{QueryBuilder, Sqlite}; +use tokens::tokens::sqlite::tasks::TaskId; + +pub struct ListTasksArgs<'a> { + pub db: &'a TaskDbConnection, + pub provider: GenerationProvider, + pub provider_job_ids: Option>, +} + +pub struct TaskList { + pub tasks: Vec, +} + +#[derive(Debug, Clone)] +pub struct Task { + pub id: TaskId, + pub status: TaskStatus, + pub task_type: TaskType, + pub provider: GenerationProvider, + pub provider_job_id: Option, + pub frontend_subscriber_id: Option, + pub frontend_subscriber_payload: Option, +} + +#[derive(Debug)] +#[derive(sqlx::FromRow)] +struct RawTask { + id: String, + task_status: String, + task_type: String, + provider: String, + provider_job_id: Option, + frontend_subscriber_id: Option, + frontend_subscriber_payload: Option, +} + +pub async fn list_tasks_by_provider_and_tokens( + args: ListTasksArgs<'_>, +) -> Result { + + let mut query_builder: QueryBuilder = QueryBuilder::new(r#" + SELECT + id, + task_status, + task_type, + provider, + provider_job_id, + frontend_subscriber_id, + frontend_subscriber_payload + FROM tasks + WHERE provider = + "#); + + // TODO(bt,2025-07-15): Fix this. The sqlx mysql queries never required temporaries + let provider = args.provider.to_string(); + + query_builder.push_bind(provider); + + if let Some(provider_job_ids) = args.provider_job_ids { + query_builder.push(" AND provider_job_id IN ("); + let mut separated = query_builder.separated(", "); + + for job_id in provider_job_ids.into_iter() { + separated.push_bind(job_id); + } + + separated.push_unseparated(") "); + } + + let query = query_builder.build_query_as::(); + + // info!("Query: {:?}", query.sql()); + + 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) + .map_err(|err| SqliteTasksError::TaskParseError(err))?, + task_type: TaskType::from_str(&task.task_type) + .map_err(|err| SqliteTasksError::TaskParseError(err))?, + provider: GenerationProvider::from_str(&task.provider) + .map_err(|err| SqliteTasksError::TaskParseError(err))?, + provider_job_id: task.provider_job_id, + frontend_subscriber_id: task.frontend_subscriber_id, + frontend_subscriber_payload: task.frontend_subscriber_payload, + }); + } + + Ok(TaskList { + tasks, + }) +} diff --git a/crates/schema/database/sqlite_tasks/src/queries/mod.rs b/crates/schema/database/sqlite_tasks/src/queries/mod.rs index cd1a4daa46..1c9eb91946 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/mod.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/mod.rs @@ -1 +1,3 @@ -pub mod create_task; \ No newline at end of file +pub mod create_task; +pub mod list_tasks_by_provider_and_tokens; +pub mod update_task_status; \ No newline at end of file diff --git a/crates/schema/database/sqlite_tasks/src/queries/update_task_status.rs b/crates/schema/database/sqlite_tasks/src/queries/update_task_status.rs new file mode 100644 index 0000000000..ead78ece46 --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/update_task_status.rs @@ -0,0 +1,39 @@ +use crate::connection::TaskDbConnection; +use crate::error::SqliteTasksError; +use enums::tauri::tasks::task_status::TaskStatus; +use tokens::tokens::sqlite::tasks::TaskId; + +pub struct UpdateTaskArgs<'a> { + pub db: &'a TaskDbConnection, + pub task_id: &'a TaskId, + pub status: TaskStatus, +} + +/// Returns true if rows were updated. +pub async fn update_task_status( + args: UpdateTaskArgs<'_>, +) -> Result { + + // TODO(bt,2025-07-12): Fix this. The sqlx mysql queries never required temporaries + let task_id_temp = args.task_id.as_str(); + let status_temp = args.status.to_str(); + + // info!("Update task id: {}, status: {}", task_id_temp, status_temp); + + // TODO(bt,2025-07-15): We can't set a LIMIT without a certain compiler flag for SQLite ? + let query = sqlx::query!(r#" + UPDATE tasks + SET task_status = ? + WHERE id = ? + "#, + status_temp, + task_id_temp, + ); + + // info!("query: {:?}", query.sql()); + + let res = query.execute(args.db.get_pool()).await?; + let rows_updated = res.rows_affected() > 0; + + Ok(rows_updated) +} diff --git a/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs b/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs index d6d572043b..1ff8f3229d 100644 --- a/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs +++ b/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs @@ -1,9 +1,10 @@ +use serde::Deserialize; use serde::Serialize; use tokens::tokens::media_files::MediaFileToken; use utoipa::ToSchema; /// Details about submitted live portrait jobs (request arguments only) -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct JobDetailsLivePortraitRequest { /// Source media token /// This can be an image or a video media file token @@ -20,7 +21,7 @@ pub struct JobDetailsLivePortraitRequest { } /// Details about submitted lipsync jobs (request arguments only) -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct JobDetailsLipsyncRequest { /// Media file token for the source audio. /// This is probably an audio file, but in the future we might pull audio from video. diff --git a/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs b/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs index a8868a65e1..367dae629d 100644 --- a/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs +++ b/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs @@ -1,9 +1,10 @@ +use serde::Deserialize; use serde::Serialize; -use utoipa::ToSchema; use url::Url; +use utoipa::ToSchema; /// Links to media file locations (bucket, CDN, etc.) -#[derive(Serialize, ToSchema, Debug, Clone, Eq, PartialEq)] +#[derive(Serialize, Deserialize, ToSchema, Debug, Clone, Eq, PartialEq)] pub struct MediaLinks { /// Primary link to the asset via the CDN. pub cdn_url: Url, @@ -20,7 +21,7 @@ pub struct MediaLinks { pub maybe_video_previews: Option, } -#[derive(Serialize, ToSchema, Debug, Clone, Eq, PartialEq)] +#[derive(Serialize, Deserialize, ToSchema, Debug, Clone, Eq, PartialEq)] pub struct VideoPreviews { /// A static single frame preview image of the video. pub still: Url, diff --git a/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs b/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs index f4b85fcc38..281fc20a2d 100644 --- a/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs +++ b/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs @@ -27,7 +27,7 @@ pub struct ListSessionJobsQueryParams { // TODO(bt,2024-04-23): Add the ability for users to dismiss completed/dead jobs from view. } -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionJobsSuccessResponse { pub success: bool, @@ -35,7 +35,7 @@ pub struct ListSessionJobsSuccessResponse { pub jobs: Vec, } -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionJobsItem { pub job_token: InferenceJobToken, @@ -48,7 +48,7 @@ pub struct ListSessionJobsItem { } /// Details about what the user requested for generation -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionRequestDetailsResponse { pub inference_category: InferenceCategory, pub maybe_model_type: Option, @@ -73,7 +73,7 @@ pub struct ListSessionRequestDetailsResponse { } /// Details about the ongoing job status -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionStatusDetailsResponse { /// Primary status from the database (a state machine). pub status: JobStatusPlus, @@ -107,7 +107,7 @@ pub struct ListSessionStatusDetailsResponse { } /// Details about the completed result (if any) -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionResultDetailsResponse { pub entity_type: String, pub entity_token: String,