diff --git a/crates/desktop/artcraft/src/services/grok/commands/grok_get_credential_info_command.rs b/crates/desktop/artcraft/src/services/grok/commands/grok_get_credential_info_command.rs --- a/crates/desktop/artcraft/src/services/grok/commands/grok_get_credential_info_command.rs +++ b/crates/desktop/artcraft/src/services/grok/commands/grok_get_credential_info_command.rs @@ -1,9 +1,11 @@ use crate::core::commands::response::shorthand::{Response, ResponseOrErrorMessage}; use crate::core::commands::response::success_response_wrapper::SerializeMarker; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; -use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; +use crate::services::grok::util::get_or_upgrade_grok_full_credentials::get_or_update_grok_full_credentials; use errors::AnyhowResult; -use log::{error, info}; +use grok_consumer_client::credentials::grok_full_credentials::GrokFullCredentials; +use grok_consumer_client::credentials::grok_user_data::GrokUserData; +use log::{error, info, warn}; use serde_derive::Serialize; use tauri::State; @@ -22,6 +24,7 @@ pub async fn grok_get_credential_info_command( info!("grok_get_credential_info_command called"); let info = get_info(&creds_manager) + .await .map_err(|err| { error!("Error getting info: {:?}", err); "error getting info" @@ -30,25 +33,57 @@ pub async fn grok_get_credential_info_command( Ok(info.into()) } -fn get_info( +async fn get_info( creds: &GrokCredentialManager, ) -> AnyhowResult { - 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_user_data = creds.maybe_copy_user_data()?; + + let can_clear_state = maybe_cookies.is_some() + || maybe_full_credentials.is_some() + || maybe_user_data.is_some(); + + let mut maybe_email = email_from_full_credentials(maybe_full_credentials.as_ref()) + .or_else(|| email_from_user_data(maybe_user_data.as_ref())); + + // Best-effort: cookies exist but we still lack an email — try upgrading secrets. + let needs_identity_upgrade = maybe_cookies.is_some() + && maybe_email.is_none() + && maybe_full_credentials.is_none(); + + if needs_identity_upgrade { + match get_or_update_grok_full_credentials(creds).await { + Ok(full_creds) => { + maybe_email = email_from_full_credentials(Some(&full_creds)); + } + Err(err) => { + warn!( + "Best-effort Grok secrets upgrade failed while reading credential info: {:?}", + err + ); + } + } } - - 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(GrokGetCredentialInfoResponse { maybe_email, can_clear_state, }) } + +fn email_from_full_credentials( + maybe_full_credentials: Option<&GrokFullCredentials>, +) -> Option { + maybe_full_credentials + .and_then(|full_creds| full_creds.client_secrets.user_email.as_ref()) + .map(|email| email.to_string()) +} + +fn email_from_user_data( + maybe_user_data: Option<&GrokUserData>, +) -> Option { + maybe_user_data + .and_then(|data| data.user_email.as_ref()) + .map(|email| email.to_string()) +} diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs b/crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs --- a/crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_credential_holder.rs @@ -1,5 +1,4 @@ use cookie_store::cookie_store::CookieStore; -use grok_consumer_client::credentials::grok_client_secrets::GrokClientSecrets; use grok_consumer_client::credentials::grok_full_credentials::GrokFullCredentials; use grok_consumer_client::credentials::grok_user_data::GrokUserData; @@ -10,6 +9,10 @@ pub struct GrokCredentialHolder { /// The Grok client consumes a string-only form (rather than this cookie jar) pub browser_cookies: Option, + /// Persisted identity (user id / email) restored from disk. + /// Survives restarts even when full credentials are cleared. + pub maybe_user_data: Option, + /// Full credentials. /// NOT PERSISTED TO DISK. pub grok_full_credentials: Option, @@ -19,6 +22,7 @@ impl GrokCredentialHolder { pub fn empty() -> Self { Self { browser_cookies: None, + maybe_user_data: None, grok_full_credentials: None, } } diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs --- a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs @@ -2,13 +2,11 @@ use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::services::grok::state::grok_credential_holder::GrokCredentialHolder; use crate::services::grok::state::grok_serializable_state::{GrokSerializableState, SERIALIZABLE_GROK_STATE_VERSION}; use cookie_store::cookie_store::CookieStore; -use errors::AnyhowResult; use grok_consumer_client::credentials::grok_client_secrets::GrokClientSecrets; use grok_consumer_client::credentials::grok_cookies::GrokCookies; use grok_consumer_client::credentials::grok_full_credentials::GrokFullCredentials; use grok_consumer_client::credentials::grok_user_data::GrokUserData; use log::{info, warn}; -use std::fs::read_to_string; use std::sync::{Arc, RwLock}; #[derive(Clone)] @@ -16,7 +14,7 @@ pub struct GrokCredentialManager { // TODO: Put last write/last read timestamps on these. credential_data: Arc>, - //user_info: Arc>>, + // user_info: Arc>>, app_data_root: AppDataRoot, } @@ -29,7 +27,7 @@ impl GrokCredentialManager { } pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { - let mut credential_data; + let credential_data; let result = GrokSerializableState::read_from_disk(app_data_root); match result { @@ -44,13 +42,13 @@ impl GrokCredentialManager { let maybe_cookies = state.user_cookies .map(|cookies| cookies.to_cookie_store()); - let mut user_data = None; - if let Some(user_id) = state.user_id { - user_data = Some(GrokUserData::from_id_and_email(user_id, state.user_email)); - } + let maybe_user_data = state.user_id.map(|user_id| { + GrokUserData::from_id_and_email(user_id, state.user_email) + }); credential_data = Arc::new(RwLock::new(GrokCredentialHolder { browser_cookies: maybe_cookies, + maybe_user_data, grok_full_credentials: None, // NB: We don't want to keep this on disk. It goes stale. })); } @@ -92,6 +90,13 @@ impl GrokCredentialManager { } } + pub fn maybe_copy_user_data(&self) -> anyhow::Result> { + match self.credential_data.read() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(holder) => Ok(holder.maybe_user_data.clone()), + } + } + pub fn replace_cookie_store(&self, store: CookieStore) -> anyhow::Result<()> { match self.credential_data.write() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), @@ -102,10 +107,24 @@ impl GrokCredentialManager { } } + pub fn replace_user_data(&self, user_data: GrokUserData) -> anyhow::Result<()> { + match self.credential_data.write() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut holder) => { + holder.maybe_user_data = Some(user_data); + Ok(()) + } + } + } + pub fn replace_full_credentials(&self, creds: GrokFullCredentials) -> anyhow::Result<()> { match self.credential_data.write() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), Ok(mut holder) => { + holder.maybe_user_data = Some(GrokUserData { + user_id: creds.client_secrets.user_id.clone(), + user_email: creds.client_secrets.user_email.clone(), + }); holder.grok_full_credentials = Some(creds); Ok(()) } @@ -116,6 +135,10 @@ impl GrokCredentialManager { match self.credential_data.write() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), Ok(mut holder) => { + holder.maybe_user_data = Some(GrokUserData { + user_id: secrets.user_id.clone(), + user_email: secrets.user_email.clone(), + }); let mut maybe_full_creds = holder.grok_full_credentials.clone(); if let Some(creds) = maybe_full_creds.as_mut() { creds.client_secrets = secrets; @@ -183,19 +206,33 @@ impl GrokCredentialManager { Ok(store) => store.clone(), }; + let user_id = creds.grok_full_credentials + .as_ref() + .map(|full| full.client_secrets.user_id.to_string()) + .or_else(|| { + creds.maybe_user_data + .as_ref() + .map(|data| data.user_id.to_string()) + }); + + let user_email = creds.grok_full_credentials + .as_ref() + .and_then(|full| full.client_secrets.user_email.as_ref()) + .map(|email| email.to_string()) + .or_else(|| { + creds.maybe_user_data + .as_ref() + .and_then(|data| data.user_email.as_ref()) + .map(|email| email.to_string()) + }); + let state = GrokSerializableState { version: SERIALIZABLE_GROK_STATE_VERSION, user_cookies: creds.browser_cookies .as_ref() .map(|cookies| cookies.to_serializable()), - user_id: creds.grok_full_credentials - .as_ref() - .map(|creds| creds.client_secrets.user_id.to_string()), - user_email: creds.grok_full_credentials - .as_ref() - .map(|data| data.client_secrets.user_email.as_ref()) - .flatten() - .map(|email| email.to_string()), + user_id, + user_email, }; let path = self.app_data_root.credentials_dir().get_grok_state_path(); @@ -210,7 +247,7 @@ impl GrokCredentialManager { std::fs::write(path, cookies_header)?; } + info!("Persisted grok credential state to disk"); Ok(()) } } - diff --git a/crates/desktop/artcraft/src/services/grok/windows/grok_login_window/grok_login_window_thread.rs b/crates/desktop/artcraft/src/services/grok/windows/grok_login_window/grok_login_window_thread.rs --- a/crates/desktop/artcraft/src/services/grok/windows/grok_login_window/grok_login_window_thread.rs +++ b/crates/desktop/artcraft/src/services/grok/windows/grok_login_window/grok_login_window_thread.rs @@ -1,24 +1,16 @@ 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::grok::state::grok_credential_manager::GrokCredentialManager; use crate::services::grok::windows::grok_login_window::grok_login_webview_extract_cookies::grok_login_webview_extract_cookies; use crate::services::grok::windows::grok_login_window::grok_login_window_open::GROK_LOGIN_WINDOW_NAME; -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; -use anyhow::anyhow; -use cookie_store::cookie_store::CookieStore; use enums::common::generation_provider::GenerationProvider; use errors::AnyhowResult; -use log::{error, info}; -use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; -use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session; -use openai_sora_client::utils::has_session_cookie::{has_session_cookie, SessionCookiePresence}; +use grok_consumer_client::credentials::grok_full_credentials::GrokFullCredentials; +use grok_consumer_client::recipes::request_client_secrets::{request_client_secrets, RequestClientSecretsArgs}; +use log::{error, info, warn}; use tauri::{AppHandle, Manager, WebviewWindow}; pub async fn grok_login_window_thread( @@ -169,6 +161,14 @@ async fn check_login_window( return Ok(false); } + // Best-effort identity upgrade after cookie save. Keep cookies even if scrape fails. + if let Err(err) = maybe_upgrade_grok_identity_after_login(grok_creds_manager).await { + warn!( + "Grok post-login secrets scrape failed; keeping cookies. err={:?}", + err + ); + } + let event = RefreshAccountStateEvent { provider: Some(GenerationProvider::Grok), }; @@ -180,3 +180,31 @@ async fn check_login_window( Ok(true) } +async fn maybe_upgrade_grok_identity_after_login( + grok_creds_manager: &GrokCredentialManager, +) -> AnyhowResult<()> { + let cookies = match grok_creds_manager.maybe_copy_typed_cookies()? { + Some(cookies) => cookies, + None => { + warn!("No Grok cookies available for post-login secrets upgrade"); + return Ok(()); + } + }; + + info!("Requesting Grok client secrets after login..."); + + let secrets = request_client_secrets(RequestClientSecretsArgs { + cookies: &cookies, + }) + .await + .map_err(|err| anyhow::anyhow!("request_client_secrets failed: {}", err))?; + + let full_creds = GrokFullCredentials::from_cookies_and_client_secrets(cookies, secrets); + // replace_full_credentials also syncs maybe_user_data from secrets. + grok_creds_manager.replace_full_credentials(full_creds)?; + grok_creds_manager.persist_to_disk()?; + + info!("Grok identity upgraded after login"); + Ok(()) +} +