diff --git a/.sqlx/query-49c878f5ea50d249260ce4877a6c2072ede9145b83935561563b3dcbd8f79b5d.json b/.sqlx/query-49c878f5ea50d249260ce4877a6c2072ede9145b83935561563b3dcbd8f79b5d.json new file mode 100644 index 0000000000..825aa8c125 --- /dev/null +++ b/.sqlx/query-49c878f5ea50d249260ce4877a6c2072ede9145b83935561563b3dcbd8f79b5d.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "\n UPDATE tasks\n SET \n task_status = ?,\n on_complete_primary_media_file_token = ?,\n on_complete_batch_token = ?,\n on_complete_primary_media_file_thumbnail_url_template = ?\n WHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 5 + }, + "nullable": [] + }, + "hash": "49c878f5ea50d249260ce4877a6c2072ede9145b83935561563b3dcbd8f79b5d" +} diff --git a/.sqlx/query-a142dc3fdf5b6e741c268d8f2e9996a28c214f5835681297d58a9431da8de939.json b/.sqlx/query-5e6a6c099fbf4dbc949d3d46d885f4072e1d0761edc109ff4e1fc92b01caf3c0.json similarity index 61% rename from .sqlx/query-a142dc3fdf5b6e741c268d8f2e9996a28c214f5835681297d58a9431da8de939.json rename to .sqlx/query-5e6a6c099fbf4dbc949d3d46d885f4072e1d0761edc109ff4e1fc92b01caf3c0.json index 4b6a33696a..788735d5c0 100644 --- a/.sqlx/query-a142dc3fdf5b6e741c268d8f2e9996a28c214f5835681297d58a9431da8de939.json +++ b/.sqlx/query-5e6a6c099fbf4dbc949d3d46d885f4072e1d0761edc109ff4e1fc92b01caf3c0.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 created_at as \"created_at: DateTime\",\n updated_at as \"updated_at: DateTime\",\n completed_at as \"completed_at: DateTime\"\n FROM tasks\n WHERE is_dismissed_by_user == 0\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 frontend_subscriber_id,\n frontend_subscriber_payload,\n on_complete_primary_media_file_token,\n on_complete_primary_media_file_type,\n on_complete_batch_token,\n on_complete_primary_media_file_thumbnail_url_template,\n created_at as \"created_at: DateTime\",\n updated_at as \"updated_at: DateTime\",\n completed_at as \"completed_at: DateTime\"\n FROM tasks\n WHERE is_dismissed_by_user == 0\n ", "describe": { "columns": [ { @@ -49,18 +49,38 @@ "type_info": "Text" }, { - "name": "created_at: DateTime", + "name": "on_complete_primary_media_file_token", "ordinal": 9, + "type_info": "Text" + }, + { + "name": "on_complete_primary_media_file_type", + "ordinal": 10, + "type_info": "Text" + }, + { + "name": "on_complete_batch_token", + "ordinal": 11, + "type_info": "Text" + }, + { + "name": "on_complete_primary_media_file_thumbnail_url_template", + "ordinal": 12, + "type_info": "Text" + }, + { + "name": "created_at: DateTime", + "ordinal": 13, "type_info": "Int64" }, { "name": "updated_at: DateTime", - "ordinal": 10, + "ordinal": 14, "type_info": "Int64" }, { "name": "completed_at: DateTime", - "ordinal": 11, + "ordinal": 15, "type_info": "Int64" } ], @@ -77,10 +97,14 @@ true, true, true, + true, + true, + true, + true, false, false, true ] }, - "hash": "a142dc3fdf5b6e741c268d8f2e9996a28c214f5835681297d58a9431da8de939" + "hash": "5e6a6c099fbf4dbc949d3d46d885f4072e1d0761edc109ff4e1fc92b01caf3c0" } diff --git a/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql b/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql index f63423b175..21c952f57f 100644 --- a/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql +++ b/_sql/artcraft_migrations/20250708000000_create_tasks_table.sql @@ -10,6 +10,12 @@ -- tasks_v2.sqlite - added model_type (nullable) -- tasks_v3.sqlite - added frontend_caller, comments -- tasks_v4.sqlite - added is_dismissed_by_user +-- tasks_v5.sqlite - added four fields: +-- on_complete_primary_media_file_token, +-- on_complete_primary_media_file_type, +-- on_complete_batch_token, +-- on_complete_primary_media_file_thumbnail_url_template, +-- CREATE TABLE tasks ( -- Task auto-incrementing primary key. @@ -54,6 +60,25 @@ CREATE TABLE tasks ( -- Whether the user has dismissed the task from view. is_dismissed_by_user INTEGER NOT NULL DEFAULT 0, + -- OPTIONAL. + -- When the generation completes, the media token of the "primary" media file. + -- We'll probably always generate homogenous types, but just in case this is the first or primary item in a batch. + on_complete_primary_media_file_token TEXT, + + -- OPTIONAL. + -- When the generation completes, the file type of the primary media file. + -- We'll probably always generate homogenous types, but just in case this is the first or primary item in a batch. + on_complete_primary_media_file_type TEXT, + + -- OPTIONAL. + -- When the generation completes, the batch token (if any). + -- This will point to a collection of multiple files. + on_complete_batch_token TEXT, + + -- OPTIONAL. + -- The URL template for the thumbnail of the primary media file. + on_complete_primary_media_file_thumbnail_url_template TEXT, + created_at INTEGER NOT NULL DEFAULT (unixepoch('now')), updated_at INTEGER NOT NULL DEFAULT (unixepoch('now')), completed_at INTEGER DEFAULT NULL diff --git a/crates/desktop/artcraft/src/core/commands/task_queue/get_task_queue_command.rs b/crates/desktop/artcraft/src/core/commands/task_queue/get_task_queue_command.rs index b5c9f480d7..035b732a40 100644 --- a/crates/desktop/artcraft/src/core/commands/task_queue/get_task_queue_command.rs +++ b/crates/desktop/artcraft/src/core/commands/task_queue/get_task_queue_command.rs @@ -34,11 +34,19 @@ pub struct TaskQueueItem { pub model_type: Option, pub provider: Option, pub provider_job_id: Option, + pub completed_item: Option, pub created_at: DateTime, pub updated_at: DateTime, pub completed_at: Option>, } +#[derive(Serialize)] +pub struct CompletedItemData { + pub primary_media_file_token: MediaFileToken, + pub thumbnail_template: Option, + //pub cdn_url: Option, +} + impl SerializeMarker for GetTaskQueueCommandResponse {} #[tauri::command] @@ -74,17 +82,35 @@ pub async fn handle_request( let tasks = list_tasks_for_frontend(task_database.get_connection()) .await?; - Ok(tasks.tasks.into_iter() - .map(|task| TaskQueueItem { - id: task.id, - task_status: task.status, - task_type: task.task_type, - model_type: task.model_type, - provider: task.provider, - provider_job_id: task.provider_job_id, - created_at: task.created_at, - updated_at: task.updated_at, - completed_at: task.completed_at, - }) - .collect()) + let mut transformed_tasks = Vec::with_capacity(tasks.tasks.len()); + + for task in tasks.tasks.into_iter() { + let mut completed_item = None; + + if task.status == TaskStatus::CompleteSuccess { + if let Some(primary_media_file_token) = task.on_complete_primary_media_file_token { + completed_item = Some(CompletedItemData { + primary_media_file_token, + thumbnail_template: task.on_complete_primary_media_file_thumbnail_url_template, + }); + } else { + warn!("Task {} is marked complete but has no primary media file token", task.id); + } + } + + transformed_tasks.push(TaskQueueItem { + id: task.id, + task_status: task.status, + task_type: task.task_type, + model_type: task.model_type, + provider: task.provider, + provider_job_id: task.provider_job_id, + created_at: task.created_at, + updated_at: task.updated_at, + completed_at: task.completed_at, + completed_item, + }) + } + + Ok(transformed_tasks) } 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 eef293ca86..a12fd0196d 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 = 4; +pub const TASK_DATABASE_VERSION: u32 = 5; 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 index d6e74011c8..ffd81b9f54 100644 --- 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 @@ -25,18 +25,20 @@ 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_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs}; 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::credentials::storyteller_credential_set::StorytellerCredentialSet; -use storyteller_client::error::api_error::ApiError; -use storyteller_client::error::storyteller_error::StorytellerError; use storyteller_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States}; use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; use tauri::AppHandle; use tempdir::TempDir; +use tokens::tokens::media_files::MediaFileToken; pub async fn storyteller_task_polling_thread( app_handle: AppHandle, @@ -134,10 +136,22 @@ async fn polling_loop( _ => {} // Fall-through. } - let updated = update_task_status(UpdateTaskArgs { + let maybe_primary_media_file_token = job.maybe_result + .as_ref() + .map(|result| MediaFileToken::new_from_str(&result.entity_token)); + + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { db: task_database.get_connection(), task_id: &task.id, - status: TaskStatus::CompleteSuccess, + maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), + maybe_batch_token: job.maybe_result + .as_ref() + .map(|result| result.maybe_batch_token.as_ref()) + .flatten(), + maybe_primary_media_file_thumbnail_url_template: job.maybe_result + .as_ref() + .map(|result| result.media_links.maybe_thumbnail_template.as_deref()) + .flatten(), }).await?; if !updated { diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_for_frontend.rs b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_for_frontend.rs index cfb2576d1e..1c1528ba5c 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_for_frontend.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_for_frontend.rs @@ -6,6 +6,8 @@ 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 tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::sqlite::tasks::TaskId; pub struct TaskList { @@ -22,6 +24,10 @@ pub struct TaskItem { pub frontend_caller: Option, pub frontend_subscriber_id: Option, pub frontend_subscriber_payload: Option, + pub on_complete_primary_media_file_token: Option, + pub on_complete_primary_media_file_type: Option, + pub on_complete_batch_token: Option, + pub on_complete_primary_media_file_thumbnail_url_template: Option, pub created_at: DateTime, pub updated_at: DateTime, pub completed_at: Option>, @@ -43,6 +49,10 @@ pub async fn list_tasks_for_frontend( frontend_caller, frontend_subscriber_id, frontend_subscriber_payload, + on_complete_primary_media_file_token, + on_complete_primary_media_file_type, + on_complete_batch_token, + on_complete_primary_media_file_thumbnail_url_template, created_at as "created_at: DateTime", updated_at as "updated_at: DateTime", completed_at as "completed_at: DateTime" @@ -72,6 +82,10 @@ pub async fn list_tasks_for_frontend( .transpose()?, frontend_subscriber_id: raw.frontend_subscriber_id, frontend_subscriber_payload: raw.frontend_subscriber_payload, + on_complete_primary_media_file_token: raw.on_complete_primary_media_file_token.map(|t| MediaFileToken::new_from_str(&t)), + on_complete_primary_media_file_type: raw.on_complete_primary_media_file_type, + on_complete_batch_token: raw.on_complete_batch_token.map(|t| BatchGenerationToken::new_from_str(&t)), + on_complete_primary_media_file_thumbnail_url_template: raw.on_complete_primary_media_file_thumbnail_url_template, created_at: raw.created_at, updated_at: raw.updated_at, completed_at: raw.completed_at, @@ -93,6 +107,10 @@ struct TaskItemRaw { frontend_caller: Option, frontend_subscriber_id: Option, frontend_subscriber_payload: Option, + on_complete_primary_media_file_token: Option, + on_complete_primary_media_file_type: Option, + on_complete_batch_token: Option, + on_complete_primary_media_file_thumbnail_url_template: Option, created_at: DateTime, updated_at: DateTime, completed_at: Option>, diff --git a/crates/schema/database/sqlite_tasks/src/queries/mod.rs b/crates/schema/database/sqlite_tasks/src/queries/mod.rs index 083cdda18b..a318edea1c 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/mod.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/mod.rs @@ -3,4 +3,5 @@ pub mod list_tasks_by_provider_and_status; pub mod list_tasks_by_provider_and_tokens; pub mod list_tasks_for_frontend; pub mod mark_task_as_dismissed; +pub mod update_successful_task_status_with_metadata; pub mod update_task_status; \ No newline at end of file diff --git a/crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata.rs b/crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata.rs new file mode 100644 index 0000000000..da510a7bd2 --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata.rs @@ -0,0 +1,51 @@ +use crate::connection::TaskDbConnection; +use crate::error::SqliteTasksError; +use enums::tauri::tasks::task_status::TaskStatus; +use tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::sqlite::tasks::TaskId; + +const SUCCESSFUL_STATUS : &str = TaskStatus::CompleteSuccess.to_str(); + +pub struct UpdateSuccessfulTaskArgs<'a> { + pub db: &'a TaskDbConnection, + pub task_id: &'a TaskId, + pub maybe_primary_media_file_token: Option<&'a MediaFileToken>, + pub maybe_batch_token: Option<&'a BatchGenerationToken>, + pub maybe_primary_media_file_thumbnail_url_template: Option<&'a str>, +} + +/// Returns true if rows were updated. +pub async fn update_successful_task_status_with_metadata( + args: UpdateSuccessfulTaskArgs<'_>, +) -> 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 maybe_primary_media_token_temp = args.maybe_primary_media_file_token.map(|t| t.as_str()); + let maybe_batch_token_temp = args.maybe_batch_token.map(|t| t.as_str()); + + // 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 = ?, + on_complete_primary_media_file_token = ?, + on_complete_batch_token = ?, + on_complete_primary_media_file_thumbnail_url_template = ? + WHERE id = ? + "#, + SUCCESSFUL_STATUS, + maybe_primary_media_token_temp, + maybe_batch_token_temp, + args.maybe_primary_media_file_thumbnail_url_template, + 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/enums/src/tauri/tasks/task_status.rs b/crates/schema/public/enums/src/tauri/tasks/task_status.rs index 8e4dbec150..cc130c8c77 100644 --- a/crates/schema/public/enums/src/tauri/tasks/task_status.rs +++ b/crates/schema/public/enums/src/tauri/tasks/task_status.rs @@ -35,7 +35,7 @@ impl Default for TaskStatus { } impl TaskStatus { - pub fn to_str(&self) -> &'static str { + pub const fn to_str(&self) -> &'static str { match self { Self::Pending => "pending", Self::Started => "started", diff --git a/frontend/apps/editor3d/app/src/components/signaled/TopBar/TaskQueue.tsx b/frontend/apps/editor3d/app/src/components/signaled/TopBar/TaskQueue.tsx index 36504ec922..42caa90d26 100644 --- a/frontend/apps/editor3d/app/src/components/signaled/TopBar/TaskQueue.tsx +++ b/frontend/apps/editor3d/app/src/components/signaled/TopBar/TaskQueue.tsx @@ -219,44 +219,58 @@ export const TaskQueue = () => { const load = async () => { try { const result = await GetTaskQueue(); + if (cancelled) return; console.log("TaskQueue:GetTaskQueue result", result); + const { tasks } = result; // Fetch media files for thumbnails, using dummy tokens from API for now const tokens = result.results as string[] | undefined; + let imageUrls: string[] = []; - let mediaTokens: string[] | undefined = undefined; - if (tokens && tokens.length > 0) { - try { - const api = new MediaFilesApi(); - const batchToken = tokens.find((t) => t.startsWith("batch_")); - if (batchToken) { - const resp = await api.GetMediaFilesByBatchToken({ - batchToken, - }); - if (resp.success && resp.data) { - imageUrls = resp.data - .filter((m) => !!m?.media_links?.cdn_url) - .map((m) => m.media_links.cdn_url); - mediaTokens = resp.data.map((m) => m.token); - } - } else { - const resp = await api.ListMediaFilesByTokens({ - mediaTokens: tokens, - }); - if (resp.success && resp.data) { - imageUrls = resp.data - .filter((m) => !!m?.media_links?.cdn_url) - .map((m) => m.media_links.cdn_url); - mediaTokens = resp.data.map((m) => m.token); - } + //let mediaTokens: string[] | undefined = undefined; + let mediaTokens: string[] = []; + + for (const task of tasks) { + if (!!task.completed_item) { + mediaTokens.push(task.completed_item.primary_media_file_token); + if (!!task.completed_item.thumbnail_template) { + imageUrls.push(task.completed_item.thumbnail_template.replace("{WIDTH}", "250")); } - } catch (_) { - // ignore } } + //if (tokens && tokens.length > 0) { + // //try { + // // const api = new MediaFilesApi(); + // // const batchToken = tokens.find((t) => t.startsWith("batch_")); + // // if (batchToken) { + // // const resp = await api.GetMediaFilesByBatchToken({ + // // batchToken, + // // }); + // // if (resp.success && resp.data) { + // // imageUrls = resp.data + // // .filter((m) => !!m?.media_links?.cdn_url) + // // .map((m) => m.media_links.cdn_url); + // // mediaTokens = resp.data.map((m) => m.token); + // // } + // // } else { + // // const resp = await api.ListMediaFilesByTokens({ + // // mediaTokens: tokens, + // // }); + // // if (resp.success && resp.data) { + // // imageUrls = resp.data + // // .filter((m) => !!m?.media_links?.cdn_url) + // // .map((m) => m.media_links.cdn_url); + // // mediaTokens = resp.data.map((m) => m.token); + // // } + // // } + // //} catch (_) { + // // // ignore + // //} + //} + const now = Date.now(); const inProg = tasks .filter( diff --git a/frontend/libs/tauri-api/src/lib/artcraft/task_queue/GetTaskQueue.ts b/frontend/libs/tauri-api/src/lib/artcraft/task_queue/GetTaskQueue.ts index 93aa514a5c..e211f36951 100644 --- a/frontend/libs/tauri-api/src/lib/artcraft/task_queue/GetTaskQueue.ts +++ b/frontend/libs/tauri-api/src/lib/artcraft/task_queue/GetTaskQueue.ts @@ -19,11 +19,17 @@ export interface TaskQueueItem { model_type?: TaskModelType; provider?: GenerationProvider; provider_job_id?: string; + completed_item?: TaskQueueCompletedItem; created_at: Date; updated_at: Date; completed_at?: Date; } +export interface TaskQueueCompletedItem { + primary_media_file_token: string, + thumbnail_template?: string, +} + interface GetTaskQueueSuccess extends CommandResult { payload: GetTaskQueueResponse; } @@ -44,6 +50,7 @@ export const GetTaskQueue = async (): Promise => { model_type: task.model_type, provider: task.provider, provider_job_id: task.provider_job_id, + completed_item: task.completed_item, created_at: new Date(task.created_at), updated_at: new Date(task.updated_at), completed_at: task.completed_at ? new Date(task.completed_at) : undefined,