From d9d8b6057c4a21a98c990ed58b1fd52aee0ad76f Mon Sep 17 00:00:00 2001 From: Oleksiy Syvokon Date: Tue, 26 May 2026 14:48:15 +0300 Subject: [PATCH 1/5] Clear LLM token and model list on sign out --- Cargo.lock | 1 + crates/client/src/client.rs | 88 +++++++++++++++++++ crates/client/src/llm_token.rs | 5 ++ crates/cloud_api_client/src/llm_token.rs | 4 + crates/language_models/Cargo.toml | 1 + crates/language_models/src/provider/cloud.rs | 78 +++++++++++++++- .../src/language_models_cloud.rs | 7 ++ 7 files changed, 183 insertions(+), 1 deletion(-) 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/client.rs b/crates/client/src/client.rs index 093eb894483a75..b1687786fcef5f 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -1582,11 +1582,24 @@ impl Client { }) } + async fn ensure_authenticated_for_llm_token(&self, llm_token: &LlmApiToken) -> Result<()> { + let is_signed_out = + self.state.read().credentials.is_none() || self.status().borrow().is_signed_out(); + if is_signed_out { + llm_token.clear().await; + anyhow::bail!("not signed in"); + } + + Ok(()) + } + pub async fn cached_llm_token( &self, llm_token: &LlmApiToken, organization_id: Option, ) -> Result { + self.ensure_authenticated_for_llm_token(llm_token).await?; + let system_id = self.telemetry().system_id().map(|x| x.to_string()); let cloud_client = self.cloud_client(); match llm_token @@ -1632,6 +1645,8 @@ impl Client { llm_token: &LlmApiToken, organization_id: Option, ) -> Result { + self.ensure_authenticated_for_llm_token(llm_token).await?; + let system_id = self.telemetry().system_id().map(|x| x.to_string()); let cloud_client = self.cloud_client(); match llm_token @@ -1670,6 +1685,9 @@ impl Client { pub async fn sign_out(self: &Arc, cx: &AsyncApp) { self.state.write().credentials = None; self.cloud_client.clear_credentials(); + if let Some(llm_token) = cx.update(|cx| try_global_llm_token(cx)) { + llm_token.clear().await; + } self.disconnect(cx); if self.has_credentials(cx).await { @@ -2221,6 +2239,76 @@ mod tests { assert_eq!(credentials.access_token, "2"); } + #[gpui::test] + async fn test_cached_llm_token_is_cleared_after_sign_out(cx: &mut TestAppContext) { + init_test(cx); + let llm_token_request_count = Arc::new(Mutex::new(0)); + let http_client = FakeHttpClient::create({ + let llm_token_request_count = llm_token_request_count.clone(); + move |request| { + let llm_token_request_count = llm_token_request_count.clone(); + async move { + assert_eq!(request.uri().path(), "/client/llm_tokens"); + let token = { + let mut request_count = llm_token_request_count.lock(); + *request_count += 1; + format!("llm-token-{}", *request_count) + }; + Ok(http_client::Response::builder() + .status(200) + .body( + serde_json::to_string(&cloud_api_client::CreateLlmTokenResponse { + token: cloud_api_client::LlmToken(token), + }) + .expect("failed to serialize LLM token response") + .into(), + ) + .expect("failed to build LLM token response")) + } + } + }); + let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx)); + let llm_token = LlmApiToken::default(); + + let authenticate_client = |client: &Arc, cx: &TestAppContext| { + client.state.write().credentials = Some(Credentials { + user_id: 1, + access_token: "account-token".into(), + }); + client + .cloud_client + .set_credentials(1, "account-token".into()); + let cx = cx.to_async(); + client.set_status(Status::Authenticated, &cx); + }; + + authenticate_client(&client, cx); + assert_eq!( + client + .cached_llm_token(&llm_token, None) + .await + .expect("initial LLM token request should succeed"), + "llm-token-1" + ); + assert_eq!(*llm_token_request_count.lock(), 1); + + client.sign_out(&cx.to_async()).await; + client + .cached_llm_token(&llm_token, None) + .await + .expect_err("signed-out clients should not reuse cached LLM tokens"); + + authenticate_client(&client, cx); + assert_eq!( + client + .cached_llm_token(&llm_token, None) + .await + .expect("LLM token should be fetched again after reauthentication"), + "llm-token-2" + ); + assert_eq!(*llm_token_request_count.lock(), 2); + } + #[gpui::test(iterations = 10)] async fn test_authenticating_more_than_once( cx: &mut TestAppContext, diff --git a/crates/client/src/llm_token.rs b/crates/client/src/llm_token.rs index 058be7905fa12d..ad312682f84165 100644 --- a/crates/client/src/llm_token.rs +++ b/crates/client/src/llm_token.rs @@ -32,6 +32,11 @@ pub fn global_llm_token(cx: &App) -> LlmApiToken { .clone() } +pub fn try_global_llm_token(cx: &App) -> Option { + cx.try_global::() + .map(|listener| listener.0.read(cx).llm_api_token.clone()) +} + struct GlobalRefreshLlmTokenListener(Entity); impl Global for GlobalRefreshLlmTokenListener {} 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, From 5cf3327a6bbcafdd39fc508b26bc84d245cb5d07 Mon Sep 17 00:00:00 2001 From: Oleksiy Syvokon Date: Tue, 26 May 2026 14:54:33 +0300 Subject: [PATCH 2/5] Refactor --- crates/client/src/client.rs | 3 --- crates/client/src/llm_token.rs | 25 ++++++++++++++++++------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs index b1687786fcef5f..a0a0783d0802be 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -1685,9 +1685,6 @@ impl Client { pub async fn sign_out(self: &Arc, cx: &AsyncApp) { self.state.write().credentials = None; self.cloud_client.clear_credentials(); - if let Some(llm_token) = cx.update(|cx| try_global_llm_token(cx)) { - llm_token.clear().await; - } self.disconnect(cx); if self.has_credentials(cx).await { diff --git a/crates/client/src/llm_token.rs b/crates/client/src/llm_token.rs index ad312682f84165..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; @@ -32,11 +33,6 @@ pub fn global_llm_token(cx: &App) -> LlmApiToken { .clone() } -pub fn try_global_llm_token(cx: &App) -> Option { - cx.try_global::() - .map(|listener| listener.0.read(cx).llm_api_token.clone()) -} - struct GlobalRefreshLlmTokenListener(Entity); impl Global for GlobalRefreshLlmTokenListener {} @@ -47,6 +43,7 @@ pub struct RefreshLlmTokenListener { client: Arc, user_store: Entity, llm_api_token: LlmApiToken, + _clear_llm_token_on_sign_out: Task<()>, _subscription: Subscription, } @@ -78,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, } } From 4a371ab418592de1d2434b5ac01874d971c3ee2a Mon Sep 17 00:00:00 2001 From: Oleksiy Syvokon Date: Tue, 26 May 2026 15:07:52 +0300 Subject: [PATCH 3/5] Refactor test --- crates/client/src/client.rs | 43 ++++++++++++++++++------------------- 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs index a0a0783d0802be..f41e37a9204038 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -2239,18 +2239,15 @@ mod tests { #[gpui::test] async fn test_cached_llm_token_is_cleared_after_sign_out(cx: &mut TestAppContext) { init_test(cx); - let llm_token_request_count = Arc::new(Mutex::new(0)); + let llm_token_request_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let http_client = FakeHttpClient::create({ let llm_token_request_count = llm_token_request_count.clone(); move |request| { let llm_token_request_count = llm_token_request_count.clone(); async move { assert_eq!(request.uri().path(), "/client/llm_tokens"); - let token = { - let mut request_count = llm_token_request_count.lock(); - *request_count += 1; - format!("llm-token-{}", *request_count) - }; + let request_count = llm_token_request_count.fetch_add(1, Ordering::SeqCst) + 1; + let token = format!("llm-token-{request_count}"); Ok(http_client::Response::builder() .status(200) .body( @@ -2265,21 +2262,20 @@ mod tests { } }); let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx)); + client.override_authenticate(|cx| { + cx.background_spawn(async { + Ok(Credentials { + user_id: 1, + access_token: "account-token".into(), + }) + }) + }); let llm_token = LlmApiToken::default(); - let authenticate_client = |client: &Arc, cx: &TestAppContext| { - client.state.write().credentials = Some(Credentials { - user_id: 1, - access_token: "account-token".into(), - }); - client - .cloud_client - .set_credentials(1, "account-token".into()); - let cx = cx.to_async(); - client.set_status(Status::Authenticated, &cx); - }; - - authenticate_client(&client, cx); + client + .sign_in(false, &cx.to_async()) + .await + .expect("initial sign-in should succeed"); assert_eq!( client .cached_llm_token(&llm_token, None) @@ -2287,7 +2283,7 @@ mod tests { .expect("initial LLM token request should succeed"), "llm-token-1" ); - assert_eq!(*llm_token_request_count.lock(), 1); + assert_eq!(llm_token_request_count.load(Ordering::SeqCst), 1); client.sign_out(&cx.to_async()).await; client @@ -2295,7 +2291,10 @@ mod tests { .await .expect_err("signed-out clients should not reuse cached LLM tokens"); - authenticate_client(&client, cx); + client + .sign_in(false, &cx.to_async()) + .await + .expect("reauthentication should succeed"); assert_eq!( client .cached_llm_token(&llm_token, None) @@ -2303,7 +2302,7 @@ mod tests { .expect("LLM token should be fetched again after reauthentication"), "llm-token-2" ); - assert_eq!(*llm_token_request_count.lock(), 2); + assert_eq!(llm_token_request_count.load(Ordering::SeqCst), 2); } #[gpui::test(iterations = 10)] From 0826c62f2dd34d015194ec8a4d4413bfbe039bff Mon Sep 17 00:00:00 2001 From: Oleksiy Syvokon Date: Wed, 27 May 2026 01:56:14 +0300 Subject: [PATCH 4/5] Remove explicit auth check --- crates/client/src/client.rs | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs index f41e37a9204038..b846bb4a51afa7 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -1582,24 +1582,11 @@ impl Client { }) } - async fn ensure_authenticated_for_llm_token(&self, llm_token: &LlmApiToken) -> Result<()> { - let is_signed_out = - self.state.read().credentials.is_none() || self.status().borrow().is_signed_out(); - if is_signed_out { - llm_token.clear().await; - anyhow::bail!("not signed in"); - } - - Ok(()) - } - pub async fn cached_llm_token( &self, llm_token: &LlmApiToken, organization_id: Option, ) -> Result { - self.ensure_authenticated_for_llm_token(llm_token).await?; - let system_id = self.telemetry().system_id().map(|x| x.to_string()); let cloud_client = self.cloud_client(); match llm_token @@ -1645,8 +1632,6 @@ impl Client { llm_token: &LlmApiToken, organization_id: Option, ) -> Result { - self.ensure_authenticated_for_llm_token(llm_token).await?; - let system_id = self.telemetry().system_id().map(|x| x.to_string()); let cloud_client = self.cloud_client(); match llm_token From 43a62aabdca43947028e797255a41a3ddbd235c7 Mon Sep 17 00:00:00 2001 From: Oleksiy Syvokon Date: Wed, 27 May 2026 02:08:30 +0300 Subject: [PATCH 5/5] Remove a test Without an explicit auth check, the test becomes too complex --- crates/client/src/client.rs | 69 ------------------------------------- 1 file changed, 69 deletions(-) diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs index b846bb4a51afa7..093eb894483a75 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -2221,75 +2221,6 @@ mod tests { assert_eq!(credentials.access_token, "2"); } - #[gpui::test] - async fn test_cached_llm_token_is_cleared_after_sign_out(cx: &mut TestAppContext) { - init_test(cx); - let llm_token_request_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let http_client = FakeHttpClient::create({ - let llm_token_request_count = llm_token_request_count.clone(); - move |request| { - let llm_token_request_count = llm_token_request_count.clone(); - async move { - assert_eq!(request.uri().path(), "/client/llm_tokens"); - let request_count = llm_token_request_count.fetch_add(1, Ordering::SeqCst) + 1; - let token = format!("llm-token-{request_count}"); - Ok(http_client::Response::builder() - .status(200) - .body( - serde_json::to_string(&cloud_api_client::CreateLlmTokenResponse { - token: cloud_api_client::LlmToken(token), - }) - .expect("failed to serialize LLM token response") - .into(), - ) - .expect("failed to build LLM token response")) - } - } - }); - let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx)); - client.override_authenticate(|cx| { - cx.background_spawn(async { - Ok(Credentials { - user_id: 1, - access_token: "account-token".into(), - }) - }) - }); - let llm_token = LlmApiToken::default(); - - client - .sign_in(false, &cx.to_async()) - .await - .expect("initial sign-in should succeed"); - assert_eq!( - client - .cached_llm_token(&llm_token, None) - .await - .expect("initial LLM token request should succeed"), - "llm-token-1" - ); - assert_eq!(llm_token_request_count.load(Ordering::SeqCst), 1); - - client.sign_out(&cx.to_async()).await; - client - .cached_llm_token(&llm_token, None) - .await - .expect_err("signed-out clients should not reuse cached LLM tokens"); - - client - .sign_in(false, &cx.to_async()) - .await - .expect("reauthentication should succeed"); - assert_eq!( - client - .cached_llm_token(&llm_token, None) - .await - .expect("LLM token should be fetched again after reauthentication"), - "llm-token-2" - ); - assert_eq!(llm_token_request_count.load(Ordering::SeqCst), 2); - } - #[gpui::test(iterations = 10)] async fn test_authenticating_more_than_once( cx: &mut TestAppContext,