From 16a55e429a74a4340ad8c9e5694de18d88ec9898 Mon Sep 17 00:00:00 2001 From: Neel Date: Wed, 2 Sep 2026 13:50:06 +0000 Subject: [PATCH] language_models: Allow cancelling ChatGPT Subscription setup (#63627) On `main`, the only way to interrupt ChatGPT Subscription setup is to restart Zed. --- Release Notes: - Add a button to cancel ChatGPT Subscription setup --- .../src/provider/openai_subscribed.rs | 33 +++-- .../src/openai_subscribed.rs | 127 ++++++++++++++++-- 2 files changed, 143 insertions(+), 17 deletions(-) diff --git a/crates/language_models/src/provider/openai_subscribed.rs b/crates/language_models/src/provider/openai_subscribed.rs index 024489457f..bac7831a6d 100644 --- a/crates/language_models/src/provider/openai_subscribed.rs +++ b/crates/language_models/src/provider/openai_subscribed.rs @@ -211,8 +211,10 @@ impl Render for ConfigurationView { let last_auth_error = state.last_auth_error(); let provider_state = self.state.clone(); + let cancel_provider_state = self.state.clone(); let is_signing_in = state.is_signing_in(); + let is_sign_in_cancellable = state.is_sign_in_cancellable(); let button_label = if is_signing_in { "Signing in…" } else { @@ -225,14 +227,29 @@ impl Render for ConfigurationView { this.child(Label::new(SUBSCRIPTION_DESCRIPTION)) }) .child( - Button::new("sign-in", button_label) - .when(!self.compact, |this| this.full_width()) - .style(ButtonStyle::Outlined) - .size(ButtonSize::Medium) - .loading(is_signing_in) - .disabled(is_signing_in) - .on_click(move |_, _window, cx| { - provider_state.update(cx, |state, cx| state.sign_in(cx)); + h_flex() + .gap_2() + .child( + Button::new("sign-in", button_label) + .when(!self.compact, |this| this.full_width()) + .style(ButtonStyle::Outlined) + .size(ButtonSize::Medium) + .loading(is_signing_in) + .disabled(is_signing_in) + .on_click(move |_, _window, cx| { + provider_state.update(cx, |state, cx| state.sign_in(cx)); + }), + ) + .when(is_sign_in_cancellable, |this| { + this.child( + Button::new("cancel-sign-in", "Cancel") + .style(ButtonStyle::Subtle) + .size(ButtonSize::Medium) + .on_click(move |_, _window, cx| { + cancel_provider_state + .update(cx, |state, cx| state.cancel_sign_in(cx)); + }), + ) }), ) .when_some(last_auth_error, |this, error| { diff --git a/crates/openai_subscribed/src/openai_subscribed.rs b/crates/openai_subscribed/src/openai_subscribed.rs index 431bab98bd..9d8de4e093 100644 --- a/crates/openai_subscribed/src/openai_subscribed.rs +++ b/crates/openai_subscribed/src/openai_subscribed.rs @@ -62,9 +62,15 @@ impl CodexCredentials { } } +enum SignInState { + Idle, + Authorizing(Task>), + PersistingCredentials { _task: Task> }, +} + pub struct State { credentials: Option, - sign_in_task: Option>>, + sign_in_state: SignInState, refresh_task: Option>>>>, load_task: Option>>>>, credentials_provider: Arc, @@ -154,7 +160,7 @@ impl State { Self { credentials: None, - sign_in_task: None, + sign_in_state: SignInState::Idle, refresh_task: None, load_task: Some(load_task), credentials_provider, @@ -177,7 +183,22 @@ impl State { } pub fn is_signing_in(&self) -> bool { - self.sign_in_task.is_some() + !matches!(self.sign_in_state, SignInState::Idle) + } + + pub fn is_sign_in_cancellable(&self) -> bool { + matches!(self.sign_in_state, SignInState::Authorizing(_)) + } + + fn begin_persisting_credentials(&mut self, cx: &mut Context) { + let sign_in_state = std::mem::replace(&mut self.sign_in_state, SignInState::Idle); + self.sign_in_state = match sign_in_state { + SignInState::Authorizing(task) => { + cx.notify(); + SignInState::PersistingCredentials { _task: task } + } + sign_in_state => sign_in_state, + }; } pub fn last_auth_error(&self) -> Option { @@ -280,6 +301,10 @@ impl State { let task = cx.spawn(async move |this, cx| { match do_oauth_flow(http_client, cx).await { Ok(creds) => { + this.update(cx, |state, cx| { + state.begin_persisting_credentials(cx); + })?; + let persist_result = async { let credentials_provider = this.read_with(cx, |state, _| state.credentials_provider.clone())?; @@ -303,7 +328,7 @@ impl State { log::warn!("Failed to refresh ChatGPT models: {error:#}"); } this.update(cx, |state, cx| { - state.sign_in_task = None; + state.sign_in_state = SignInState::Idle; cx.notify(); })?; } @@ -312,7 +337,7 @@ impl State { "ChatGPT subscription sign-in failed to persist credentials: {err:?}" ); this.update(cx, |state, cx| { - state.sign_in_task = None; + state.sign_in_state = SignInState::Idle; state.last_auth_error = Some("Failed to save credentials. Please try again.".into()); cx.notify(); @@ -324,7 +349,7 @@ impl State { Err(err) => { log::error!("ChatGPT subscription sign-in failed: {err:?}"); this.update(cx, |state, cx| { - state.sign_in_task = None; + state.sign_in_state = SignInState::Idle; state.last_auth_error = Some("Sign-in failed. Please try again.".into()); cx.notify(); }) @@ -335,17 +360,24 @@ impl State { }); self.last_auth_error = None; - self.sign_in_task = Some(task); + self.sign_in_state = SignInState::Authorizing(task); cx.notify(); } + pub fn cancel_sign_in(&mut self, cx: &mut Context) { + if matches!(self.sign_in_state, SignInState::Authorizing(_)) { + self.sign_in_state = SignInState::Idle; + cx.notify(); + } + } + /// Clears credentials and in-flight work immediately (so observers see the /// sign-out right away); the returned task deletes the persisted /// credentials. pub fn sign_out(&mut self, cx: &mut Context) -> Task> { self.auth_generation += 1; self.credentials = None; - self.sign_in_task = None; + self.sign_in_state = SignInState::Idle; self.refresh_task = None; self.last_auth_error = None; self.reset_model_catalog(); @@ -1510,6 +1542,83 @@ mod tests { }); } + #[gpui::test] + async fn test_cancel_sign_in_drops_pending_task(cx: &mut TestAppContext) { + let http: Arc = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::default())?) + }); + let state = make_state(http, None, cx); + let (continue_sign_in_tx, continue_sign_in_rx) = futures::channel::oneshot::channel::<()>(); + + state.update(cx, |state, cx| { + let task = cx.spawn(async move |_this, _cx| { + continue_sign_in_rx.await?; + anyhow::Ok(()) + }); + state.sign_in_state = SignInState::Authorizing(task); + }); + + cx.read(|cx| assert!(state.read(cx).is_signing_in())); + state.update(cx, |state, cx| state.cancel_sign_in(cx)); + cx.run_until_parked(); + cx.read(|cx| assert!(!state.read(cx).is_signing_in())); + assert!( + continue_sign_in_tx.send(()).is_err(), + "canceling sign-in should drop the task" + ); + } + + #[gpui::test] + async fn test_sign_in_task_remains_alive_while_persisting_credentials(cx: &mut TestAppContext) { + let http: Arc = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::default())?) + }); + let state = make_state(http, None, cx); + let (begin_persisting_tx, begin_persisting_rx) = futures::channel::oneshot::channel::<()>(); + let (finish_persisting_tx, finish_persisting_rx) = + futures::channel::oneshot::channel::<()>(); + + state.update(cx, |state, cx| { + let task = cx.spawn(async move |this, cx| { + begin_persisting_rx.await?; + this.update(cx, |state, cx| { + state.begin_persisting_credentials(cx); + })?; + finish_persisting_rx.await?; + this.update(cx, |state, cx| { + state.sign_in_state = SignInState::Idle; + cx.notify(); + })?; + anyhow::Ok(()) + }); + state.sign_in_state = SignInState::Authorizing(task); + }); + + begin_persisting_tx + .send(()) + .expect("sign-in task should be waiting to persist credentials"); + cx.run_until_parked(); + cx.read(|cx| { + let state = state.read(cx); + assert!(state.is_signing_in()); + assert!(!state.is_sign_in_cancellable()); + }); + + state.update(cx, |state, cx| state.cancel_sign_in(cx)); + cx.run_until_parked(); + cx.read(|cx| assert!(state.read(cx).is_signing_in())); + + finish_persisting_tx + .send(()) + .expect("sign-in task should still be persisting credentials"); + cx.run_until_parked(); + cx.read(|cx| assert!(!state.read(cx).is_signing_in())); + } + #[gpui::test] async fn test_sign_out_during_refresh_discards_result(cx: &mut TestAppContext) { let (gate_tx, gate_rx) = futures::channel::oneshot::channel::<()>(); @@ -2219,7 +2328,7 @@ mod tests { ) -> Entity { cx.new(|_cx| State { credentials, - sign_in_task: None, + sign_in_state: SignInState::Idle, refresh_task: None, load_task: None, credentials_provider,