polling thread

crates/api_clients/midjourney_client/src/endpoints/mod.rs

@@ -1,5 +1,5 @@
 pub mod get_index_page_html;
+pub mod imagine;
 pub mod storage_list;
 pub mod submit_job;
-pub mod user_connections;
-mod imagine;
+pub mod user_connections;

crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs

@@ -1,7 +1,7 @@
 use crate::core::lifecycle::startup::tasks::bootstrap_task_database::bootstrap_task_database;
+use crate::core::lifecycle::startup::tasks::initially_size_and_position_windows::initially_size_and_position_windows;
 use crate::core::lifecycle::startup::tasks::load_provider_priority_state::load_provider_priority_state;
 use crate::core::lifecycle::startup::tasks::set_app_log_level::set_app_log_level;
-use crate::core::lifecycle::startup::tasks::initially_size_and_position_windows::initially_size_and_position_windows;
 use crate::core::lifecycle::startup::tasks::spawn_discord_presence_thread::spawn_discord_presence_thread;
 use crate::core::lifecycle::startup::tasks::spawn_fal_task_polling_thread::spawn_fal_task_polling_thread;
 use crate::core::lifecycle::startup::tasks::spawn_main_window_thread::spawn_main_window_thread;
@@ -11,6 +11,8 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
 use crate::services::fal::state::fal_task_queue::FalTaskQueue;
+use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
+use crate::services::midjourney::threads::midjourney_long_polling_thread::midjourney_long_polling_thread;
 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;
@@ -26,6 +28,7 @@ pub async fn handle_tauri_startup(
   sora_task_queue: SoraTaskQueue,
   fal_credential_manager: FalCredentialManager,
   fal_task_queue: FalTaskQueue,
+  mj_creds_manager: MidjourneyCredentialManager,
 ) -> AnyhowResult<()> {
 
   set_app_log_level(
@@ -72,6 +75,13 @@ pub async fn handle_tauri_startup(
     &storyteller_creds_manager,
   )?;
 
+  tauri::async_runtime::spawn(midjourney_long_polling_thread(
+    app.clone(),
+    app_env_configs.clone(),
+    task_database.clone(),
+    mj_creds_manager.clone(),
+  ));
+
   spawn_discord_presence_thread()?;
 
   initially_size_and_position_windows(&app, &root);

crates/desktop/artcraft/src/lib.rs

@@ -91,6 +91,7 @@ pub fn run() {
   let app_env_configs_2 = app_env_configs.clone();
 
   let midjourney_creds_manager = MidjourneyCredentialManager::initialize_from_disk_infallible(&app_data_root);
+  let midjourney_creds_manager_2 = midjourney_creds_manager.clone();
 
   println!("Initializing backend runtime...");
 
@@ -129,6 +130,7 @@ pub fn run() {
           sora_tasks,
           fal_creds,
           fal_tasks,
+          midjourney_creds_manager_2,
         ).await;
 
         if let Err(err) = result {

crates/desktop/artcraft/src/services/midjourney/state/midjourney_credential_manager.rs

@@ -8,6 +8,7 @@ use errors::AnyhowResult;
 use log::{info, warn};
 use std::fs::read_to_string;
 use std::sync::{Arc, RwLock};
+use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
 use storyteller_client::credentials::storyteller_avt_cookie::StorytellerAvtCookie;
 use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
 use storyteller_client::credentials::storyteller_session_cookie::StorytellerSessionCookie;
@@ -93,7 +94,41 @@ impl MidjourneyCredentialManager {
       }
     }
   }
-  
+
+  // NB: This is just a heuristic. We'll add better checks later.
+  pub fn session_appears_active(&self) -> anyhow::Result<bool> {
+    let maybe_cookies = match self.cookies.read() {
+      Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
+      Ok(store) => store.clone(),
+    };
+
+    let maybe_user_info = match self.user_info.read() {
+      Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
+      Ok(info) => info.clone(),
+    };
+
+    if maybe_cookies.is_none() || maybe_user_info.is_none() {
+      return Ok(false);
+    }
+
+    let cookies = maybe_cookies.unwrap();
+    let user_info = maybe_user_info.unwrap();
+
+    // NB: We consider the session active if we have auth cookies and a user id.
+    let has_user_id = user_info.user_id.is_some();
+    let maybe_has_auth_cookies = cookie_store_has_auth_cookies(&cookies);
+
+    // Misc cookies without login cookies are ~1055 length
+    // AUTH_I is ~1500 length
+    // AUTH_R is ~500 length
+    let maybe_has_big_enough_cookie = cookies.calculate_approx_cookie_character_length() > 2100;
+
+    // TODO: Heuristic should count.
+    // TODO: Consolidate with "login window thread" logic.
+    // TODO: Check timestamp of last valid request.
+    Ok(has_user_id && maybe_has_auth_cookies && maybe_has_big_enough_cookie && maybe_has_big_enough_cookie)
+  }
+
   pub fn clear_credentials(&self) -> anyhow::Result<()> {
     match self.cookies.write() {
       Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),

crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs

@@ -0,0 +1,83 @@
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::core::state::task_database::TaskDatabase;
+use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+use errors::AnyhowResult;
+use log::error;
+use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
+use midjourney_client::endpoints::imagine::{imagine, ImagineRequest};
+use tauri::AppHandle;
+
+/// This thread is responsible for picking up tasks that fell through the cracks of
+/// the faster websocket thread.
+pub async fn midjourney_long_polling_thread(
+  app_handle: AppHandle,
+  app_env_configs: AppEnvConfigs,
+  task_database: TaskDatabase,
+  creds: MidjourneyCredentialManager,
+) -> ! {
+  loop {
+    let res = polling_loop(
+      &app_handle,
+      &app_env_configs,
+      &task_database,
+      &creds,
+    ).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,
+  creds: &MidjourneyCredentialManager,
+) -> AnyhowResult<()> {
+  loop {
+    if !creds.session_appears_active()? {
+      tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
+      continue;
+    }
+
+    let cookies = creds.maybe_copy_cookie_store()?;
+
+    let cookies = match cookies {
+      Some(cookies) => cookies,
+      None => {
+        tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
+        continue;
+      }
+    };
+
+    let user_info = creds.maybe_copy_user_info()?;
+
+    let maybe_user_id = user_info
+        .map(|info| info.user_id)
+        .flatten();
+
+    let user_id = match maybe_user_id {
+      Some(user_id) => user_id,
+      None => {
+        tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
+        continue;
+      }
+    };
+
+    let cookie_header = cookies.to_cookie_string();
+
+    let result = imagine(ImagineRequest {
+      hostname: MidjourneyHostname::Standard,
+      cookie_header,
+      user_id,
+      page_size: None,
+    }).await?;
+
+    println!("Response: {:?}\n\n", result);
+
+    tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
+  }
+}

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

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