diff --git a/crates/desktop/artcraft/src/core/events/functional_events/mod.rs b/crates/desktop/artcraft/src/core/events/functional_events/mod.rs index 33fc716932..b4253b3ecb 100644 --- a/crates/desktop/artcraft/src/core/events/functional_events/mod.rs +++ b/crates/desktop/artcraft/src/core/events/functional_events/mod.rs @@ -1 +1,2 @@ -pub mod canvas_background_removal_complete_event; \ No newline at end of file +pub mod canvas_background_removal_complete_event; +pub mod refresh_account_state_event; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/events/functional_events/refresh_account_state_event.rs b/crates/desktop/artcraft/src/core/events/functional_events/refresh_account_state_event.rs new file mode 100644 index 0000000000..56d55ad94f --- /dev/null +++ b/crates/desktop/artcraft/src/core/events/functional_events/refresh_account_state_event.rs @@ -0,0 +1,19 @@ +use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent}; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::ux::tauri_event_name::TauriEventName; +use serde_derive::Serialize; + +/// Send a signal to the frontend that an account state has changed: login, +/// logout, invalidated, etc. +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct RefreshAccountStateEvent { + /// The frontend can optionally decide whether it wants to refresh based + /// on the impacted provider account. + pub provider: Option, +} + +impl BasicSendableEvent for RefreshAccountStateEvent { + const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::RefreshAccountStateEvent; + const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success; +} diff --git a/crates/desktop/artcraft/src/core/utils/window/get_webview_window_hostname.rs b/crates/desktop/artcraft/src/core/utils/window/get_webview_window_hostname.rs new file mode 100644 index 0000000000..4fd10cd58f --- /dev/null +++ b/crates/desktop/artcraft/src/core/utils/window/get_webview_window_hostname.rs @@ -0,0 +1,11 @@ +use anyhow::anyhow; +use tauri::WebviewWindow; +use errors::AnyhowResult; + +pub fn get_webview_window_hostname(webview: &WebviewWindow) -> AnyhowResult { + let url = webview.url()?; + let url_hostname= url.host() + .ok_or(anyhow!("no host in url"))? + .to_string(); + Ok(url_hostname) +} diff --git a/crates/desktop/artcraft/src/core/utils/window/mod.rs b/crates/desktop/artcraft/src/core/utils/window/mod.rs index a6de92d14f..11abd59ee4 100644 --- a/crates/desktop/artcraft/src/core/utils/window/mod.rs +++ b/crates/desktop/artcraft/src/core/utils/window/mod.rs @@ -1,3 +1,4 @@ +pub mod get_webview_window_hostname; pub mod get_window_size_heuristic; pub mod position_main_window; pub mod resize_main_window; diff --git a/crates/desktop/artcraft/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs b/crates/desktop/artcraft/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs index bc6fd814ea..9c348c2de8 100644 --- a/crates/desktop/artcraft/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs +++ b/crates/desktop/artcraft/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs @@ -1,8 +1,8 @@ -use std::collections::HashSet; use cookie_store::cookie_store::CookieStore; use errors::AnyhowResult; use once_cell::sync::Lazy; use reqwest::Url; +use std::collections::HashSet; use tauri::webview::Cookie; use tauri::WebviewWindow; diff --git a/crates/desktop/artcraft/src/services/midjourney/windows/midjourney_login_window_thread.rs b/crates/desktop/artcraft/src/services/midjourney/windows/midjourney_login_window_thread.rs index 8445553dd1..55a267eb51 100644 --- a/crates/desktop/artcraft/src/services/midjourney/windows/midjourney_login_window_thread.rs +++ b/crates/desktop/artcraft/src/services/midjourney/windows/midjourney_login_window_thread.rs @@ -1,5 +1,10 @@ +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::functional_events::refresh_account_state_event::RefreshAccountStateEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; use crate::core::events::sendable_event_trait::SendableEvent; use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::utils::window::get_webview_window_hostname::get_webview_window_hostname; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo; use crate::services::midjourney::windows::extract_midjourney_webview_cookies::extract_midjourney_webview_cookies; @@ -9,6 +14,7 @@ use crate::services::sora::state::sora_credential_manager::SoraCredentialManager use crate::services::sora::windows::sora_login_window::extract_sora_webview_cookies::extract_sora_webview_cookies; use anyhow::anyhow; use cookie_store::cookie_store::CookieStore; +use enums::common::generation_provider::GenerationProvider; use errors::AnyhowResult; use log::{error, info}; use midjourney_client::client::midjourney_hostname::MidjourneyHostname; @@ -77,7 +83,7 @@ async fn check_login_window( 3. Done / Landing: https://www.midjourney.com/... */ - let hostname = get_hostname(webview_window)?; + let hostname = get_webview_window_hostname(webview_window)?; let mut maybe_at_destination = false; @@ -166,13 +172,14 @@ async fn check_login_window( return Ok(false); } + let event = RefreshAccountStateEvent { + provider: Some(GenerationProvider::Midjourney), + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send RefreshAccountStateEvent: {:?}", err); // Fail open + } + Ok(true) } -fn get_hostname(webview: &WebviewWindow) -> AnyhowResult { - let url = webview.url()?; - let url_hostname= url.host() - .ok_or(anyhow!("no host in url"))? - .to_string(); - Ok(url_hostname) -} diff --git a/crates/desktop/artcraft/src/services/sora/windows/sora_login_window/sora_login_thread.rs b/crates/desktop/artcraft/src/services/sora/windows/sora_login_window/sora_login_thread.rs index d6fe3581a3..d85a3133fa 100644 --- a/crates/desktop/artcraft/src/services/sora/windows/sora_login_window/sora_login_thread.rs +++ b/crates/desktop/artcraft/src/services/sora/windows/sora_login_window/sora_login_thread.rs @@ -1,5 +1,6 @@ use crate::core::events::sendable_event_trait::SendableEvent; use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::utils::window::get_webview_window_hostname::get_webview_window_hostname; use crate::services::sora::events::sora_login_success_event::SoraLoginSuccessEvent; use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; use crate::services::sora::windows::sora_login_window::extract_sora_webview_cookies::extract_sora_webview_cookies; @@ -73,7 +74,7 @@ async fn check_login_window( - Check if we're on the correct domain, if not exit? (or inverse, that we're not in login flow) */ - let hostname = get_hostname(webview_window)?; + let hostname = get_webview_window_hostname(webview_window)?; let mut maybe_at_destination = false; @@ -138,11 +139,3 @@ async fn check_login_window( Ok(true) } - -fn get_hostname(webview: &WebviewWindow) -> AnyhowResult { - let url = webview.url()?; - let url_hostname= url.host() - .ok_or(anyhow!("no host in url"))? - .to_string(); - Ok(url_hostname) -} diff --git a/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs b/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs index 68eb5dbcc3..8e2253e74f 100644 --- a/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs +++ b/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs @@ -42,6 +42,11 @@ pub enum TauriEventName { /// Background removal complete #[serde(rename = "canvas_bg_removed_event")] CanvasBgRemovedEvent, + + /// Special event: + /// Refresh account states + #[serde(rename = "refresh_account_state_event")] + RefreshAccountStateEvent, } impl_enum_display_and_debug_using_to_str!(TauriEventName); @@ -58,6 +63,7 @@ impl TauriEventName { Self::GenerationCompleteEvent => "generation-complete-event", Self::GenerationFailedEvent => "generation-failed-event", Self::CanvasBgRemovedEvent => "canvas_bg_removed_event", + Self::RefreshAccountStateEvent => "refresh_account_state_event", } } @@ -68,6 +74,7 @@ impl TauriEventName { "generation-complete-event" => Ok(Self::GenerationCompleteEvent), "generation-failed-event" => Ok(Self::GenerationFailedEvent), "canvas_bg_removed_event" => Ok(Self::CanvasBgRemovedEvent), + "refresh_account_state_event" => Ok(Self::RefreshAccountStateEvent), _ => Err(format!("invalid tauri_event_name: {:?}", job_status)), } } @@ -81,6 +88,7 @@ impl TauriEventName { Self::GenerationCompleteEvent, Self::GenerationFailedEvent, Self::CanvasBgRemovedEvent, + Self::RefreshAccountStateEvent, ]) } } @@ -100,6 +108,7 @@ mod tests { assert_serialization(TauriEventName::GenerationCompleteEvent, "generation-complete-event"); assert_serialization(TauriEventName::GenerationFailedEvent, "generation-failed-event"); assert_serialization(TauriEventName::CanvasBgRemovedEvent, "canvas_bg_removed_event"); + assert_serialization(TauriEventName::RefreshAccountStateEvent, "refresh_account_state_event"); } #[test] @@ -109,9 +118,7 @@ mod tests { assert_eq!(TauriEventName::GenerationCompleteEvent.to_str(), "generation-complete-event"); assert_eq!(TauriEventName::GenerationFailedEvent.to_str(), "generation-failed-event"); assert_eq!(TauriEventName::CanvasBgRemovedEvent.to_str(), "canvas_bg_removed_event"); - - - + assert_eq!(TauriEventName::RefreshAccountStateEvent.to_str(), "refresh_account_state_event"); } #[test] @@ -121,17 +128,19 @@ mod tests { assert_eq!(TauriEventName::from_str("generation-complete-event").unwrap(), TauriEventName::GenerationCompleteEvent); assert_eq!(TauriEventName::from_str("generation-failed-event").unwrap(), TauriEventName::GenerationFailedEvent); assert_eq!(TauriEventName::from_str("canvas_bg_removed_event").unwrap(), TauriEventName::CanvasBgRemovedEvent); + assert_eq!(TauriEventName::from_str("refresh_account_state_event").unwrap(), TauriEventName::RefreshAccountStateEvent); } #[test] fn all_variants() { let mut variants = TauriEventName::all_variants(); - assert_eq!(variants.len(), 5); + assert_eq!(variants.len(), 6); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueSuccessEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueFailureEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationCompleteEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationFailedEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::CanvasBgRemovedEvent)); + assert_eq!(variants.pop_first(), Some(TauriEventName::RefreshAccountStateEvent)); assert_eq!(variants.pop_first(), None); } }