Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Acquire LLM token from Cloud instead of Collab for Edit Predictions (#35431)

Browse source 
This PR updates the Zed Edit Prediction provider to acquire the LLM
token from Cloud instead of Collab to allow using Edit Predictions even
when disconnected from or unable to connect to the Collab server.

Release Notes:

- N/A

---------

Co-authored-by: Richard Feldman <oss@rtfeldman.com>
This commit is contained in:
Marshall Bowers 2025-07-31 18:12:04 -04:00 • committed by GitHub
parent 8e7f1899e1
commit 410348deb0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 211 additions and 125 deletions
Download patch file Download diff file Expand all files Collapse all files

2
crates/zed/src/main.rs
View file

@ -564,7 +564,7 @@ pub fn main() {
snippet_provider::init(cx);
inline_completion_registry::init(
app_state.client.clone(),
app_state.user_store.clone(),
app_state.cloud_user_store.clone(),
cx,
);
let prompt_builder = PromptBuilder::load(app_state.fs.clone(), stdout_is_a_pty(), cx);

31
crates/zed/src/zed/inline_completion_registry.rs
View file

@ -1,4 +1,4 @@
use client::{Client, DisableAiSettings, UserStore};
use client::{Client, CloudUserStore, DisableAiSettings};
use collections::HashMap;
use copilot::{Copilot, CopilotCompletionProvider};
use editor::Editor;
@ -13,12 +13,12 @@ use util::ResultExt;
use workspace::Workspace;
use zeta::{ProviderDataCollection, ZetaInlineCompletionProvider};
pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &mut App) {
let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
cx.observe_new({
let editors = editors.clone();
let client = client.clone();
let user_store = user_store.clone();
let cloud_user_store = cloud_user_store.clone();
move |editor: &mut Editor, window, cx: &mut Context<Editor>| {
if !editor.mode().is_full() {
return;
@ -48,7 +48,7 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
editor,
provider,
&client,
user_store.clone(),
cloud_user_store.clone(),
window,
cx,
);
@ -60,7 +60,7 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
let mut provider = all_language_settings(None, cx).edit_predictions.provider;
cx.spawn({
let user_store = user_store.clone();
let cloud_user_store = cloud_user_store.clone();
let editors = editors.clone();
let client = client.clone();
@ -72,7 +72,7 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
&editors,
provider,
&client,
user_store.clone(),
cloud_user_store.clone(),
cx,
);
})
@ -85,15 +85,12 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
cx.observe_global::<SettingsStore>({
let editors = editors.clone();
let client = client.clone();
let user_store = user_store.clone();
let cloud_user_store = cloud_user_store.clone();
move |cx| {
let new_provider = all_language_settings(None, cx).edit_predictions.provider;
if new_provider != provider {
let tos_accepted = user_store
.read(cx)
.current_user_has_accepted_terms()
.unwrap_or(false);
let tos_accepted = cloud_user_store.read(cx).has_accepted_tos();
telemetry::event!(
"Edit Prediction Provider Changed",
@ -107,7 +104,7 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
&editors,
provider,
&client,
user_store.clone(),
cloud_user_store.clone(),
cx,
);
@ -148,7 +145,7 @@ fn assign_edit_prediction_providers(
editors: &Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>>,
provider: EditPredictionProvider,
client: &Arc<Client>,
user_store: Entity<UserStore>,
cloud_user_store: Entity<CloudUserStore>,
cx: &mut App,
) {
for (editor, window) in editors.borrow().iter() {
@ -158,7 +155,7 @@ fn assign_edit_prediction_providers(
editor,
provider,
&client,
user_store.clone(),
cloud_user_store.clone(),
window,
cx,
);
@ -213,7 +210,7 @@ fn assign_edit_prediction_provider(
editor: &mut Editor,
provider: EditPredictionProvider,
client: &Arc<Client>,
user_store: Entity<UserStore>,
cloud_user_store: Entity<CloudUserStore>,
window: &mut Window,
cx: &mut Context<Editor>,
) {
@ -244,7 +241,7 @@ fn assign_edit_prediction_provider(
}
}
EditPredictionProvider::Zed => {
if client.status().borrow().is_connected() {
if cloud_user_store.read(cx).is_authenticated() {
let mut worktree = None;
if let Some(buffer) = &singleton_buffer {
@ -266,7 +263,7 @@ fn assign_edit_prediction_provider(
.map(|workspace| workspace.downgrade());
let zeta =
zeta::Zeta::register(workspace, worktree, client.clone(), user_store, cx);
zeta::Zeta::register(workspace, worktree, client.clone(), cloud_user_store, cx);
if let Some(buffer) = &singleton_buffer {
if buffer.read(cx).file().is_some() {