Fix RefCell panic in cloud model token counting (#54188)
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.
This commit is contained in:
1 parent
87b47a4b08
commit
76aab0c35c
2 files changed
+4
-5
No files matched your search
@@ -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()
|
||||
|
||||
@@ -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?;
|
||||
|
||||
|
||||
Reference in new issue
Block a user