diff --git a/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs b/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs index 9ab01960d9..d5c0341995 100644 --- a/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs +++ b/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs @@ -62,7 +62,7 @@ pub async fn begin_object_image_upload(args: BeginObjectImageUploadArgs<'_>) -> .header(ACCEPT_ENCODING, "gzip, deflate, br, zstd") .header(REFERER, REFERER_VALUE) .header(CONTENT_TYPE, CONTENT_TYPE_APPLICATION_JSON) - .header(AUTHORIZATION, args.bearer_token.to_bearer_token_string()) + .header(AUTHORIZATION, args.bearer_token.to_bearer_token_header_string()) .header(ORIGIN, ORIGIN_VALUE) .header(SEC_GPC, "1") .header(CONNECTION, CONNECTION_KEEP_ALIVE) diff --git a/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs b/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs index 05a4da77e9..11d9b9f3dc 100644 --- a/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs +++ b/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs @@ -54,7 +54,7 @@ pub async fn create_run_object(args: CreateRunObjectArgs<'_>) -> Result) -> Result .header(ACCEPT_ENCODING, "gzip, deflate, br, zstd") .header(REFERER, REFERER_VALUE) .header(CONTENT_TYPE, CONTENT_TYPE_APPLICATION_JSON) - .header(AUTHORIZATION, args.bearer_token.to_bearer_token_string()) + .header(AUTHORIZATION, args.bearer_token.to_bearer_token_header_string()) .header(ORIGIN, ORIGIN_VALUE) .header(SEC_GPC, "1") .header(CONNECTION, CONNECTION_KEEP_ALIVE) diff --git a/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs b/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs index 64aed8760f..dc9d2db97c 100644 --- a/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs +++ b/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs @@ -58,7 +58,7 @@ pub async fn recaption_image(args: RecaptionImageArgs<'_>) -> Result) -> Result) -> Result String { + pub fn to_bearer_token_header_string(&self) -> String { format!("Bearer {}", self.bearer_token) } } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs index ac76434588..926c472750 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs @@ -35,6 +35,7 @@ impl TaskEnqueueSuccess{ GenerationProvider::Grok => GenerationServiceProvider::Grok, GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney, GenerationProvider::Sora => GenerationServiceProvider::Sora, + GenerationProvider::WorldLabs => GenerationServiceProvider::WorldLabs, } } diff --git a/crates/desktop/artcraft/src/core/events/generation_events/common.rs b/crates/desktop/artcraft/src/core/events/generation_events/common.rs index eb6ec67c53..ebc9306dae 100644 --- a/crates/desktop/artcraft/src/core/events/generation_events/common.rs +++ b/crates/desktop/artcraft/src/core/events/generation_events/common.rs @@ -8,6 +8,7 @@ pub enum GenerationServiceProvider { Grok, Midjourney, Sora, + WorldLabs, } #[derive(Debug, Copy, Clone, Serialize)] diff --git a/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs b/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs index 1ec6cc92c3..991747a68e 100644 --- a/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs +++ b/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs @@ -9,5 +9,6 @@ pub fn to_generation_service_provider(provider: GenerationProvider) -> Generatio GenerationProvider::Grok => GenerationServiceProvider::Grok, GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney, GenerationProvider::Sora => GenerationServiceProvider::Sora, + GenerationProvider::WorldLabs => GenerationServiceProvider::WorldLabs, } } diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index ccf8ef10d0..d80c490314 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -32,8 +32,8 @@ use crate::core::threads::main_window_thread::main_window_thread::main_window_th use crate::services::grok::commands::grok_clear_credentials_command::grok_clear_credentials_command; use crate::services::grok::commands::grok_get_credential_info_command::grok_get_credential_info_command; use crate::services::grok::commands::grok_open_login_command::grok_open_login_command; -use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; use crate::services::midjourney::commands::midjourney_clear_credentials_command::midjourney_clear_credentials_command; use crate::services::midjourney::commands::midjourney_get_credential_info_command::midjourney_get_credential_info_command; use crate::services::midjourney::commands::midjourney_open_login_command::midjourney_open_login_command; @@ -55,6 +55,9 @@ use crate::services::storyteller::commands::stripe_customer_portal::storyteller_ use crate::services::storyteller::commands::stripe_customer_portal::storyteller_open_customer_portal_switch_plan_command::storyteller_open_customer_portal_switch_plan_command; use crate::services::storyteller::commands::stripe_customer_portal::storyteller_open_customer_portal_update_payment_method_command::storyteller_open_customer_portal_update_payment_method_command; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::worldlabs::commands::worldlabs_clear_credentials_command::worldlabs_clear_credentials_command; +use crate::services::worldlabs::commands::worldlabs_get_credential_info_command::worldlabs_get_credential_info_command; +use crate::services::worldlabs::commands::worldlabs_open_login_command::worldlabs_open_login_command; use log::error; use tauri_plugin_dialog; @@ -207,6 +210,9 @@ pub fn run() { storyteller_open_subscription_purchase_command, storyteller_purge_credentials_command, update_app_preferences_command, + worldlabs_clear_credentials_command, + worldlabs_get_credential_info_command, + worldlabs_open_login_command, ]); builder.run(tauri::generate_context!("tauri.conf.json")) diff --git a/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_get_credential_info_command.rs b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_get_credential_info_command.rs index ea9ce016a9..33386f1899 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_get_credential_info_command.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_get_credential_info_command.rs @@ -1,10 +1,10 @@ use crate::core::commands::response::shorthand::{Response, ResponseOrErrorMessage}; use crate::core::commands::response::success_response_wrapper::SerializeMarker; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; use errors::AnyhowResult; use log::{error, info}; use serde_derive::Serialize; use tauri::State; -use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; #[derive(Serialize)] pub struct WorldlabsGetCredentialInfoResponse { @@ -35,19 +35,19 @@ fn get_info( let mut can_clear_state = true; let maybe_cookies = creds.maybe_copy_cookie_store()?; - let maybe_full_credentials = creds.maybe_copy_full_credentials()?; - - if maybe_cookies.is_none() && maybe_full_credentials.is_none() { - can_clear_state = false; - } - - let maybe_email = maybe_full_credentials - .map(|full_creds| full_creds.client_secrets.user_email) - .map(|maybe_email| maybe_email.map(|email| email.to_string())) - .flatten(); + //let maybe_full_credentials = creds.user()?; + // + //if maybe_cookies.is_none() && maybe_full_credentials.is_none() { + // can_clear_state = false; + //} + // + //let maybe_email = maybe_full_credentials + // .map(|full_creds| full_creds.client_secrets.user_email) + // .map(|maybe_email| maybe_email.map(|email| email.to_string())) + // .flatten(); Ok(WorldlabsGetCredentialInfoResponse { - maybe_email, + maybe_email: None, can_clear_state, }) } diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/state/mod.rs index 17f409ec37..5f55a7c416 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/state/mod.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/state/mod.rs @@ -1,7 +1,5 @@ pub mod worldlabs_credential_manager; -pub mod worldlabs_image_prompt_queue; pub mod worldlabs_serializable_state; -pub mod worldlabs_websocket_manager; // Private mod worldlabs_credential_holder; diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_holder.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_holder.rs index 23182a2e19..04d82a6522 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_holder.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_holder.rs @@ -1,5 +1,6 @@ use cookie_store::cookie_store::CookieStore; -use grok_client::credentials::grok_full_credentials::GrokFullCredentials; +use world_labs_client::credentials::world_labs_bearer_token::WorldLabsBearerToken; +use world_labs_client::credentials::world_labs_cookies::WorldLabsCookies; #[derive(Clone)] pub struct WorldlabsCredentialHolder { @@ -8,16 +9,21 @@ pub struct WorldlabsCredentialHolder { /// The Worldlabs client consumes a string-only form (rather than this cookie jar) pub browser_cookies: Option, - /// Full credentials. + /// Copy of cookies, specifically for the client library. /// NOT PERSISTED TO DISK. - pub worldlabs_full_credentials: Option, + pub world_labs_cookies: Option, + + /// Bearer token + /// This needs to be persisted to disk. + pub world_labs_bearer_token: Option, } impl WorldlabsCredentialHolder { pub fn empty() -> Self { Self { browser_cookies: None, - worldlabs_full_credentials: None, + world_labs_cookies: None, + world_labs_bearer_token: None, } } } diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs index a6b2c8e9cd..59e0d99c37 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs @@ -1,12 +1,13 @@ use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::services::worldlabs::state::worldlabs_credential_holder::WorldlabsCredentialHolder; +use crate::services::worldlabs::state::worldlabs_serializable_state::{WorldlabsSerializableState, SERIALIZABLE_WORLDLABS_STATE_VERSION}; use cookie_store::cookie_store::CookieStore; use errors::AnyhowResult; use log::{info, warn}; use std::fs::read_to_string; use std::sync::{Arc, RwLock}; +use world_labs_client::credentials::world_labs_bearer_token::WorldLabsBearerToken; use world_labs_client::credentials::world_labs_cookies::WorldLabsCookies; -use crate::services::worldlabs::state::worldlabs_credential_holder::WorldlabsCredentialHolder; -use crate::services::worldlabs::state::worldlabs_serializable_state::{WorldlabsSerializableState, SERIALIZABLE_WORLDLABS_STATE_VERSION}; #[derive(Clone)] pub struct WorldlabsCredentialManager { @@ -38,17 +39,22 @@ impl WorldlabsCredentialManager { credential_data = Arc::new(RwLock::new(WorldlabsCredentialHolder::empty())); } Ok(Some(state)) => { - let maybe_cookies = state.user_cookies + let maybe_browser_cookies = state.user_cookies + .as_ref() .map(|cookies| cookies.to_cookie_store()); - - let mut user_data = None; - if let Some(user_id) = state.user_id { - user_data = Some(WorldlabsUserData::from_id_and_email(user_id, state.user_email)); - } + + let maybe_bearer = state.bearer_token + .map(|bearer| WorldLabsBearerToken::new(bearer)); + + let maybe_cookies = maybe_browser_cookies + .as_ref() + .map(|cookies| cookies.to_cookie_string()) + .map(|cookies| WorldLabsCookies::new(cookies)); credential_data = Arc::new(RwLock::new(WorldlabsCredentialHolder { - browser_cookies: maybe_cookies, - worldlabs_full_credentials: None, // NB: We don't want to keep this on disk. It goes stale. + browser_cookies: maybe_browser_cookies, + world_labs_bearer_token: maybe_bearer, + world_labs_cookies: maybe_cookies, })); } }; @@ -61,7 +67,7 @@ impl WorldlabsCredentialManager { pub fn maybe_copy_cookie_store(&self) -> anyhow::Result> { match self.credential_data.read() { - Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), Ok(holder) => Ok(holder.browser_cookies.clone()), } } @@ -82,10 +88,10 @@ impl WorldlabsCredentialManager { Ok(maybe_cookies) } - pub fn maybe_copy_full_credentials(&self) -> anyhow::Result> { + pub fn maybe_copy_bearer_token(&self) -> anyhow::Result> { match self.credential_data.read() { Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), - Ok(holder) => Ok(holder.worldlabs_full_credentials.clone()), + Ok(holder) => Ok(holder.world_labs_bearer_token.clone()), } } @@ -99,27 +105,11 @@ impl WorldlabsCredentialManager { } } - pub fn replace_full_credentials(&self, creds: WorldlabsFullCredentials) -> anyhow::Result<()> { + pub fn replace_bearer_token(&self, creds: WorldLabsBearerToken) -> anyhow::Result<()> { match self.credential_data.write() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), Ok(mut holder) => { - holder.worldlabs_full_credentials = Some(creds); - Ok(()) - } - } - } - - pub fn replace_client_secrets_only(&self, secrets: WorldLabsClientSecrets) -> anyhow::Result<()> { - match self.credential_data.write() { - Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), - Ok(mut holder) => { - let mut maybe_full_creds = holder.worldlabs_full_credentials.clone(); - if let Some(creds) = maybe_full_creds.as_mut() { - creds.client_secrets = secrets; - } else { - warn!("No existing credentials to replace secrets into."); - } - holder.worldlabs_full_credentials = maybe_full_creds; + holder.world_labs_bearer_token = Some(creds); Ok(()) } } @@ -137,7 +127,7 @@ impl WorldlabsCredentialManager { Ok(store) => store.clone(), }; - if holder.worldlabs_full_credentials.is_some() { + if holder.world_labs_bearer_token.is_some() { return Ok(true); } @@ -185,17 +175,12 @@ impl WorldlabsCredentialManager { user_cookies: creds.browser_cookies .as_ref() .map(|cookies| cookies.to_serializable()), - user_id: creds.worldlabs_full_credentials - .as_ref() - .map(|creds| creds.client_secrets.user_id.to_string()), - user_email: creds.worldlabs_full_credentials - .as_ref() - .map(|data| data.client_secrets.user_email.as_ref()) - .flatten() - .map(|email| email.to_string()), + user_id: None, // TODO + user_email: None, // TODO + bearer_token: creds.world_labs_bearer_token.map(|token| token.to_raw_string()) }; - let path = self.app_data_root.credentials_dir().get_grok_state_path(); + let path = self.app_data_root.credentials_dir().get_worldlabs_state_path(); let serialized = serde_json::to_string(&state)?; std::fs::write(path, serialized)?; @@ -203,7 +188,7 @@ impl WorldlabsCredentialManager { // TODO: This is just for building the client and testing. if let Some(cookies) = creds.browser_cookies.as_ref() { let cookies_header = cookies.to_cookie_string(); - let path = self.app_data_root.credentials_dir().get_grok_cookies_path(); + let path = self.app_data_root.credentials_dir().get_worldlabs_state_path(); std::fs::write(path, cookies_header)?; } diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_serializable_state.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_serializable_state.rs index 036de888ac..09e75d1565 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_serializable_state.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_serializable_state.rs @@ -17,6 +17,9 @@ pub (super) struct WorldlabsSerializableState { #[serde(skip_serializing_if = "Option::is_none")] pub (super) user_cookies: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) bearer_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub (super) user_id: Option,