Fix RefCell panic in cloud model token counting (#54188) (cherry-pick to preview) (#54191)

zed-zippy[bot] and Richard Feldman created 4 days ago

Cherry-pick of #54188 to preview

----
Fixes #54140

When `RulesLibrary::count_tokens` calls
`CloudLanguageModel::count_tokens` for Google cloud models, it does so
inside a `cx.update` closure, which holds a mutable borrow on the global
`AppCell`. The Google provider branch then called
`token_provider.auth_context(&cx.to_async())`, which created a new
`AsyncApp` handle and tried to take a shared borrow on the same
`RefCell` — causing a "RefCell already mutably borrowed" panic.

This only affects Google models because they are the only provider that
counts tokens server-side via an HTTP request (requiring
authentication). The other providers (Anthropic, OpenAI, xAI) count
tokens locally using tiktoken, so they never call `auth_context` during
`count_tokens`.

The fix makes `CloudLlmTokenProvider::auth_context` generic over `impl
AppContext` instead of requiring `&AsyncApp`. This allows the
`count_tokens` call site to pass `&App` directly (which reads entities
without re-borrowing the `RefCell`), while all other call sites that
already pass `&AsyncApp` (e.g. `stream_completion`, `refresh_models`)
continue to work unchanged.

Release Notes:

- Fixed a crash ("RefCell already mutably borrowed") that could occur
when counting tokens with Google cloud language models.

Co-authored-by: Richard Feldman <[e-mail]>

Change summary

crates/language_models/src/provider/cloud.rs              | 5 ++---
crates/language_models_cloud/src/language_models_cloud.rs | 4 ++--
2 files changed, 4 insertions(+), 5 deletions(-)

Detailed changes

crates/language_models/src/provider/cloud.rs

@@ -6,8 +6,7 @@ use cloud_api_types::OrganizationId;
 use cloud_api_types::Plan;
 use futures::StreamExt;
 use futures::future::BoxFuture;
-use gpui::AsyncApp;
-use gpui::{AnyElement, AnyView, App, Context, Entity, Subscription, Task};
+use gpui::{AnyElement, AnyView, App, AppContext, Context, Entity, Subscription, Task};
 use language_model::{
     AuthenticateError, IconOrSvg, LanguageModel, LanguageModelProvider, LanguageModelProviderId,
     LanguageModelProviderName, LanguageModelProviderState, ZED_CLOUD_PROVIDER_ID,
@@ -34,7 +33,7 @@ struct ClientTokenProvider {
 impl CloudLlmTokenProvider for ClientTokenProvider {
     type AuthContext = Option<OrganizationId>;
 
-    fn auth_context(&self, cx: &AsyncApp) -> Self::AuthContext {
+    fn auth_context(&self, cx: &impl AppContext) -> Self::AuthContext {
         self.user_store.read_with(cx, |user_store, _| {
             user_store
                 .current_organization()

crates/language_models_cloud/src/language_models_cloud.rs

@@ -57,7 +57,7 @@ const PROVIDER_NAME: LanguageModelProviderName = ZED_CLOUD_PROVIDER_NAME;
 pub trait CloudLlmTokenProvider: Send + Sync {
     type AuthContext: Clone + Send + 'static;
 
-    fn auth_context(&self, cx: &AsyncApp) -> Self::AuthContext;
+    fn auth_context(&self, cx: &impl AppContext) -> Self::AuthContext;
     fn acquire_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>;
     fn refresh_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>;
 }
@@ -405,7 +405,7 @@ impl<TP: CloudLlmTokenProvider + 'static> LanguageModel for CloudLanguageModel<T
                 let model_id = self.model.id.to_string();
                 let generate_content_request =
                     into_google(request, model_id.clone(), GoogleModelMode::Default);
-                let auth_context = token_provider.auth_context(&cx.to_async());
+                let auth_context = token_provider.auth_context(cx);
                 async move {
                     let token = token_provider.acquire_token(auth_context).await?;