Add zeta2 to registry

Cargo.lock

@@ -20643,6 +20643,7 @@ dependencies = [
  "zed_actions",
  "zed_env_vars",
  "zeta",
+ "zeta2",
  "zlog",
  "zlog_settings",
 ]
@@ -20919,6 +20920,19 @@ dependencies = [
  "zlog",
 ]
 
+[[package]]
+name = "zeta2"
+version = "0.1.0"
+dependencies = [
+ "client",
+ "edit_prediction",
+ "gpui",
+ "language",
+ "project",
+ "util",
+ "workspace-hack",
+]
+
 [[package]]
 name = "zeta_cli"
 version = "0.1.0"

Cargo.toml

@@ -197,6 +197,7 @@ members = [
     "crates/zed_actions",
     "crates/zed_env_vars",
     "crates/zeta",
+    "crates/zeta2",
     "crates/zeta_cli",
     "crates/zlog",
     "crates/zlog_settings",
@@ -427,6 +428,7 @@ zed = { path = "crates/zed" }
 zed_actions = { path = "crates/zed_actions" }
 zed_env_vars = { path = "crates/zed_env_vars" }
 zeta = { path = "crates/zeta" }
+zeta2 = { path = "crates/zeta2" }
 zlog = { path = "crates/zlog" }
 zlog_settings = { path = "crates/zlog_settings" }
 

crates/edit_prediction_button/src/edit_prediction_button.rs

@@ -233,7 +233,7 @@ impl Render for EditPredictionButton {
                 )
             }
 
