diff --git a/crates/api_clients/midjourney_client/src/endpoints/mod.rs b/crates/api_clients/midjourney_client/src/endpoints/mod.rs index 853a601f08..6abc6a8e51 100644 --- a/crates/api_clients/midjourney_client/src/endpoints/mod.rs +++ b/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; \ No newline at end of file +pub mod user_connections; \ No newline at end of file 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 56eaae0e23..e6b82f2d2e 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,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); diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 1a68fcef75..d5f34f4f68 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/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 { diff --git a/crates/desktop/artcraft/src/services/midjourney/state/midjourney_credential_manager.rs b/crates/desktop/artcraft/src/services/midjourney/state/midjourney_credential_manager.rs index 91523919ac..28b9f4cc88 100644 --- a/crates/desktop/artcraft/src/services/midjourney/state/midjourney_credential_manager.rs +++ b/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 { + 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)), diff --git a/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs b/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs new file mode 100644 index 0000000000..b56509e1e6 --- /dev/null +++ b/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; + } +} diff --git a/crates/desktop/artcraft/src/services/midjourney/threads/mod.rs b/crates/desktop/artcraft/src/services/midjourney/threads/mod.rs index e69de29bb2..13b80af3cc 100644 --- a/crates/desktop/artcraft/src/services/midjourney/threads/mod.rs +++ b/crates/desktop/artcraft/src/services/midjourney/threads/mod.rs @@ -0,0 +1 @@ +pub mod midjourney_long_polling_thread; \ No newline at end of file