Clear LLM token and model list on sign out (#57724)

Closes #FR-20

Release Notes:

- Fixed Zed provider models remaining available after sign-out

Cargo.lock

@@ -9735,6 +9735,7 @@ dependencies = [
  "clock",
  "cloud_api_client",
  "cloud_api_types",
+ "cloud_llm_client",
  "collections",
  "component",
  "convert_case 0.11.0",

crates/client/src/llm_token.rs

@@ -2,9 +2,10 @@ use super::{Client, UserStore};
 use cloud_api_client::LlmApiToken;
 use cloud_api_types::websocket_protocol::MessageToClient;
 use cloud_llm_client::{EXPIRED_LLM_TOKEN_HEADER_NAME, OUTDATED_LLM_TOKEN_HEADER_NAME};
+use futures::StreamExt;
 use gpui::{
     App, AppContext as _, Context, Entity, EventEmitter, Global, ReadGlobal as _, Subscription,
-    TaskExt,
+    Task, TaskExt,
 };
 use std::sync::Arc;
 
@@ -42,6 +43,7 @@ pub struct RefreshLlmTokenListener {
     client: Arc<Client>,
     user_store: Entity<UserStore>,
     llm_api_token: LlmApiToken,
+    _clear_llm_token_on_sign_out: Task<()>,
     _subscription: Subscription,
 }
 
@@ -73,10 +75,24 @@ impl RefreshLlmTokenListener {
             }
         });
 
+        let llm_api_token = LlmApiToken::default();
+        let mut status = client.status();
+        let clear_llm_token_on_sign_out = cx.spawn({
+            let llm_api_token = llm_api_token.clone();
+            async move |_this, _cx| {
+                while let Some(status) = status.next().await {
+                    if status.is_signed_out() {
+                        llm_api_token.clear().await;
+                    }
+                }
+            }
+        });
+
         Self {
             client,
             user_store,
-            llm_api_token: LlmApiToken::default(),
+            llm_api_token,
+            _clear_llm_token_on_sign_out: clear_llm_token_on_sign_out,
             _subscription: subscription,
         }
     }

crates/cloud_api_client/src/llm_token.rs

@@ -42,6 +42,10 @@ impl LlmApiToken {
         Self::fetch(self.0.write().await, client, system_id, organization_id).await
     }
 
+    pub async fn clear(&self) {
+        *self.0.write().await = None;
+    }
+
     /// Clears the existing token before attempting to fetch a new one.
     ///
     /// Used when switching organizations so that a failed refresh doesn't

crates/language_models/Cargo.toml

@@ -71,6 +71,7 @@ x_ai = { workspace = true, features = ["schemars"] }
 [dev-dependencies]
 client = { workspace = true, features = ["test-support"] }
 clock = { workspace = true, features = ["test-support"] }
+cloud_llm_client.workspace = true
 db = { workspace = true, features = ["test-support"] }
 feature_flags.workspace = true
 gpui = { workspace = true, features = ["test-support"] }

crates/language_models/src/provider/cloud.rs

@@ -173,7 +173,7 @@ impl State {
     }
 
     fn is_signed_out(&self, cx: &App) -> bool {
-        self.user_store.read(cx).current_user().is_none()
+        self.status.is_signed_out() || self.user_store.read(cx).current_user().is_none()
     }
 
     fn sign_in(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
@@ -229,6 +229,12 @@ impl CloudLanguageModelProvider {
                     _ = this.update(cx, |this, cx| {
                         if this.status != status {
                             this.status = status;
+                            if status.is_signed_out() {
+                                this.provider.update(cx, |provider, cx| {
+                                    provider.clear_models();
+                                    cx.notify();
+                                });
+                            }
                             cx.notify();
                         }
                     });
@@ -706,6 +712,76 @@ mod tests {
             .expect_err("provider authentication should fail when sign-in fails");
         assert!(error.to_string().contains("AuthenticationError"));
     }
+
+    #[gpui::test]
+    async fn sign_out_hides_cached_cloud_models(cx: &mut TestAppContext) {
+        let (client, _user_store, provider) = cx.update(init_test);
+        let (authenticate_tx, authenticate_rx) = futures::channel::oneshot::channel();
+        let (authenticated_user_tx, authenticated_user_rx) = futures::channel::oneshot::channel();
+        override_authenticate(&client, authenticate_rx);
+        respond_to_authenticated_user_after(&client, authenticated_user_rx);
+
+        let sign_in_task = sign_in_until_authenticating(client.clone(), cx).await;
+        authenticate_tx
+            .send(Ok(Credentials {
+                user_id: TEST_USER_ID,
+                access_token: "token".to_string(),
+            }))
+            .expect("authenticate receiver dropped");
+        authenticated_user_tx
+            .send(())
+            .expect("authenticated user receiver dropped");
+        sign_in_task.await.expect("sign-in should complete");
+        cx.executor().run_until_parked();
+
+        let model_id = cloud_llm_client::LanguageModelId(Arc::from("test-model"));
+        cx.update(|cx| {
+            let cloud_model_provider = provider.state.read(cx).provider.clone();
+            cloud_model_provider.update(cx, |cloud_model_provider, cx| {
+                cloud_model_provider.update_models(cloud_llm_client::ListModelsResponse {
+                    models: vec![cloud_llm_client::LanguageModel {
+                        provider: cloud_llm_client::LanguageModelProvider::Anthropic,
+                        id: model_id.clone(),
+                        display_name: "Test Model".to_string(),
+                        is_latest: true,
+                        max_token_count: 200_000,
+                        max_token_count_in_max_mode: None,
+                        max_output_tokens: 8_192,
+                        supports_tools: true,
+                        supports_images: false,
+                        supports_thinking: false,
+                        supports_fast_mode: false,
+                        supported_effort_levels: Vec::new(),
+                        supports_streaming_tools: false,
+                        supports_parallel_tool_calls: false,
+                    }],
+                    default_model: Some(model_id.clone()),
+                    default_fast_model: None,
+                    recommended_models: vec![model_id],
+                });
+                cx.notify();
+            });
+        });
+
+        assert!(cx.read(|cx| provider.is_authenticated(cx)));
+        assert_eq!(cx.read(|cx| provider.provided_models(cx).len()), 1);
+        assert!(cx.read(|cx| provider.default_model(cx).is_some()));
+        assert_eq!(cx.read(|cx| provider.recommended_models(cx).len()), 1);
+
+        cx.update(|cx| {
+            cx.spawn({
+                let client = client.clone();
+                async move |cx| client.sign_out(cx).await
+            })
+        })
+        .await;
+        cx.executor().run_until_parked();
+
+        assert!(!cx.read(|cx| provider.is_authenticated(cx)));
+        assert!(cx.read(|cx| provider.provided_models(cx).is_empty()));
+        assert!(cx.read(|cx| provider.default_model(cx).is_none()));
+        assert!(cx.read(|cx| provider.recommended_models(cx).is_empty()));
+    }
 }
 
 impl Component for ZedAiConfiguration {

crates/language_models_cloud/src/language_models_cloud.rs

@@ -699,6 +699,13 @@ impl<TP: CloudLlmTokenProvider + 'static> CloudModelProvider<TP> {
         self.models = models;
     }
 
+    pub fn clear_models(&mut self) {
+        self.models.clear();
+        self.default_model = None;
+        self.default_fast_model = None;
+        self.recommended_models.clear();
+    }
+
     pub fn create_model(
         &self,
         model: &Arc<cloud_llm_client::LanguageModel>,