-            EditPredictionProvider::Zed => {
+            EditPredictionProvider::Zed | EditPredictionProvider::Zed2 => {
                 let enabled = self.editor_enabled.unwrap_or(true);
 
                 let zeta_icon = if enabled {

crates/language/src/language_settings.rs

@@ -220,12 +220,13 @@ pub enum EditPredictionProvider {
     Copilot,
     Supermaven,
     Zed,
+    Zed2,
 }
 
 impl EditPredictionProvider {
     pub fn is_zed(&self) -> bool {
         match self {
-            EditPredictionProvider::Zed => true,
+            EditPredictionProvider::Zed | EditPredictionProvider::Zed2 => true,
             EditPredictionProvider::None
             | EditPredictionProvider::Copilot
             | EditPredictionProvider::Supermaven => false,

crates/zed/Cargo.toml

@@ -165,6 +165,7 @@ workspace.workspace = true
 zed_actions.workspace = true
 zed_env_vars.workspace = true
 zeta.workspace = true
+zeta2.workspace = true
 zlog.workspace = true
 zlog_settings.workspace = true
 

crates/zed/src/zed/edit_prediction_registry.rs

@@ -220,5 +220,10 @@ fn assign_edit_prediction_provider(
                 editor.set_edit_prediction_provider(Some(provider), window, cx);
             }
         }
+        EditPredictionProvider::Zed2 => {
+            let provider = cx.new(|_| zeta2::Zeta2EditPredictionProvider::new());
+
+            editor.set_edit_prediction_provider(Some(provider), window, cx);
+        }
     }
 }

crates/zeta2/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "zeta2"
+version = "0.1.0"
+edition.workspace = true
+publish.workspace = true
+license = "GPL-3.0-or-later"
+
+[lints]
+workspace = true
+
+[lib]
+path = "src/zeta2.rs"
+
+[dependencies]
+client.workspace = true
+edit_prediction.workspace = true
+gpui.workspace = true
+language.workspace = true
+project.workspace = true
+workspace-hack.workspace = true
+util.workspace = true

crates/zeta2/LICENSE-GPL

@@ -0,0 +1 @@
+../../LICENSE-GPL

crates/zeta2/src/zeta2.rs

@@ -0,0 +1,154 @@
+use std::{ops::Range, sync::Arc};
+
+use gpui::{App, Entity, EntityId, Task, prelude::*};
+
+use edit_prediction::{DataCollectionState, Direction, EditPrediction, EditPredictionProvider};
+use language::{Anchor, ToPoint};
+
+pub struct Zeta2EditPredictionProvider {
+    current: Option<CurrentEditPrediction>,
+    pending: Option<Task<()>>,
+}
+
+impl Zeta2EditPredictionProvider {
+    pub fn new() -> Self {
+        Self {
+            current: None,
+            pending: None,
+        }
+    }
+}
+
+#[derive(Clone)]
+struct CurrentEditPrediction {
+    buffer_id: EntityId,
+    prediction: EditPrediction,
+}
+
+impl EditPredictionProvider for Zeta2EditPredictionProvider {
+    fn name() -> &'static str {
+        // TODO [zeta2]
+        "zed-predict2"
+    }
+
+    fn display_name() -> &'static str {
+        "Zed's Edit Predictions 2"
+    }
+
+    fn show_completions_in_menu() -> bool {
+        true
+    }
+
+    fn show_tab_accept_marker() -> bool {
+        true
+    }
+
+    fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
+        // TODO [zeta2]
+        DataCollectionState::Unsupported
+    }
+
+    fn toggle_data_collection(&mut self, _cx: &mut App) {
+        // TODO [zeta2]
+    }
+
+    fn usage(&self, _cx: &App) -> Option<client::EditPredictionUsage> {
+        // TODO [zeta2]
+        None
+    }
+
+    fn is_enabled(
+        &self,
+        _buffer: &Entity<language::Buffer>,
+        _cursor_position: language::Anchor,
+        _cx: &App,
+    ) -> bool {
+        true
+    }
+
+    fn is_refreshing(&self) -> bool {
+        self.pending.is_some()
+    }
+
+    fn refresh(
+        &mut self,
+        _project: Option<Entity<project::Project>>,
+        buffer: Entity<language::Buffer>,
+        cursor_position: language::Anchor,
+        _debounce: bool,
+        cx: &mut Context<Self>,
+    ) {
+        // TODO [zeta2] check account
+        // TODO [zeta2] actually request completion / interpolate
+
+        let snapshot = buffer.read(cx).snapshot();
+        let point = cursor_position.to_point(&snapshot);
+        let end_anchor = snapshot.anchor_before(language::Point::new(
+            point.row,
+            snapshot.line_len(point.row),
+        ));
+
+        let edits: Arc<[(Range<Anchor>, String)]> =
+            vec![(cursor_position..end_anchor, "👻".to_string())].into();
+        let edits_preview_task = buffer.read(cx).preview_edits(edits.clone(), cx);
+
+        // TODO [zeta2] throttle
+        // TODO [zeta2] keep 2 requests
+        self.pending = Some(cx.spawn(async move |this, cx| {
+            let edits_preview = edits_preview_task.await;
+
+            this.update(cx, |this, cx| {
+                this.current = Some(CurrentEditPrediction {
+                    buffer_id: buffer.entity_id(),
+                    prediction: EditPrediction {
+                        // TODO! [zeta2] request id?
+                        id: None,
+                        edits: edits.to_vec(),
+                        edit_preview: Some(edits_preview),
+                    },
+                });
+                this.pending.take();
+                cx.notify();
+            })
+            .ok();
+        }));
+        cx.notify();
+    }
+
+    fn cycle(
+        &mut self,
+        _buffer: Entity<language::Buffer>,
+        _cursor_position: language::Anchor,
+        _direction: Direction,
+        _cx: &mut Context<Self>,
+    ) {
+    }
+
+    fn accept(&mut self, _cx: &mut Context<Self>) {
+        // TODO [zeta2] report accept
+        self.current.take();
+        self.pending.take();
+    }
+
+    fn discard(&mut self, _cx: &mut Context<Self>) {
+        self.current.take();
+        self.pending.take();
+    }
+
+    fn suggest(
+        &mut self,
+        buffer: &Entity<language::Buffer>,
+        _cursor_position: language::Anchor,
+        _cx: &mut Context<Self>,
+    ) -> Option<EditPrediction> {
+        let current_prediction = self.current.take()?;
+
+        if current_prediction.buffer_id != buffer.entity_id() {
+            return None;
+        }
+
+        // TODO [zeta2] interpolate
+
+        Some(current_prediction.prediction)
+    }
+}