diff --git a/app/src/ai/chatgpt_auth.rs b/app/src/ai/chatgpt_auth.rs index ce5a746b..f82300d0 100644 --- a/app/src/ai/chatgpt_auth.rs +++ b/app/src/ai/chatgpt_auth.rs @@ -2,12 +2,17 @@ use std::io::{BufRead, BufReader, Write}; use std::net::{TcpListener, TcpStream}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use base64::Engine; use futures::channel::oneshot; +use futures::StreamExt; +use galaxy_agent_rig::ChatGPTCredentialResolver; use galaxyui::{Entity, ModelContext, SingletonEntity}; +use galaxyui_extras::secure_storage::AppContextExt as _; use instant::Instant; use rand::RngCore; use sha2::{Digest, Sha256}; @@ -21,6 +26,10 @@ const CALLBACK_HOST: &str = "localhost"; const CALLBACK_PORTS: [u16; 2] = [1455, 1457]; const CALLBACK_PATH: &str = "/auth/callback"; const CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60); +const AUTH_STORAGE_KEY: &str = "galaxy.chatgpt.oauth"; +const SIGN_IN_REQUIRED: &str = "ChatGPT sign-in required. Connect or reconnect ChatGPT in Providers."; +type Reply = oneshot::Sender>; +static RESOLVER: OnceLock>>> = OnceLock::new(); /// Current state of the local ChatGPT subscription connection. #[derive(Clone, Debug, PartialEq, Eq)] @@ -45,26 +54,42 @@ struct OAuthCallback { /// Coordinates browser-based OAuth authorization for ChatGPT subscriptions. pub(crate) struct ChatGPTAuthModel { state: ChatGPTAuthState, + record: Option, + refresh: RefreshCoordinator, + sign_in_generation: u64, + sign_in_cancel: Option>, } impl ChatGPTAuthModel { - pub(crate) fn new() -> Self { - let state = match has_or_import_auth_credentials() { - Ok(()) => ChatGPTAuthState::Connected, - Err(error) => { - log::debug!( - "[chatgpt/auth] No usable persisted ChatGPT credentials at startup: {error}" - ); - ChatGPTAuthState::NotConnected - } + pub(crate) fn new(ctx: &mut ModelContext) -> Self { + let (tx, rx) = async_channel::bounded(64); + *RESOLVER.get_or_init(|| Mutex::new(None)).lock().unwrap() = Some(tx); + ctx.spawn_stream_local(rx, |model, reply, ctx| model.resolve(reply, ctx), |_, _| {}); + let record = ctx.secure_storage().read_value(AUTH_STORAGE_KEY).ok() + .and_then(|json| serde_json::from_str::(&json).ok()) + .filter(|record| record.credentials(now()).is_ok() || record.can_refresh()); + let state = if record.is_some() { + ChatGPTAuthState::Connected + } else { + ChatGPTAuthState::NotConnected }; - Self { state } + Self { + state, + record, + refresh: RefreshCoordinator::default(), + sign_in_generation: 0, + sign_in_cancel: None, + } } #[cfg(test)] pub(crate) fn new_for_test() -> Self { Self { state: ChatGPTAuthState::NotConnected, + record: None, + refresh: RefreshCoordinator::default(), + sign_in_generation: 0, + sign_in_cancel: None, } } @@ -72,22 +97,26 @@ impl ChatGPTAuthModel { &self.state } - /// Reuses persisted credentials when possible, otherwise starts browser OAuth with PKCE. - pub(crate) fn connect(&mut self, ctx: &mut ModelContext) { + fn begin_sign_in(&mut self) -> Option { if matches!( self.state, ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken ) { - return; + return None; } - if has_or_import_auth_credentials().is_ok() { - self.state = ChatGPTAuthState::Connected; - ctx.emit(ChatGPTAuthModelEvent::StateChanged); - return; - } + self.refresh.invalidate(); + self.sign_in_generation += 1; + Some(self.sign_in_generation) + } - let (redirect_uri, callback_rx) = match start_oauth_callback_server() { + /// Explicit connection always starts a fresh browser OAuth flow with PKCE. + pub(crate) fn connect(&mut self, ctx: &mut ModelContext) { + let Some(generation) = self.begin_sign_in() else { + return; + }; + let cancelled = Arc::new(AtomicBool::new(false)); + let (redirect_uri, callback_rx) = match start_oauth_callback_server(cancelled.clone()) { Ok(server) => server, Err(error) => { self.fail(error, ctx); @@ -106,8 +135,8 @@ impl ChatGPTAuthModel { } }; - self.state = ChatGPTAuthState::AwaitingBrowser; - ctx.emit(ChatGPTAuthModelEvent::StateChanged); + self.sign_in_cancel = Some(cancelled); + self.set_state(ChatGPTAuthState::AwaitingBrowser, ctx); ctx.open_url(authorize_url.as_str()); let _ = ctx.spawn( @@ -117,6 +146,9 @@ impl ChatGPTAuthModel { })? }, move |model, callback, ctx| { + if model.sign_in_generation != generation { + return; + } let callback = match callback { Ok(callback) => callback, Err(error) => { @@ -129,19 +161,27 @@ impl ChatGPTAuthModel { return; } - model.state = ChatGPTAuthState::ExchangingToken; - ctx.emit(ChatGPTAuthModelEvent::StateChanged); + model.set_state(ChatGPTAuthState::ExchangingToken, ctx); let _ = ctx.spawn( async move { exchange_code_for_tokens(&callback.code, &code_verifier, &redirect_uri) .await }, - |model, result, ctx| match result { - Ok(()) => { - model.state = ChatGPTAuthState::Connected; - ctx.emit(ChatGPTAuthModelEvent::StateChanged); + move |model, result, ctx| { + if model.sign_in_generation != generation { + return; + } + match result.and_then(|record| { + save_auth_record(&record, ctx)?; + Ok(record) + }) { + Ok(record) => { + model.record = Some(record); + model.sign_in_cancel = None; + model.set_state(ChatGPTAuthState::Connected, ctx); + } + Err(error) => model.fail(error, ctx), } - Err(error) => model.fail(error, ctx), }, ); }, @@ -149,8 +189,77 @@ impl ChatGPTAuthModel { } fn fail(&mut self, error: String, ctx: &mut ModelContext) { - self.state = ChatGPTAuthState::Failed(error); - ctx.emit(ChatGPTAuthModelEvent::StateChanged); + if let Some(cancelled) = self.sign_in_cancel.take() { + cancelled.store(true, Ordering::SeqCst); + } + self.set_state(ChatGPTAuthState::Failed(error), ctx); + } + + fn set_state(&mut self, state: ChatGPTAuthState, ctx: &mut ModelContext) { + if self.state != state { + self.state = state; + ctx.emit(ChatGPTAuthModelEvent::StateChanged); + } + } + + pub(crate) fn cancel_sign_in(&mut self, ctx: &mut ModelContext) { + if let Some(cancelled) = self.sign_in_cancel.take() { + cancelled.store(true, Ordering::SeqCst); + self.sign_in_generation += 1; + self.refresh.invalidate(); + self.set_state(ChatGPTAuthState::NotConnected, ctx); + } + } + + fn resolve(&mut self, reply: Reply, ctx: &mut ModelContext) { + if reply.is_canceled() { + return; + } + if matches!(self.state, ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken) { + let _ = reply.send(Err("ChatGPT sign-in is in progress.".to_owned())); + return; + } + let Some(record) = self.record.as_ref() else { + let _ = reply.send(Err(SIGN_IN_REQUIRED.to_owned())); + return; + }; + if let Ok(credentials) = record.credentials(now()) { + let _ = reply.send(Ok(credentials)); + return; + } + if !record.can_refresh() { + let _ = reply.send(Err(SIGN_IN_REQUIRED.to_owned())); + self.fail(SIGN_IN_REQUIRED.to_owned(), ctx); + return; + } + let record = record.clone(); + let Some(generation) = self.refresh.begin(reply) else { + return; + }; + let _ = ctx.spawn( + async move { refresh_tokens(&record).await }, + move |model, result, ctx| { + let completed = model.refresh.complete(generation, result, |record| { + save_auth_record(record, ctx) + }); + match completed { + None => {} + Some(Ok(record)) => { + model.record = Some(record); + model.set_state(ChatGPTAuthState::Connected, ctx); + } + Some(Err(error)) => { + if error.reauthenticate { + model.record = None; + if save_auth_record(&AuthRecord::default(), ctx).is_err() { + log::warn!("[chatgpt/auth] Could not persist reconnect-required state"); + } + } + model.fail(error.message, ctx); + } + } + }, + ); } } @@ -194,6 +303,7 @@ fn random_hex_string(byte_count: usize) -> String { } fn start_oauth_callback_server( + cancelled: Arc, ) -> Result<(String, oneshot::Receiver>), String> { let mut bind_errors = Vec::new(); let mut bound = None; @@ -222,6 +332,10 @@ fn start_oauth_callback_server( .spawn(move || { let deadline = Instant::now() + CALLBACK_TIMEOUT; loop { + if cancelled.load(Ordering::SeqCst) { + let _ = callback_tx.send(Err("ChatGPT sign-in was cancelled.".to_owned())); + return; + } match listener.accept() { Ok((stream, _)) => { let _ = callback_tx.send(handle_callback_connection(stream)); @@ -323,26 +437,19 @@ fn parse_callback_request_line(request_line: &str) -> Result code = Some(value.into_owned()), "state" if !value.is_empty() => state = Some(value.into_owned()), - "error" if !value.is_empty() => oauth_error = Some(value.into_owned()), - "error_description" if !value.is_empty() => { - error_description = Some(value.into_owned()) - } + "error" if !value.is_empty() => oauth_error = true, "code" | "state" | "error" | "error_description" => {} _ => {} } } - if let Some(error) = oauth_error { - return Err(format!( - "ChatGPT authorization failed: {error} ({})", - error_description.as_deref().unwrap_or("no description") - )); + if oauth_error { + return Err("ChatGPT browser authorization failed. Please try signing in again.".to_owned()); } Ok(OAuthCallback { code: code.ok_or("ChatGPT OAuth callback did not include an authorization code.")?, @@ -358,157 +465,101 @@ fn callback_response_page(title: &str, message: &str, is_error: bool) -> String .replace("__ACCENT__", accent) } -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct ChatGPTAuthCredentials { - pub(crate) access_token: String, - pub(crate) account_id: Option, -} +pub(crate) use galaxy_agent_rig::ChatGPTSubscriptionCredentials as ChatGPTAuthCredentials; -/// Attempts to read tokens from `~/.codex/auth.json` and write them to Rig's auth file. -/// Returns `Ok(())` if credentials were found and successfully imported. -pub(crate) fn import_codex_credentials() -> Result<(), String> { - let codex_path = codex_auth_file_path().ok_or("Cannot determine codex auth path")?; - let bytes = std::fs::read(&codex_path).map_err(|e| format!("{e}"))?; - let doc: serde_json::Value = serde_json::from_slice(&bytes).map_err(|e| format!("{e}"))?; - - let tokens = doc.get("tokens").ok_or("No tokens object")?; - let access_token = tokens - .get("access_token") - .and_then(|value| value.as_str()) - .filter(|token| !token.is_empty()); - let refresh_token = tokens - .get("refresh_token") - .and_then(|value| value.as_str()) - .filter(|token| !token.is_empty()) - .map(ToOwned::to_owned); - if access_token.is_none() && refresh_token.is_none() { - return Err("Codex auth file does not contain an access or refresh token".to_string()); - } - let expires_at = access_token.and_then(extract_expiration_timestamp); - let id_token = tokens - .get("id_token") - .and_then(|v| v.as_str()) - .map(ToOwned::to_owned); - let account_id = tokens - .get("account_id") - .and_then(|v| v.as_str()) - .map(ToOwned::to_owned) - .or_else(|| extract_account_id(id_token.as_deref())) - .or_else(|| extract_account_id(access_token)); - - let record = AuthRecord { - access_token: access_token.map(ToOwned::to_owned), - refresh_token, - id_token, - expires_at, - account_id, - }; - - write_auth_file(&record) +pub(crate) fn credential_resolver() -> ChatGPTCredentialResolver { + Arc::new(|| Box::pin(load_or_refresh_auth_credentials())) } pub(crate) async fn load_or_refresh_auth_credentials() -> Result { - if let Ok(credentials) = load_auth_credentials() { - return Ok(credentials); - } - - let _ = import_codex_credentials(); - if let Ok(credentials) = load_auth_credentials() { - return Ok(credentials); - } - - galaxy_agent_rig::refresh_chatgpt_subscription_credentials().await?; - load_auth_credentials() -} - -fn has_or_import_auth_credentials() -> Result<(), String> { - if has_auth_credentials() { - return Ok(()); - } - import_codex_credentials()?; - if has_auth_credentials() { - Ok(()) - } else { - Err("ChatGPT credentials are missing an access or refresh token.".to_string()) + let (tx, rx) = oneshot::channel(); + RESOLVER.get_or_init(|| Mutex::new(None)).lock().unwrap() + .as_ref().ok_or("ChatGPT authorization service is unavailable.")? + .try_send(tx).map_err(|_| "ChatGPT authorization service is busy or unavailable.")?; + match futures::future::select( + Box::pin(rx), + Box::pin(async_io::Timer::after(Duration::from_secs(35))), + ).await { + futures::future::Either::Left((result, _)) => { + result.map_err(|_| "ChatGPT authorization service stopped.".to_owned())? + } + futures::future::Either::Right(_) => { + Err("ChatGPT credential refresh timed out. Retry or sign in again.".to_owned()) + } } } -fn has_auth_credentials() -> bool { - let Some(path) = auth_file_path() else { - return false; - }; - let Ok(bytes) = std::fs::read(path) else { - return false; - }; - let Ok(record) = serde_json::from_slice::(&bytes) else { - return false; - }; - record - .refresh_token - .as_deref() - .is_some_and(|token| !token.trim().is_empty()) - || record.access_token.as_deref().is_some_and(|token| { - !token.trim().is_empty() - && !token_is_expired( - record - .expires_at - .or_else(|| extract_expiration_timestamp(token)), - ) - }) +#[derive(Clone, Debug)] +struct AuthFailure { + message: String, + reauthenticate: bool, } -fn load_auth_credentials() -> Result { - let path = auth_file_path().ok_or("Cannot determine ChatGPT auth file path")?; - let bytes = std::fs::read(&path) - .map_err(|error| format!("Failed to read {}: {error}", path.display()))?; - let record: AuthRecord = serde_json::from_slice(&bytes) - .map_err(|error| format!("Failed to parse {}: {error}", path.display()))?; - let access_token = record - .access_token - .as_deref() - .filter(|token| !token.trim().is_empty()) - .ok_or("ChatGPT auth file does not contain an access token")?; +#[derive(Default)] +struct RefreshCoordinator { + generation: u64, + active: Option<(u64, Vec)>, +} - let expires_at = record - .expires_at - .or_else(|| extract_expiration_timestamp(access_token)); - if token_is_expired(expires_at) { - return Err("ChatGPT access token is expired".to_string()); +impl RefreshCoordinator { + fn begin(&mut self, reply: Reply) -> Option { + if let Some((_, replies)) = &mut self.active { + replies.push(reply); + return None; + } + self.generation += 1; + self.active = Some((self.generation, vec![reply])); + Some(self.generation) } - let account_id = record - .account_id - .clone() - .filter(|account_id| !account_id.trim().is_empty()) - .or_else(|| extract_account_id(record.id_token.as_deref())) - .or_else(|| extract_account_id(Some(access_token))); + fn invalidate(&mut self) { + self.generation += 1; + if let Some((_, replies)) = self.active.take() { + for reply in replies { + let _ = reply.send(Err("ChatGPT session was replaced. Please sign in again.".to_owned())); + } + } + } - Ok(ChatGPTAuthCredentials { - access_token: access_token.to_string(), - account_id, - }) -} - -fn codex_auth_file_path() -> Option { - if let Some(codex_home) = std::env::var_os("CODEX_HOME") { - return Some(std::path::PathBuf::from(codex_home).join("auth.json")); + fn complete( + &mut self, + generation: u64, + result: Result, + persist: impl FnOnce(&AuthRecord) -> Result<(), String>, + ) -> Option> { + if self.active.as_ref().is_none_or(|(active, _)| *active != generation) { + return None; + } + let (_, replies) = self.active.take()?; + let result = result.and_then(|record| { + let credentials = record.access_context().map_err(|message| AuthFailure { + message, reauthenticate: true, + })?; + persist(&record).map_err(|_| AuthFailure { + message: "Could not securely save refreshed ChatGPT credentials. Please sign in again.".to_owned(), + reauthenticate: true, + })?; + Ok((record, credentials)) + }); + for reply in replies { + let response = result.as_ref().map(|(_, credentials)| credentials.clone()) + .map_err(|error| error.message.clone()); + let _ = reply.send(response); + } + Some(result.map(|(record, _)| record)) } - std::env::var_os("HOME").map(|h| std::path::PathBuf::from(h).join(".codex").join("auth.json")) } -fn token_is_expired(expires_at: Option) -> bool { - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|duration| duration.as_secs() as i64) - .unwrap_or_default(); - expires_at.is_none_or(|expires_at| now >= expires_at - 60) +fn now() -> i64 { + std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_secs() as i64).unwrap_or_default() } async fn exchange_code_for_tokens( code: &str, code_verifier: &str, redirect_uri: &str, -) -> Result<(), String> { +) -> Result { let body = url::form_urlencoded::Serializer::new(String::new()) .append_pair("grant_type", "authorization_code") .append_pair("client_id", CHATGPT_CLIENT_ID) @@ -516,7 +567,34 @@ async fn exchange_code_for_tokens( .append_pair("redirect_uri", redirect_uri) .append_pair("code_verifier", code_verifier) .finish(); - let response = reqwest::Client::new() + let tokens = token_request(body).await.map_err(|error| error.message)?; + record_from_tokens(tokens, None, now()).map_err(|error| error.message) +} + +async fn refresh_tokens(previous: &AuthRecord) -> Result { + let refresh_token = previous.refresh_token.as_deref().filter(|token| !token.trim().is_empty()) + .ok_or_else(|| AuthFailure { + message: SIGN_IN_REQUIRED.to_owned(), + reauthenticate: true, + })?; + let body = url::form_urlencoded::Serializer::new(String::new()) + .append_pair("grant_type", "refresh_token") + .append_pair("client_id", CHATGPT_CLIENT_ID) + .append_pair("refresh_token", refresh_token) + .finish(); + let tokens = token_request(body).await?; + record_from_tokens(tokens, Some(previous), now()) +} + +async fn token_request(body: String) -> Result { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .redirect(reqwest::redirect::Policy::none()) + .build().map_err(|_| AuthFailure { + message: "Could not create the ChatGPT authorization client.".to_owned(), + reauthenticate: false, + })?; + let response = client .post(CHATGPT_TOKEN_URL) .header( reqwest::header::CONTENT_TYPE, @@ -525,67 +603,83 @@ async fn exchange_code_for_tokens( .body(body) .send() .await - .map_err(|error| format!("ChatGPT token exchange request failed: {error}"))?; - - let status = response.status(); - if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - return Err(format!("ChatGPT token exchange failed: {status} {body}")); + .map_err(|_| AuthFailure { + message: "ChatGPT authorization request failed. Retry the connection.".to_owned(), + reauthenticate: false, + })?; + let status = response.status().as_u16(); + let mut bytes = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|_| AuthFailure { + message: "Could not read the ChatGPT authorization response.".to_owned(), + reauthenticate: (200..300).contains(&status), + })?; + if bytes.len().saturating_add(chunk.len()) > 64 * 1024 { + return Err(AuthFailure { + message: "ChatGPT authorization returned an oversized response.".to_owned(), + reauthenticate: (200..300).contains(&status), + }); + } + bytes.extend_from_slice(&chunk); } - let tokens: TokenResponse = response - .json() - .await - .map_err(|error| format!("Could not parse the ChatGPT token response: {error}"))?; + parse_token_response(status, &bytes) +} + +fn parse_token_response(status: u16, bytes: &[u8]) -> Result { + if (200..300).contains(&status) { + return serde_json::from_slice(bytes).map_err(|_| AuthFailure { + message: "ChatGPT returned invalid credentials. Please sign in again.".to_owned(), + reauthenticate: true, + }); + } + let body = serde_json::from_slice::(bytes).unwrap_or_default(); + let code = body.get("error").and_then(|error| { + error.as_str().or_else(|| error.get("code").and_then(|code| code.as_str())) + }); + let reauthenticate = matches!(status, 401 | 403) + || (status == 400 && matches!(code, + Some("invalid_grant" | "refresh_token_reused" | "refresh_token_expired" | "refresh_token_invalid") + )); + let message = if reauthenticate { + format!("ChatGPT authorization expired or was rejected (HTTP {status}). Please sign in again.") + } else { + format!("ChatGPT authorization failed (HTTP {status}). Retry the connection.") + }; + Err(AuthFailure { message, reauthenticate }) +} + +fn record_from_tokens( + tokens: TokenResponse, + previous: Option<&AuthRecord>, + now: i64, +) -> Result { let expires_at = extract_expiration_timestamp(&tokens.access_token).or_else(|| { - tokens.expires_in.map(|expires_in| { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|duration| duration.as_secs() as i64) - .unwrap_or_default() - + expires_in as i64 - }) + tokens.expires_in.and_then(|expires| i64::try_from(expires).ok()) + .map(|expires| now.saturating_add(expires)) }); let account_id = extract_account_id(tokens.id_token.as_deref()) - .or_else(|| extract_account_id(Some(&tokens.access_token))); - write_auth_file(&AuthRecord { + .or_else(|| extract_account_id(Some(&tokens.access_token))) + .or_else(|| previous.and_then(|record| record.account_id.clone())); + let record = AuthRecord { access_token: Some(tokens.access_token), - refresh_token: tokens.refresh_token, - id_token: tokens.id_token, + refresh_token: tokens.refresh_token.or_else(|| previous.and_then(|record| record.refresh_token.clone())), + id_token: tokens.id_token.or_else(|| previous.and_then(|record| record.id_token.clone())), expires_at, account_id, - }) + }; + record.credentials(now).map_err(|_| AuthFailure { + message: "ChatGPT returned invalid or expired credentials. Please sign in again.".to_owned(), + reauthenticate: true, + })?; + Ok(record) } -fn write_auth_file(record: &AuthRecord) -> Result<(), String> { - let path = auth_file_path().ok_or("Cannot determine auth file path")?; - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent) - .map_err(|e| format!("Failed to create auth directory: {e}"))?; - } - let json = - serde_json::to_vec_pretty(record).map_err(|e| format!("Failed to serialize auth: {e}"))?; - std::fs::write(&path, json).map_err(|e| format!("Failed to write auth file: {e}"))?; - Ok(()) -} - -fn auth_file_path() -> Option { - #[cfg(target_os = "windows")] - { - std::env::var_os("APPDATA").map(|d| { - std::path::PathBuf::from(d) - .join("chatgpt") - .join("auth.json") - }) - } - #[cfg(not(target_os = "windows"))] - { - std::env::var_os("XDG_CONFIG_HOME") - .map(std::path::PathBuf::from) - .or_else(|| { - std::env::var_os("HOME").map(|h| std::path::PathBuf::from(h).join(".config")) - }) - .map(|d| d.join("chatgpt").join("auth.json")) - } +fn save_auth_record(record: &AuthRecord, ctx: &galaxyui::AppContext) -> Result<(), String> { + let json = serde_json::to_string(record) + .map_err(|_| "Could not encode ChatGPT credentials.".to_owned())?; + ctx.secure_storage().write_value_with_owner_only_fallback(AUTH_STORAGE_KEY, &json) + .map_err(|_| "Could not securely save ChatGPT credentials. Please sign in again.".to_owned()) } fn extract_expiration_timestamp(token: &str) -> Option { @@ -633,7 +727,7 @@ struct TokenResponse { expires_in: Option, } -#[derive(serde::Deserialize, serde::Serialize)] +#[derive(Clone, Default, serde::Deserialize, serde::Serialize)] struct AuthRecord { access_token: Option, refresh_token: Option, @@ -641,3 +735,41 @@ struct AuthRecord { expires_at: Option, account_id: Option, } + +impl std::fmt::Debug for AuthRecord { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChatGPTAuthRecord") + .field("credentials", &"") + .field("expires_at", &self.expires_at) + .finish() + } +} + +impl AuthRecord { + fn can_refresh(&self) -> bool { + self.refresh_token.as_deref().is_some_and(|token| !token.trim().is_empty()) + } + + fn access_context(&self) -> Result { + let access_token = self.access_token.as_deref().filter(|token| !token.trim().is_empty()) + .ok_or_else(|| SIGN_IN_REQUIRED.to_owned())?; + let account_id = self.account_id.clone().filter(|id| !id.trim().is_empty()) + .or_else(|| extract_account_id(self.id_token.as_deref())) + .or_else(|| extract_account_id(Some(access_token))); + Ok(ChatGPTAuthCredentials { access_token: access_token.to_owned(), account_id }) + } + + fn credentials(&self, now: i64) -> Result { + let expiry = self.expires_at.or_else(|| { + self.access_token.as_deref().and_then(extract_expiration_timestamp) + }); + if expiry.is_none_or(|expiry| now >= expiry.saturating_sub(60)) { + return Err(SIGN_IN_REQUIRED.to_owned()); + } + self.access_context() + } +} + +#[cfg(test)] +#[path = "chatgpt_auth_tests.rs"] +mod tests; diff --git a/app/src/ai/chatgpt_auth_tests.rs b/app/src/ai/chatgpt_auth_tests.rs new file mode 100644 index 00000000..280e69be --- /dev/null +++ b/app/src/ai/chatgpt_auth_tests.rs @@ -0,0 +1,171 @@ +use std::cell::Cell; + +use futures::FutureExt; + +use super::*; + +fn record(access: &str, refresh: &str) -> AuthRecord { + AuthRecord { + access_token: Some(access.to_owned()), + refresh_token: Some(refresh.to_owned()), + id_token: None, + expires_at: Some(10_000), + account_id: Some("test-account".to_owned()), + } +} + +#[test] +fn concurrent_requests_share_one_refresh_and_wait_for_persistence() { + let mut coordinator = RefreshCoordinator::default(); + let (first_tx, mut first_rx) = oneshot::channel(); + let (second_tx, mut second_rx) = oneshot::channel(); + let generation = coordinator.begin(first_tx).unwrap(); + assert_eq!(coordinator.begin(second_tx), None); + let saved = Cell::new(false); + let result = coordinator.complete(generation, Ok(record("rotated-access", "rotated-refresh")), |record| { + assert_eq!(first_rx.try_recv().unwrap(), None); + assert_eq!(second_rx.try_recv().unwrap(), None); + assert_eq!(record.refresh_token.as_deref(), Some("rotated-refresh")); + saved.set(true); + Ok(()) + }); + assert!(result.unwrap().is_ok()); + assert!(saved.get()); + let first = first_rx.now_or_never().unwrap().unwrap().unwrap(); + let second = second_rx.now_or_never().unwrap().unwrap().unwrap(); + assert_eq!(first, second); + assert_eq!(first.access_token, "rotated-access"); +} + +#[test] +fn failed_persistence_never_publishes_rotated_credentials() { + let mut coordinator = RefreshCoordinator::default(); + let (tx, rx) = oneshot::channel(); + let generation = coordinator.begin(tx).unwrap(); + let result = coordinator.complete(generation, Ok(record("rotated", "new-refresh")), |_| { + Err("storage unavailable".to_owned()) + }).unwrap().unwrap_err(); + assert!(result.reauthenticate); + assert!(rx.now_or_never().unwrap().unwrap().is_err()); +} + +#[test] +fn replaced_sign_in_rejects_late_refresh_without_overwriting_new_session() { + let mut coordinator = RefreshCoordinator::default(); + let (old_tx, old_rx) = oneshot::channel(); + let old = coordinator.begin(old_tx).unwrap(); + coordinator.invalidate(); + assert!(old_rx.now_or_never().unwrap().unwrap().is_err()); + let (new_tx, new_rx) = oneshot::channel(); + let current = coordinator.begin(new_tx).unwrap(); + assert_ne!(old, current); + assert!(coordinator.complete(old, Ok(record("stale", "stale-refresh")), |_| { + panic!("A stale refresh must not write credentials") + }).is_none()); + assert!(coordinator.complete(current, Ok(record("current", "current-refresh")), |_| Ok(())) + .unwrap().is_ok()); + assert_eq!(new_rx.now_or_never().unwrap().unwrap().unwrap().access_token, "current"); +} + +#[test] +fn reconnect_starts_a_new_generation_even_with_a_cached_refresh_token() { + let mut model = ChatGPTAuthModel::new_for_test(); + model.record = Some(record("old-access", "already-used-refresh")); + model.state = ChatGPTAuthState::Connected; + let first = model.begin_sign_in().unwrap(); + model.state = ChatGPTAuthState::AwaitingBrowser; + assert_eq!(model.begin_sign_in(), None); + model.state = ChatGPTAuthState::Failed("refresh_token_reused".to_owned()); + let replacement = model.begin_sign_in().unwrap(); + assert!(replacement > first); +} + +#[test] +fn rotated_tokens_replace_previous_tokens_and_omitted_refresh_tokens_are_preserved() { + let previous = record("old-access", "old-refresh"); + for refresh in [None, Some("new-refresh".to_owned())] { + let result = record_from_tokens(TokenResponse { + access_token: "new-access".to_owned(), + refresh_token: refresh.clone(), + id_token: None, + expires_in: Some(3600), + }, Some(&previous), 100).unwrap(); + assert_eq!(result.access_token.as_deref(), Some("new-access")); + assert_eq!(result.refresh_token, refresh.or_else(|| Some("old-refresh".to_owned()))); + assert_eq!(result.expires_at, Some(3700)); + assert_eq!(result.account_id.as_deref(), Some("test-account")); + } +} + +#[test] +fn rejected_refresh_tokens_require_reauthentication_without_echoing_tokens() { + for code in ["refresh_token_reused", "invalid_grant", "refresh_token_expired"] { + let body = serde_json::to_vec(&serde_json::json!({ + "error": {"code": code, "message": "sensitive-refresh-token"} + })).unwrap(); + let error = parse_token_response(401, &body).err().unwrap(); + assert!(error.reauthenticate); + assert!(error.message.contains("sign in")); + assert!(!error.message.contains("sensitive-refresh-token")); + } + let error = parse_token_response(400, br#"{"error":"invalid_grant"}"#).err().unwrap(); + assert!(error.reauthenticate); +} + +#[test] +fn transient_refresh_failure_preserves_the_session_and_does_not_save_rejected_data() { + let error = parse_token_response(503, br#"{"message":"sensitive-data"}"#).err().unwrap(); + assert!(!error.reauthenticate); + assert!(!error.message.contains("sensitive-data")); + let mut coordinator = RefreshCoordinator::default(); + let (tx, rx) = oneshot::channel(); + let generation = coordinator.begin(tx).unwrap(); + let result = coordinator.complete(generation, Err(error), |_| { + panic!("Transient failures must not overwrite stored credentials") + }).unwrap().unwrap_err(); + assert!(!result.reauthenticate); + assert!(rx.now_or_never().unwrap().unwrap().is_err()); +} + +#[test] +fn credential_debug_output_is_redacted_and_expiry_has_a_refresh_margin() { + let credentials = ChatGPTAuthCredentials { + access_token: "sensitive-access-token".to_owned(), + account_id: Some("account".to_owned()), + }; + assert!(!format!("{credentials:?}").contains("sensitive-access-token")); + let credentials = record("sensitive-access-token", "sensitive-refresh-token"); + assert!(!format!("{credentials:?}").contains("sensitive-access-token")); + assert!(!format!("{credentials:?}").contains("sensitive-refresh-token")); + assert!(credentials.credentials(100).is_ok()); + assert!(credentials.credentials(9940).is_err()); + assert!(AuthRecord::default().credentials(100).is_err()); +} + +#[test] +fn invalid_token_payloads_are_not_saved_or_exposed() { + assert!(parse_token_response(200, b"invalid-sensitive-data").is_err()); + assert!(record_from_tokens(TokenResponse { + access_token: String::new(), + refresh_token: Some("refresh".to_owned()), + id_token: None, + expires_in: Some(3600), + }, None, 100).is_err()); + assert!(record_from_tokens(TokenResponse { + access_token: "opaque-token-without-expiry".to_owned(), + refresh_token: None, + id_token: None, + expires_in: None, + }, None, 100).is_err()); +} + +#[test] +fn browser_authorization_keeps_pkce_and_state_validation() { + let url = build_authorize_url("http://localhost:1455/auth/callback", "challenge", "state").unwrap(); + let params = url.query_pairs().collect::>(); + assert_eq!(params.get("code_challenge_method").map(|v| v.as_ref()), Some("S256")); + assert_eq!(params.get("state").map(|v| v.as_ref()), Some("state")); + let callback = parse_callback_request_line("GET /auth/callback?code=code&state=state HTTP/1.1").unwrap(); + assert_eq!(callback.state, "state"); + assert!(parse_callback_request_line("GET /auth/callback?code=code HTTP/1.1").is_err()); +} diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index c1aa4d82..59f6ded1 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -215,9 +215,8 @@ impl CrosscheckReviewer { model: config.model.unwrap_or(model_id), reasoning_effort: config.reasoning_effort, max_output_tokens: Some(u64::from(REVIEWER_MAX_OUTPUT_TOKENS)), - auth_file: None, }, - ); + ).with_credential_resolver(crate::ai::chatgpt_auth::credential_resolver()); runtime .complete_text( prompt::CROSSCHECK_REVIEWER_SYSTEM_PROMPT.to_string(), diff --git a/app/src/ai/mod.rs b/app/src/ai/mod.rs index c956cd27..23d8236a 100644 --- a/app/src/ai/mod.rs +++ b/app/src/ai/mod.rs @@ -90,7 +90,7 @@ pub fn init(app: &mut AppContext) { #[cfg(not(target_family = "wasm"))] app.add_singleton_model(acp::AcpRuntimeModel::new); #[cfg(not(target_family = "wasm"))] - app.add_singleton_model(|_| chatgpt_auth::ChatGPTAuthModel::new()); + app.add_singleton_model(chatgpt_auth::ChatGPTAuthModel::new); #[cfg(not(target_family = "wasm"))] app.add_singleton_model(claude_oauth::ClaudeOAuthModel::new); blocklist::keyboard_navigable_buttons::init(app); diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 61172595..c4593a09 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -314,9 +314,8 @@ pub(crate) async fn provider_runtime_for_request( model, reasoning_effort: config.reasoning_effort, max_output_tokens: config.max_output_tokens.map(u64::from), - auth_file: None, }, - )), + ).with_credential_resolver(crate::ai::chatgpt_auth::credential_resolver())), OpenAIProviderKind::Anthropic => { let credential = config.api_key.unwrap_or_default(); let stored = crate::ai::claude_subscription::is_stored_credential(&credential); diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 179ed0c7..ef1cb90e 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -220,6 +220,7 @@ pub enum ProviderSetupViewAction { ToggleModel(usize), CycleModelCapability(usize, CapabilityKey), ConnectChatGPT, + CancelChatGPTSignIn, ConnectClaude, SignOutClaude, CancelClaudeSignIn, @@ -1590,10 +1591,10 @@ impl ProviderSetupView { ); } - if matches!( - state, - ChatGPTAuthState::NotConnected | ChatGPTAuthState::Failed(_) - ) { + let pending = matches!(state, + ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken + ); + { children.push( appearance .ui_builder() @@ -1601,10 +1602,21 @@ impl ProviderSetupView { ButtonVariant::Secondary, self.chatgpt_connect_mouse_state.clone(), ) - .with_text_label("Connect ChatGPT".to_owned()) + .with_text_label(if pending { + "Cancel sign-in" + } else if matches!(state, ChatGPTAuthState::Connected) { + "Reconnect ChatGPT" + } else { + "Connect ChatGPT" + }.to_owned()) .build() - .on_click(|ctx, _, _| { - ctx.dispatch_typed_action(ProviderSetupViewAction::ConnectChatGPT); + .on_click(move |ctx, _, _| { + let action = if pending { + ProviderSetupViewAction::CancelChatGPTSignIn + } else { + ProviderSetupViewAction::ConnectChatGPT + }; + ctx.dispatch_typed_action(action); }) .finish(), ); @@ -1861,7 +1873,7 @@ impl ProviderSetupView { .finish() } - fn render_discovery(&self, appearance: &Appearance) -> Box { + fn render_discovery(&self, appearance: &Appearance, app: &AppContext) -> Box { let (message, color) = match &self.discovery_state { DiscoveryState::Idle | DiscoveryState::Loading => ( if self.provider_type == ProviderSetupProviderType::Acp @@ -1875,7 +1887,7 @@ impl ProviderSetupView { ), DiscoveryState::Failed(error) => (error.clone(), appearance.theme().ui_error_color()), }; - Flex::column() + let mut column = Flex::column() .with_spacing(12.) .with_child( Text::new( @@ -1913,8 +1925,13 @@ impl ProviderSetupView { .with_color(color) .soft_wrap(true) .finish() - }) - .finish() + }); + if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription + && matches!(self.discovery_state, DiscoveryState::Failed(_)) + { + column = column.with_child(self.render_chatgpt_auth(appearance, app)); + } + column.finish() } fn render_model_table( @@ -2468,7 +2485,7 @@ impl View for ProviderSetupView { let content = match self.step { ProviderSetupStep::ProviderType => self.render_provider_type(appearance), ProviderSetupStep::Configure => self.render_configure(appearance, app), - ProviderSetupStep::Discover => self.render_discovery(appearance), + ProviderSetupStep::Discover => self.render_discovery(appearance, app), ProviderSetupStep::Models => self.render_models(appearance), }; Flex::column() @@ -2488,6 +2505,10 @@ impl TypedActionView for ProviderSetupView { match action { ProviderSetupViewAction::SelectProvider(kind) => { if self.provider_type != *kind { + #[cfg(not(target_family = "wasm"))] + if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription { + ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx)); + } self.cancel_claude_signin(); self.draft_api_key = None; self.claude_credentials_available = false; @@ -2588,6 +2609,10 @@ impl TypedActionView for ProviderSetupView { ProviderSetupViewAction::Cancel => { self.cancel_acp_setup(); self.cancel_claude_signin(); + #[cfg(not(target_family = "wasm"))] + if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription { + ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx)); + } self.discovery_state = DiscoveryState::Idle; ctx.emit(ProviderSetupViewEvent::Close); } @@ -2633,6 +2658,10 @@ impl TypedActionView for ProviderSetupView { #[cfg(not(target_family = "wasm"))] ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.connect(ctx)); } + ProviderSetupViewAction::CancelChatGPTSignIn => { + #[cfg(not(target_family = "wasm"))] + ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx)); + } ProviderSetupViewAction::ConnectClaude => { #[cfg(not(target_family = "wasm"))] self.connect_claude(ctx); diff --git a/crates/galaxy_agent_rig/src/chatgpt.rs b/crates/galaxy_agent_rig/src/chatgpt.rs index f4af1f7d..6036a266 100644 --- a/crates/galaxy_agent_rig/src/chatgpt.rs +++ b/crates/galaxy_agent_rig/src/chatgpt.rs @@ -1,6 +1,7 @@ -use std::path::PathBuf; +use std::sync::Arc; use async_trait::async_trait; +use futures::future::BoxFuture; use futures::StreamExt; use galaxy_agent_core::{ AgentError, AgentErrorKind, AgentEvent, AgentEventStream, AgentRuntime, ConversationMessage, @@ -8,35 +9,50 @@ use galaxy_agent_core::{ TurnRequest, turn_control, }; use rig_core::client::CompletionClient; -use rig_core::providers::chatgpt; +use rig_core::providers::chatgpt::{self, ChatGPTAuth}; use crate::request::build_completion_request; use crate::stream::start_model_turn; -/// Refreshes Galaxy's cached ChatGPT subscription credentials without allowing -/// Rig to fall back to its interactive device-code flow. -pub async fn refresh_chatgpt_subscription_credentials() -> Result<(), String> { - let client = chatgpt::Client::builder() - .oauth() - .allow_device_flow(false) - .originator("galaxy") - .build() - .map_err(|error| error.to_string())?; - client.authorize().await.map_err(|error| error.to_string()) +#[derive(Clone, PartialEq, Eq)] +pub struct ChatGPTSubscriptionCredentials { + pub access_token: String, + pub account_id: Option, } +impl std::fmt::Debug for ChatGPTSubscriptionCredentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChatGPTSubscriptionCredentials") + .field("access_token", &"") + .field("account_id", &self.account_id.as_ref().map(|_| "")) + .finish() + } +} +pub type ChatGPTCredentialResolver = Arc< + dyn Fn() -> BoxFuture<'static, Result> + Send + Sync, +>; + #[derive(Clone, Debug, PartialEq, Eq)] pub struct ChatGPTSubscriptionRuntimeConfig { pub model: String, pub reasoning_effort: Option, pub max_output_tokens: Option, - pub auth_file: Option, } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct ChatGPTSubscriptionRuntime { config: ChatGPTSubscriptionRuntimeConfig, descriptor: RuntimeDescriptor, + credential_resolver: Option, +} + +impl std::fmt::Debug for ChatGPTSubscriptionRuntime { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChatGPTSubscriptionRuntime") + .field("config", &self.config) + .field("has_credential_resolver", &self.credential_resolver.is_some()) + .finish() + } } fn reasoning_additional_params(reasoning_effort: Option<&str>) -> Option { @@ -56,7 +72,24 @@ impl ChatGPTSubscriptionRuntime { kind: RuntimeKind::Provider, capabilities: RuntimeCapabilities::provider(), }; - Self { config, descriptor } + Self { config, descriptor, credential_resolver: None } + } + + pub fn with_credential_resolver(mut self, resolver: ChatGPTCredentialResolver) -> Self { + self.credential_resolver = Some(resolver); + self + } + + async fn resolve_auth(&self) -> Result { + let resolver = self.credential_resolver.as_ref().ok_or_else(|| { + AgentError::new(AgentErrorKind::Authentication, "ChatGPT sign-in required. Reconnect ChatGPT in Providers.") + })?; + let credentials = resolver().await + .map_err(|message| AgentError::new(AgentErrorKind::Authentication, message))?; + Ok(ChatGPTAuth::AccessToken { + access_token: credentials.access_token, + account_id: credentials.account_id, + }) } /// Completes a text-only request for callers such as crosscheck reviewers. @@ -117,14 +150,11 @@ impl AgentRuntime for ChatGPTSubscriptionRuntime { request: TurnRequest, control: TurnControl, ) -> Result { - let mut builder = chatgpt::Client::builder() - .oauth() + let auth = self.resolve_auth().await?; + let client = chatgpt::Client::builder() + .api_key(auth) .allow_device_flow(false) - .originator("galaxy"); - if let Some(auth_file) = &self.config.auth_file { - builder = builder.auth_file(auth_file); - } - let client = builder + .originator("galaxy") .build() .map_err(|error| AgentError::new(AgentErrorKind::Configuration, error.to_string()))?; let model = client.completion_model(&self.config.model); diff --git a/crates/galaxy_agent_rig/src/chatgpt_tests.rs b/crates/galaxy_agent_rig/src/chatgpt_tests.rs index c176c402..192dce22 100644 --- a/crates/galaxy_agent_rig/src/chatgpt_tests.rs +++ b/crates/galaxy_agent_rig/src/chatgpt_tests.rs @@ -16,7 +16,6 @@ fn runtime_descriptor_identifies_chatgpt_subscription() { model: "gpt-5.3-codex".to_string(), reasoning_effort: Some("high".to_string()), max_output_tokens: None, - auth_file: None, }); assert_eq!( @@ -25,6 +24,43 @@ fn runtime_descriptor_identifies_chatgpt_subscription() { ); } +#[tokio::test] +async fn runtime_resolves_current_credentials_without_rig_oauth() { + let runtime = ChatGPTSubscriptionRuntime::new(ChatGPTSubscriptionRuntimeConfig { + model: "test-model".to_owned(), + reasoning_effort: None, + max_output_tokens: None, + }).with_credential_resolver(std::sync::Arc::new(|| { + Box::pin(async { + Ok(ChatGPTSubscriptionCredentials { + access_token: "resolved-token".to_owned(), + account_id: Some("resolved-account".to_owned()), + }) + }) + })); + match runtime.resolve_auth().await.unwrap() { + ChatGPTAuth::AccessToken { access_token, account_id } => { + assert_eq!(access_token, "resolved-token"); + assert_eq!(account_id.as_deref(), Some("resolved-account")); + } + ChatGPTAuth::OAuth => panic!("The runtime must not own OAuth refresh"), + } +} + +#[tokio::test] +async fn missing_or_failed_resolver_requires_sign_in_without_device_flow() { + let runtime = ChatGPTSubscriptionRuntime::new(ChatGPTSubscriptionRuntimeConfig { + model: "test-model".to_owned(), + reasoning_effort: None, + max_output_tokens: None, + }); + assert_eq!(runtime.resolve_auth().await.unwrap_err().kind, AgentErrorKind::Authentication); + let runtime = runtime.with_credential_resolver(std::sync::Arc::new(|| { + Box::pin(async { Err("Sign in again".to_owned()) }) + })); + assert_eq!(runtime.resolve_auth().await.unwrap_err().kind, AgentErrorKind::Authentication); +} + #[test] fn reasoning_effort_is_encoded_for_the_responses_request() { let request = TurnRequest::new(