artcraft task polling (wip)

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<JobStatusPlus>),
+  Exclude(HashSet<JobStatusPlus>),
+}
+
+pub async fn list_session_jobs(
+  api_host: &ApiHost,
+  maybe_creds: Option<&StorytellerCredentialSet>,
+  job_states: States,
+) -> Result<ListSessionJobsSuccessResponse, StorytellerError> {
+  
+  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<JobStatusPlus>) -> String {
+  states.iter()
+    .map(|state| state.to_str())
+    .collect::<Vec<_>>()
+    .join(",")
+}

crates/api_clients/storyteller_client/src/jobs/mod.rs

@@ -0,0 +1 @@
+pub mod list_session_jobs;

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;

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, info};
+use reqwest::Client;
+use serde::de::DeserializeOwned;
+use serde::{Deserialize, Serialize};
+
+pub async fn basic_json_get_request<Res: DeserializeOwned>(
+  api_host: &ApiHost,
+  route_path: &str,
+  maybe_creds: Option<&StorytellerCredentialSet>,
+) -> Result<Res, StorytellerError> {
+
+  let url = get_route(api_host, route_path);
+
+  info!("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)
+}

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;

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,
 }

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

@@ -1 +1,2 @@
-pub mod state;
+pub mod state;
+pub mod threads;

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

@@ -0,0 +1 @@
+pub mod storyteller_task_polling_thread;

crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs

@@ -0,0 +1,144 @@
+use std::collections::HashMap;
+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::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 errors::AnyhowResult;
+use log::{error, info};
+use reqwest::Url;
+use std::fs::File;
+use std::io::Write;
+use std::path::PathBuf;
+use anyhow::anyhow;
+use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
+use tauri::AppHandle;
+use tempdir::TempDir;
+use enums::common::generation_provider::GenerationProvider;
+use enums::common::job_status_plus::JobStatusPlus;
+use enums::tauri::tasks::task_status::TaskStatus;
+use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
+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 crate::core::state::task_database::TaskDatabase;
+
+pub async fn storyteller_task_polling_thread(
+  app_handle: AppHandle,
+  app_env_configs: AppEnvConfigs,
+  app_data_root: AppDataRoot,
+  task_database: TaskDatabase,
+  storyteller_creds_manager: StorytellerCredentialManager,
+) -> ! {
+  loop {
+    let res = polling_loop(
+      &app_handle,
+      &app_env_configs,
+      &app_data_root,
+      &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,
+  app_data_root: &AppDataRoot,
+  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::<Vec<_>>();
+
+    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::<HashMap<String, _>>();
+
+    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::<HashMap<String, Task>>();
+
+    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 event = GenerationCompleteEvent {
+        //media_file_token: result.media_file_token,
+        action: Some(GenerationAction::GenerateImage), // TODO
+        service: GenerationServiceProvider::Artcraft, // TODO
+        model: None, // TODO
+      };
+
+      if let Err(err) = event.send(&app_handle) {
+        error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
+      }
+    }
+
+  }
+}

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<Sqlite>,
 }

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),
     }
   }
 }

crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs

@@ -0,0 +1,104 @@
+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};
+
+pub struct ListTasksArgs<'a> {
+  pub db: &'a TaskDbConnection,
+  pub provider: GenerationProvider,
+  pub provider_job_ids: Option<Vec<String>>,
+}
+
+pub struct TaskList {
+  pub tasks: Vec<Task>,
+}
+
+#[derive(Debug, Clone)]
+pub struct Task {
+  pub status: TaskStatus,
+  pub task_type: TaskType,
+  pub provider: GenerationProvider,
+  pub provider_job_id: Option<String>,
+  pub frontend_subscriber_id: Option<String>,
+  pub frontend_subscriber_payload: Option<String>,
+}
+
+#[derive(Debug)]
+#[derive(sqlx::FromRow)]
+struct RawTask {
+  status: String,
+  task_type: String,
+  provider: String,
+  provider_job_id: Option<String>,
+  frontend_subscriber_id: Option<String>,
+  frontend_subscriber_payload: Option<String>,
+}
+
+pub async fn list_tasks_by_provider_and_tokens(
+  args: ListTasksArgs<'_>,
+) -> Result<TaskList, SqliteTasksError> {
+
+  let mut query_builder: QueryBuilder<Sqlite> = 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 (");
+    //for (i, job_id) in provider_job_ids.iter().enumerate() {
+    //  if i > 0 {
+    //    query_builder.push(", ");
+    //  }
+    //  query_builder.push_bind(job_id);
+    //}
+    //query_builder.push(")");
+
+    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();
+
+  let query = query_builder.build_query_as::<RawTask>();
+  let results = query.fetch_all(args.db.get_pool()).await?;
+
+  let mut tasks: Vec<Task> = Vec::new();
+
+  for task in results {
+    tasks.push(Task {
+      status: TaskStatus::from_str(&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,
+  })
+}

crates/schema/database/sqlite_tasks/src/queries/mod.rs

@@ -1 +1,2 @@
-pub mod create_task;
+pub mod create_task;
+pub mod list_tasks_by_provider_and_tokens;

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.

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<VideoPreviews>,
 }
 
-#[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,

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<ListSessionJobsItem>,
 }
 
-#[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<String>,
@@ -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,