diff --git a/Cargo.lock b/Cargo.lock index 225d5ca6dbb724..e2aea8f7dc76a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9678,6 +9678,7 @@ dependencies = [ "clock", "cloud_api_client", "cloud_api_types", + "cloud_llm_client", "collections", "component", "convert_case 0.8.0", diff --git a/crates/client/src/llm_token.rs b/crates/client/src/llm_token.rs index 058be7905fa12d..e9ac56decf4a3a 100644 --- a/crates/client/src/llm_token.rs +++ b/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, user_store: Entity, 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, } } diff --git a/crates/cloud_api_client/src/llm_token.rs b/crates/cloud_api_client/src/llm_token.rs index 7baafc545b5c48..3c0983b9f051e0 100644 --- a/crates/cloud_api_client/src/llm_token.rs +++ b/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 diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 7e2d2618ea9695..b96a377f02c19b 100644 --- a/crates/language_models/Cargo.toml +++ b/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"] } diff --git a/crates/language_models/src/provider/cloud.rs b/crates/language_models/src/provider/cloud.rs index 3cb9fcf1b86975..a1c05057c23a9b 100644 --- a/crates/language_models/src/provider/cloud.rs +++ b/crates/language_models/src/provider/cloud.rs @@ -145,7 +145,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) -> Task> { @@ -183,6 +183,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(); } }); @@ -650,6 +656,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 { diff --git a/crates/language_models_cloud/src/language_models_cloud.rs b/crates/language_models_cloud/src/language_models_cloud.rs index 57ad96280b51c3..940eb51dc37d10 100644 --- a/crates/language_models_cloud/src/language_models_cloud.rs +++ b/crates/language_models_cloud/src/language_models_cloud.rs @@ -699,6 +699,13 @@ impl CloudModelProvider { 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,