Progress towards scoping zeta edit history to projects

Also:

* Removes old PredictEditsBody fields that don't have anticipated future use

* Sorts diagnostics by proximity to cursor and truncates based on json byte count

* Brings back cursor_offset

Cargo.lock

@@ -20793,6 +20793,7 @@ dependencies = [
  "language_model",
  "log",
  "menu",
+ "multi_buffer",
  "postage",
  "project",
  "rand 0.8.5",

crates/agent_ui/src/acp/message_editor.rs

@@ -438,13 +438,16 @@ impl MessageEditor {
         let Some(entry) = self.project.read(cx).entry_for_path(&project_path, cx) else {
             return Task::ready(Err(anyhow!("project entry not found")));
         };
-        let Some(worktree) = self.project.read(cx).worktree_for_entry(entry.id, cx) else {
+        let directory_path = entry.path.clone();
+        let Some(worktree) = self
+            .project
+            .read(cx)
+            .worktree_for_id(project_path.worktree_id, cx)
+        else {
             return Task::ready(Err(anyhow!("worktree not found")));
         };
         let project = self.project.clone();
         cx.spawn(async move |_, cx| {
-            let directory_path = entry.path.clone();
-
             let worktree_id = worktree.read_with(cx, |worktree, _| worktree.id())?;
             let file_paths = worktree.read_with(cx, |worktree, _cx| {
                 collect_files_in_path(worktree, &directory_path)

crates/agent_ui/src/acp/thread_view.rs

@@ -3948,15 +3948,15 @@ impl AcpThreadView {
                 MentionUri::PastedImage => {}
                 MentionUri::Directory { abs_path } => {
                     let project = workspace.project();
-                    let Some(entry) = project.update(cx, |project, cx| {
+                    let Some(entry_id) = project.update(cx, |project, cx| {
                         let path = project.find_project_path(abs_path, cx)?;
-                        project.entry_for_path(&path, cx)
+                        project.entry_for_path(&path, cx).map(|entry| entry.id)
                     }) else {
                         return;
                     };
 
                     project.update(cx, |_, cx| {
-                        cx.emit(project::Event::RevealInProjectPanel(entry.id));
+                        cx.emit(project::Event::RevealInProjectPanel(entry_id));
                     });
                 }
                 MentionUri::Symbol {
@@ -3969,11 +3969,9 @@ impl AcpThreadView {
                     line_range,
                 } => {
                     let project = workspace.project();
-                    let Some((path, _)) = project.update(cx, |project, cx| {
-                        let path = project.find_project_path(path, cx)?;
-                        let entry = project.entry_for_path(&path, cx)?;
-                        Some((path, entry))
-                    }) else {
+                    let Some(path) =
+                        project.update(cx, |project, cx| project.find_project_path(path, cx))
+                    else {
                         return;
                     };
 

crates/agent_ui/src/context_picker.rs

@@ -987,7 +987,8 @@ impl MentionLink {
                     .read(cx)
                     .project()
                     .read(cx)
-                    .entry_for_path(&project_path, cx)?;
+                    .entry_for_path(&project_path, cx)?
+                    .clone();
                 Some(MentionLink::File(project_path, entry))
             }
             Self::SYMBOL => {

crates/cloud_llm_client/src/cloud_llm_client.rs

@@ -138,22 +138,42 @@ pub enum LanguageModelProvider {
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PredictEditsBody {
-    #[serde(skip_serializing_if = "Option::is_none", default)]
-    pub outline: Option<String>,
     pub input_events: String,
     pub input_excerpt: String,
-    #[serde(skip_serializing_if = "Option::is_none", default)]
-    pub speculated_output: Option<String>,
     /// Whether the user provided consent for sampling this interaction.
     #[serde(default, alias = "data_collection_permission")]
     pub can_collect_data: bool,
+    /// Note that this is no longer sent, in favor of `PredictEditsAdditionalContext`.
     #[serde(skip_serializing_if = "Option::is_none", default)]
     pub diagnostic_groups: Option<Vec<(String, serde_json::Value)>>,
-    /// Info about the git repository state, only present when can_collect_data is true.
+    /// Info about the git repository state, only present when can_collect_data is true. Note that
+    /// this is no longer sent, in favor of `PredictEditsAdditionalContext`.
     #[serde(skip_serializing_if = "Option::is_none", default)]
     pub git_info: Option<PredictEditsGitInfo>,
 }
 
+/// Additional context only stored when can_collect_data is true for the corresponding edit
+/// predictions request.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PredictEditsAdditionalContext {
+    /// Path to the file in the repository that contains the input excerpt.
+    pub input_path: String,
+    /// Cursor position within the file that contains the input excerpt.
+    pub cursor_point: Point,
+    /// Cursor offset in bytes within the file that contains the input excerpt.
+    pub cursor_offset: usize,
+    #[serde(flatten)]
+    pub git_info: PredictEditsGitInfo,
+    /// Diagnostic near the cursor position.
+    #[serde(skip_serializing_if = "Vec::is_empty", default)]
+    pub diagnostic_groups: Vec<(String, Box<serde_json::value::RawValue>)>,
+    /// True if the diagnostics were truncated.
+    pub diagnostic_groups_truncated: bool,
+    /// Recently active files that may be within this repository.
+    #[serde(skip_serializing_if = "Vec::is_empty", default)]
+    pub recent_files: Vec<PredictEditsRecentFile>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PredictEditsGitInfo {
     /// SHA of git HEAD commit at time of prediction.
@@ -167,6 +187,32 @@ pub struct PredictEditsGitInfo {
     pub remote_upstream_url: Option<String>,
 }
 
+/// A zero-indexed point in a text buffer consisting of a row and column.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Point {
+    pub row: u32,
+    pub column: u32,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PredictEditsRecentFile {
+    /// Path to a file within the repository.
+    pub path: String,
+    /// Most recent cursor position with the file.
+    pub cursor_point: Point,
+    /// Milliseconds between the editor for this file being active and the request time.
+    pub active_to_now_ms: u32,
+    /// Number of times the editor for this file was activated.
+    pub activation_count: u32,
+    /// Rough estimate of milliseconds the user was editing the file.
+    pub cumulative_time_editing_ms: u32,
+    /// Rough estimate of milliseconds the user was navigating within the file.
+    pub cumulative_time_navigating_ms: u32,
+    /// Whether the file is a multibuffer.
+    #[serde(skip_serializing_if = "is_default", default)]
+    pub is_multibuffer: bool,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PredictEditsResponse {
     pub request_id: Uuid,
@@ -336,6 +382,10 @@ pub struct UsageData {
     pub limit: UsageLimit,
 }
 
+fn is_default<T: Default + PartialEq>(value: &T) -> bool {
+    *value == T::default()
+}
+
 #[cfg(test)]
 mod tests {
     use pretty_assertions::assert_eq;

crates/collab/src/tests/integration_tests.rs

@@ -3072,12 +3072,7 @@ async fn test_git_status_sync(
             .collect::<Vec<_>>();
         assert_eq!(repos.len(), 1);
         let repo = repos.into_iter().next().unwrap();
-        assert_eq!(
-            repo.read(cx)
-                .status_for_path(&file.into())
-                .map(|entry| entry.status),
-            status
-        );
+        assert_eq!(repo.read(cx).status_for_path(&file.into()), status);
     }
 
     project_local.read_with(cx_a, |project, cx| {

crates/editor/src/items.rs

@@ -685,7 +685,7 @@ impl Item for Editor {
                         .git_store()
                         .read(cx)
                         .repository_and_path_for_buffer_id(buffer_id, cx)?;
-                    let status = repo.read(cx).status_for_path(&repo_path)?.status;
+                    let status = repo.read(cx).status_for_path(&repo_path)?;
 
                     Some(entry_git_aware_label_color(
                         status.summary(),

crates/language/src/buffer.rs

@@ -146,7 +146,7 @@ pub struct BufferSnapshot {
     pub text: text::BufferSnapshot,
     pub(crate) syntax: SyntaxSnapshot,
     file: Option<Arc<dyn File>>,
-    diagnostics: SmallVec<[(LanguageServerId, DiagnosticSet); 2]>,
+    pub diagnostics: SmallVec<[(LanguageServerId, DiagnosticSet); 2]>,
     remote_selections: TreeMap<ReplicaId, SelectionSet>,
     language: Option<Arc<Language>>,
     non_text_state_update_count: usize,

crates/outline_panel/src/outline_panel.rs

@@ -2693,9 +2693,7 @@ impl OutlinePanel {
                         let status = git_store
                             .read(cx)
                             .repository_and_path_for_buffer_id(buffer_id, cx)
-                            .and_then(|(repo, path)| {
-                                Some(repo.read(cx).status_for_path(&path)?.status)
-                            });
+                            .and_then(|(repo, path)| repo.read(cx).status_for_path(&path));
                         buffer_excerpts
                             .entry(buffer_id)
                             .or_insert_with(|| {

crates/project/src/git_store.rs

@@ -782,7 +782,7 @@ impl GitStore {
         cx: &App,
     ) -> Option<FileStatus> {
         let (repo, repo_path) = self.repository_and_path_for_project_path(project_path, cx)?;
-        Some(repo.read(cx).status_for_path(&repo_path)?.status)
+        repo.read(cx).status_for_path(&repo_path)
     }
 
     pub fn checkpoint(&self, cx: &mut App) -> Task<Result<GitStoreCheckpoint>> {
@@ -1360,8 +1360,7 @@ impl GitStore {
 
     pub fn status_for_buffer_id(&self, buffer_id: BufferId, cx: &App) -> Option<FileStatus> {
         let (repo, path) = self.repository_and_path_for_buffer_id(buffer_id, cx)?;
-        let status = repo.read(cx).snapshot.status_for_path(&path)?;
-        Some(status.status)
+        repo.read(cx).snapshot.status_for_path(&path)
     }
 
     pub fn repository_and_path_for_buffer_id(
@@ -2810,10 +2809,10 @@ impl RepositorySnapshot {
         self.statuses_by_path.summary().item_summary
     }
 
-    pub fn status_for_path(&self, path: &RepoPath) -> Option<StatusEntry> {
+    pub fn status_for_path(&self, path: &RepoPath) -> Option<FileStatus> {
         self.statuses_by_path
             .get(&PathKey(path.0.clone()), &())
-            .cloned()
+            .map(|entry| entry.status)
     }
 
     pub fn abs_path_to_repo_path(&self, abs_path: &Path) -> Option<RepoPath> {
@@ -2840,7 +2839,7 @@ impl RepositorySnapshot {
             self.merge.conflicted_paths.contains(repo_path);
         let has_conflict_currently = self
             .status_for_path(repo_path)
-            .is_some_and(|entry| entry.status.is_conflicted());
+            .is_some_and(|status| status.is_conflicted());
         had_conflict_on_last_merge_head_change || has_conflict_currently
     }
 

crates/project/src/project.rs

@@ -4314,7 +4314,7 @@ impl Project {
         self.active_entry
     }
 
-    pub fn entry_for_path(&self, path: &ProjectPath, cx: &App) -> Option<Entry> {
+    pub fn entry_for_path<'a>(&'a self, path: &ProjectPath, cx: &'a App) -> Option<&'a Entry> {
         self.worktree_store.read(cx).entry_for_path(path, cx)
     }
 

crates/project/src/project_tests.rs

@@ -8246,7 +8246,7 @@ async fn test_repository_subfolder_git_status(
 
         assert_eq!(repository.status_for_path(&C_TXT.into()), None);
         assert_eq!(
-            repository.status_for_path(&E_TXT.into()).unwrap().status,
+            repository.status_for_path(&E_TXT.into()).unwrap(),
             FileStatus::Untracked
         );
     });
@@ -8459,15 +8459,11 @@ async fn test_rename_work_directory(cx: &mut gpui::TestAppContext) {
             root_path.join("projects/project1").as_path()
         );
         assert_eq!(
-            repository
-                .status_for_path(&"a".into())
-                .map(|entry| entry.status),
+            repository.status_for_path(&"a".into()),
             Some(StatusCode::Modified.worktree()),
         );
         assert_eq!(
-            repository
-                .status_for_path(&"b".into())
-                .map(|entry| entry.status),
+            repository.status_for_path(&"b".into()),
             Some(FileStatus::Untracked),
         );
     });
@@ -8485,11 +8481,11 @@ async fn test_rename_work_directory(cx: &mut gpui::TestAppContext) {
             root_path.join("projects/project2").as_path()
         );
         assert_eq!(
-            repository.status_for_path(&"a".into()).unwrap().status,
+            repository.status_for_path(&"a".into()).unwrap(),
             StatusCode::Modified.worktree(),
         );
         assert_eq!(
-            repository.status_for_path(&"b".into()).unwrap().status,
+            repository.status_for_path(&"b".into()).unwrap(),
             FileStatus::Untracked,
         );
     });
@@ -8562,11 +8558,11 @@ async fn test_file_status(cx: &mut gpui::TestAppContext) {
         );
 
         assert_eq!(
-            repository.status_for_path(&B_TXT.into()).unwrap().status,
+            repository.status_for_path(&B_TXT.into()).unwrap(),
             FileStatus::Untracked,
         );
         assert_eq!(
-            repository.status_for_path(&F_TXT.into()).unwrap().status,
+            repository.status_for_path(&F_TXT.into()).unwrap(),
             FileStatus::Untracked,
         );
     });
@@ -8582,7 +8578,7 @@ async fn test_file_status(cx: &mut gpui::TestAppContext) {
     // The worktree detects that the file's git status has changed.
     repository.read_with(cx, |repository, _| {
         assert_eq!(
-            repository.status_for_path(&A_TXT.into()).unwrap().status,
+            repository.status_for_path(&A_TXT.into()).unwrap(),
             StatusCode::Modified.worktree(),
         );
     });
@@ -8600,7 +8596,7 @@ async fn test_file_status(cx: &mut gpui::TestAppContext) {
     // The worktree detects that the files' git status have changed.
     repository.read_with(cx, |repository, _cx| {
         assert_eq!(
-            repository.status_for_path(&F_TXT.into()).unwrap().status,
+            repository.status_for_path(&F_TXT.into()).unwrap(),
             FileStatus::Untracked,
         );
         assert_eq!(repository.status_for_path(&B_TXT.into()), None);
@@ -8623,11 +8619,11 @@ async fn test_file_status(cx: &mut gpui::TestAppContext) {
     repository.read_with(cx, |repository, _cx| {
         assert_eq!(repository.status_for_path(&A_TXT.into()), None);
         assert_eq!(
-            repository.status_for_path(&B_TXT.into()).unwrap().status,
+            repository.status_for_path(&B_TXT.into()).unwrap(),
             FileStatus::Untracked,
         );
         assert_eq!(
-            repository.status_for_path(&E_TXT.into()).unwrap().status,
+            repository.status_for_path(&E_TXT.into()).unwrap(),
             StatusCode::Modified.worktree(),
         );
     });
@@ -8666,8 +8662,7 @@ async fn test_file_status(cx: &mut gpui::TestAppContext) {
         assert_eq!(
             repository
                 .status_for_path(&Path::new(renamed_dir_name).join(RENAMED_FILE).into())
-                .unwrap()
-                .status,
+                .unwrap(),
             FileStatus::Untracked,
         );
     });
@@ -8690,8 +8685,7 @@ async fn test_file_status(cx: &mut gpui::TestAppContext) {
         assert_eq!(
             repository
                 .status_for_path(&Path::new(renamed_dir_name).join(RENAMED_FILE).into())
-                .unwrap()
-                .status,
+                .unwrap(),
             FileStatus::Untracked,
         );
     });
@@ -9000,7 +8994,7 @@ async fn test_git_worktrees_and_submodules(cx: &mut gpui::TestAppContext) {
     barrier.await.unwrap();
     worktree_repo.update(cx, |repo, _| {
         pretty_assertions::assert_eq!(
-            repo.status_for_path(&"src/b.txt".into()).unwrap().status,
+            repo.status_for_path(&"src/b.txt".into()).unwrap(),
             StatusCode::Modified.worktree(),
         );
     });
@@ -9039,7 +9033,7 @@ async fn test_git_worktrees_and_submodules(cx: &mut gpui::TestAppContext) {
     barrier.await.unwrap();
     submodule_repo.update(cx, |repo, _| {
         pretty_assertions::assert_eq!(
-            repo.status_for_path(&"c.txt".into()).unwrap().status,
+            repo.status_for_path(&"c.txt".into()).unwrap(),
             StatusCode::Modified.worktree(),
         );
     });
@@ -9304,9 +9298,7 @@ fn assert_entry_git_state(
     let entry = tree
         .entry_for_path(path)
         .unwrap_or_else(|| panic!("entry {path} not found"));
-    let status = repository
-        .status_for_path(&path.into())
-        .map(|entry| entry.status);
+    let status = repository.status_for_path(&path.into()).map(|entry| entry);
     let expected = index_status.map(|index_status| {
         TrackedStatus {
             index_status,

crates/project/src/worktree_store.rs

@@ -203,11 +203,10 @@ impl WorktreeStore {
         })
     }
 
-    pub fn entry_for_path(&self, path: &ProjectPath, cx: &App) -> Option<Entry> {
+    pub fn entry_for_path<'a>(&'a self, path: &ProjectPath, cx: &'a App) -> Option<&'a Entry> {
         self.worktree_for_id(path.worktree_id, cx)?
             .read(cx)
             .entry_for_path(&path.path)
-            .cloned()
     }
 
     pub fn create_worktree(

crates/repl/src/notebook/notebook_ui.rs

@@ -594,9 +594,10 @@ impl project::ProjectItem for NotebookItem {
                 };
 
                 let id = project
-                    .update(cx, |project, cx| project.entry_for_path(&path, cx))?
-                    .context("Entry not found")?
-                    .id;
+                    .update(cx, |project, cx| {
+                        project.entry_for_path(&path, cx).map(|entry| entry.id)
+                    })?
+                    .context("Entry not found")?;
 
                 cx.new(|_| NotebookItem {
                     path: abs_path,

crates/search/src/project_search.rs

@@ -3188,6 +3188,7 @@ pub mod tests {
                 .read(cx)
                 .entry_for_path(&(worktree_id, "a").into(), cx)
                 .expect("no entry for /a/ directory")
+                .clone()
         });
         assert!(a_dir_entry.is_dir());
         window

crates/zed/src/zed/edit_prediction_registry.rs

@@ -207,9 +207,11 @@ fn assign_edit_prediction_provider(
 
                 if let Some(buffer) = &singleton_buffer
                     && buffer.read(cx).file().is_some()
+                    // todo!
+                    && let Some(project) = editor.project()
                 {
                     zeta.update(cx, |zeta, cx| {
-                        zeta.register_buffer(buffer, cx);
+                        zeta.register_buffer(buffer, project, cx);
                     });
                 }
 

crates/zeta/Cargo.toml

@@ -21,6 +21,7 @@ ai_onboarding.workspace = true
 anyhow.workspace = true
 arrayvec.workspace = true
 client.workspace = true
+clock.workspace = true
 cloud_llm_client.workspace = true
 collections.workspace = true
 command_palette_hooks.workspace = true
@@ -39,6 +40,7 @@ language.workspace = true
 language_model.workspace = true
 log.workspace = true
 menu.workspace = true
+multi_buffer.workspace = true
 postage.workspace = true
 project.workspace = true
 rand.workspace = true

crates/zeta_cli/src/main.rs

@@ -18,7 +18,7 @@ use std::process::exit;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
-use zeta::{GatherContextOutput, PerformPredictEditsParams, Zeta, gather_context};
+use zeta::{CanCollectData, GatherContextOutput, PerformPredictEditsParams, Zeta, gather_context};
 
 use crate::headless::ZetaCliAppState;
 
@@ -129,15 +129,15 @@ async fn get_context(
         return Err(anyhow!("Absolute paths are not supported in --cursor"));
     }
 
-    let (project, _lsp_open_handle, buffer) = if use_language_server {
-        let (project, lsp_open_handle, buffer) =
+    let (_lsp_open_handle, buffer) = if use_language_server {
+        let (_project, lsp_open_handle, buffer) =
             open_buffer_with_language_server(&worktree_path, &cursor.path, app_state, cx).await?;
-        (Some(project), Some(lsp_open_handle), buffer)
+        (Some(lsp_open_handle), buffer)
     } else {
         let abs_path = worktree_path.join(&cursor.path);
         let content = smol::fs::read_to_string(&abs_path).await?;
         let buffer = cx.new(|cx| Buffer::local(content, cx))?;
-        (None, None, buffer)
+        (None, buffer)
     };
 
     let worktree_name = worktree_path
@@ -172,22 +172,14 @@ async fn get_context(
         None => String::new(),
     };
     // Enable gathering extra data not currently needed for edit predictions
-    let can_collect_data = true;
-    let git_info = None;
-    let mut gather_context_output = cx
-        .update(|cx| {
-            gather_context(
-                project.as_ref(),
-                full_path_str,
-                &snapshot,
-                clipped_cursor,
-                move || events,
-                can_collect_data,
-                git_info,
-                cx,
-            )
-        })?
-        .await;
+    let mut gather_context_output = gather_context(
+        full_path_str,
+        snapshot,
+        clipped_cursor,
+        move || events,
+        CanCollectData(true),
+    )
+    .await;
 
     // Disable data collection for these requests, as this is currently just used for evals
     if let Ok(gather_context_output) = gather_context_output.as_mut() {