zeta2: Build edit prediction prompt and process model output in client (#41870)

Release Notes: - N/A ---------

Max Brunsfeld committed 1 year ago

Commit 784fdca1 parent fb8797232 files · +2197 −2391Browse files

Unified · Split

32 files

Cargo.lock +5 −1
crates/acp_thread/Cargo.toml +1 −0
crates/cloud_llm_client/src/cloud_llm_client.rs +0 −1
crates/cloud_llm_client/src/predict_edits_v3.rs +31 −6
crates/cloud_llm_client/src/udiff.rs +0 −294
crates/cloud_zeta2_prompt/Cargo.toml +2 −0
crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs +4 −3
crates/cloud_zeta2_prompt/src/retrieval_prompt.rs +92 −0
crates/codestral/src/codestral.rs +4 −4
crates/edit_prediction/src/edit_prediction.rs +5 −5
crates/editor/src/edit_prediction_tests.rs +4 −4
crates/editor/src/editor.rs +7 −12
crates/editor/src/editor_tests.rs +1 −1
crates/language/src/buffer.rs +6 −5
crates/language/src/buffer_tests.rs +3 −5
crates/open_ai/src/open_ai.rs +22 −5
crates/supermaven/src/supermaven_completion_provider.rs +7 −3
crates/zeta/src/zeta.rs +22 −27
crates/zeta2/Cargo.toml +2 −1
crates/zeta2/src/prediction.rs +41 −219
crates/zeta2/src/related_excerpts.rs +0 −717
crates/zeta2/src/retrieval_search.rs +194 −0
crates/zeta2/src/udiff.rs +1024 −0
crates/zeta2/src/zeta2.rs +557 −269
crates/zeta2_tools/Cargo.toml +2 −1
crates/zeta2_tools/src/zeta2_context_view.rs +10 −40
crates/zeta2_tools/src/zeta2_tools.rs +30 −63
crates/zeta_cli/src/evaluate.rs +4 −5
crates/zeta_cli/src/example.rs +64 −415
crates/zeta_cli/src/main.rs +6 −223
crates/zeta_cli/src/paths.rs +8 −0
crates/zeta_cli/src/predict.rs +39 −62

  "util",
  "uuid",
  "watch",
+ "zlog",
 ]
 [[package]]
  "indoc",
  "ordered-float 2.10.1",
  "rustc-hash 2.1.1",
+ "schemars 1.0.4",
  "serde",
+ "serde_json",
  "strum 0.27.2",
 ]
  "language_model",
  "log",
  "lsp",
+ "open_ai",
  "pretty_assertions",
  "project",
  "release_channel",
- "schemars 1.0.4",
  "serde",
  "serde_json",
  "settings",
  "uuid",
  "workspace",
  "worktree",
+ "zlog",
 ]
 [[package]]

crates/acp_thread/Cargo.toml+1−0

@@ -56,3 +56,4 @@ rand.workspace = true · 55 unmodified lines
56	tempfile.workspace = true	56	tempfile.workspace = true
57	util.workspace = true	57	util.workspace = true
58	settings.workspace = true	58	settings.workspace = true
		59	zlog.workspace = true

crates/cloud_llm_client/src/cloud_llm_client.rs+0−1

@@ -1,5 +1,4 @@
1	pub mod predict_edits_v3;	1	pub mod predict_edits_v3;
2	pub mod udiff;
3		2
4	use std::str::FromStr;	3	use std::str::FromStr;
5	use std::sync::Arc;	4	use std::sync::Arc;

crates/cloud_llm_client/src/predict_edits_v3.rs+31−6

 use chrono::Duration;
 use serde::{Deserialize, Serialize};
 use std::{
-    fmt::Display,
+    fmt::{Display, Write as _},
     ops::{Add, Range, Sub},
     path::{Path, PathBuf},
     sync::Arc,
 use crate::PredictEditsGitInfo;
-// TODO: snippet ordering within file / relative to excerpt
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PlanContextRetrievalRequest {
+    pub excerpt: String,
+    pub excerpt_path: Arc<Path>,
+    pub excerpt_line_range: Range<Line>,
+    pub cursor_file_max_row: Line,
+    pub events: Vec<Event>,
+}
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PredictEditsRequest {
                     write!(
                         f,
                         "// User accepted prediction:\n--- a/{}\n+++ b/{}\n{diff}",
-                        old_path.display(),
+                        DiffPathFmt(old_path),
-                        new_path.display()
+                        DiffPathFmt(new_path)
                     )
                 } else {
                     write!(
                         f,
                         "--- a/{}\n+++ b/{}\n{diff}",
-                        old_path.display(),
+                        DiffPathFmt(old_path),
-                        new_path.display()
+                        DiffPathFmt(new_path)
                     )
                 }
             }
     }
 }
+/// always format the Path as a unix path with `/` as the path sep in Diffs
+pub struct DiffPathFmt<'a>(pub &'a Path);
+impl<'a> std::fmt::Display for DiffPathFmt<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut is_first = true;
+        for component in self.0.components() {
+            if !is_first {
+                f.write_char('/')?;
+            } else {
+                is_first = false;
+            }
+            write!(f, "{}", component.as_os_str().display())?;
+        }
+        Ok(())
+    }
+}
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Signature {
     pub text: String,

crates/cloud_llm_client/src/udiff.rs+0−294

@@ -1,294 +0,0 @@
1	use std::{borrow::Cow, fmt::Display};
2
3	#[derive(Debug, PartialEq)]
4	pub enum DiffLine<'a> {
5	OldPath { path: Cow<'a, str> },
6	NewPath { path: Cow<'a, str> },
7	HunkHeader(Option<HunkLocation>),
8	Context(&'a str),
9	Deletion(&'a str),
10	Addition(&'a str),
11	Garbage(&'a str),
12	}
13
14	#[derive(Debug, PartialEq)]
15	pub struct HunkLocation {
16	start_line_old: u32,
17	count_old: u32,
18	start_line_new: u32,
19	count_new: u32,
20	}
21
22	impl<'a> DiffLine<'a> {
23	pub fn parse(line: &'a str) -> Self {
24	Self::try_parse(line).unwrap_or(Self::Garbage(line))
25	}
26
27	fn try_parse(line: &'a str) -> Option<Self> {
28	if let Some(header) = line.strip_prefix("---").and_then(eat_required_whitespace) {
29	let path = parse_header_path("a/", header);
30	Some(Self::OldPath { path })
31	} else if let Some(header) = line.strip_prefix("+++").and_then(eat_required_whitespace) {
32	Some(Self::NewPath {
33	path: parse_header_path("b/", header),
34	})
35	} else if let Some(header) = line.strip_prefix("@@").and_then(eat_required_whitespace) {
36	if header.starts_with("...") {
37	return Some(Self::HunkHeader(None));
38	}
39
40	let (start_line_old, header) = header.strip_prefix('-')?.split_once(',')?;
41	let mut parts = header.split_ascii_whitespace();
42	let count_old = parts.next()?;
43	let (start_line_new, count_new) = parts.next()?.strip_prefix('+')?.split_once(',')?;
44
45	Some(Self::HunkHeader(Some(HunkLocation {
46	start_line_old: start_line_old.parse::<u32>().ok()?.saturating_sub(1),
47	count_old: count_old.parse().ok()?,
48	start_line_new: start_line_new.parse::<u32>().ok()?.saturating_sub(1),
49	count_new: count_new.parse().ok()?,
50	})))
51	} else if let Some(deleted_header) = line.strip_prefix("-") {
52	Some(Self::Deletion(deleted_header))
53	} else if line.is_empty() {
54	Some(Self::Context(""))
55	} else if let Some(context) = line.strip_prefix(" ") {
56	Some(Self::Context(context))
57	} else {
58	Some(Self::Addition(line.strip_prefix("+")?))
59	}
60	}
61	}
62
63	impl<'a> Display for DiffLine<'a> {
64	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65	match self {
66	DiffLine::OldPath { path } => write!(f, "--- {path}"),
67	DiffLine::NewPath { path } => write!(f, "+++ {path}"),
68	DiffLine::HunkHeader(Some(hunk_location)) => {
69	write!(
70	f,
71	"@@ -{},{} +{},{} @@",
72	hunk_location.start_line_old + 1,
73	hunk_location.count_old,
74	hunk_location.start_line_new + 1,
75	hunk_location.count_new
76	)
77	}
78	DiffLine::HunkHeader(None) => write!(f, "@@ ... @@"),
79	DiffLine::Context(content) => write!(f, " {content}"),
80	DiffLine::Deletion(content) => write!(f, "-{content}"),
81	DiffLine::Addition(content) => write!(f, "+{content}"),
82	DiffLine::Garbage(line) => write!(f, "{line}"),
83	}
84	}
85	}
86
87	fn parse_header_path<'a>(strip_prefix: &'static str, header: &'a str) -> Cow<'a, str> {
88	if !header.contains(['"', '\\']) {
89	let path = header.split_ascii_whitespace().next().unwrap_or(header);
90	return Cow::Borrowed(path.strip_prefix(strip_prefix).unwrap_or(path));
91	}
92
93	let mut path = String::with_capacity(header.len());
94	let mut in_quote = false;
95	let mut chars = header.chars().peekable();
96	let mut strip_prefix = Some(strip_prefix);
97
98	while let Some(char) = chars.next() {
99	if char == '"' {
100	in_quote = !in_quote;
101	} else if char == '\\' {
102	let Some(&next_char) = chars.peek() else {
103	break;
104	};
105	chars.next();
106	path.push(next_char);
107	} else if char.is_ascii_whitespace() && !in_quote {
108	break;
109	} else {
110	path.push(char);
111	}
112
113	if let Some(prefix) = strip_prefix
114	&& path == prefix
115	{
116	strip_prefix.take();
117	path.clear();
118	}
119	}
120
121	Cow::Owned(path)
122	}
123
124	fn eat_required_whitespace(header: &str) -> Option<&str> {
125	let trimmed = header.trim_ascii_start();
126
127	if trimmed.len() == header.len() {
128	None
129	} else {
130	Some(trimmed)
131	}
132	}
133
134	#[cfg(test)]
135	mod tests {
136	use super::*;
137	use indoc::indoc;
138
139	#[test]
140	fn parse_lines_simple() {
141	let input = indoc! {"
142	diff --git a/text.txt b/text.txt
143	index 86c770d..a1fd855 100644
144	--- a/file.txt
145	+++ b/file.txt
146	@@ -1,2 +1,3 @@
147	context
148	-deleted
149	+inserted
150	garbage
151
152	--- b/file.txt
153	+++ a/file.txt
154	"};
155
156	let lines = input.lines().map(DiffLine::parse).collect::<Vec<_>>();
157
158	pretty_assertions::assert_eq!(
159	lines,
160	&[
161	DiffLine::Garbage("diff --git a/text.txt b/text.txt"),
162	DiffLine::Garbage("index 86c770d..a1fd855 100644"),
163	DiffLine::OldPath {
164	path: "file.txt".into()
165	},
166	DiffLine::NewPath {
167	path: "file.txt".into()
168	},
169	DiffLine::HunkHeader(Some(HunkLocation {
170	start_line_old: 0,
171	count_old: 2,
172	start_line_new: 0,
173	count_new: 3
174	})),
175	DiffLine::Context("context"),
176	DiffLine::Deletion("deleted"),
177	DiffLine::Addition("inserted"),
178	DiffLine::Garbage("garbage"),
179	DiffLine::Context(""),
180	DiffLine::OldPath {
181	path: "b/file.txt".into()
182	},
183	DiffLine::NewPath {
184	path: "a/file.txt".into()
185	},
186	]
187	);
188	}
189
190	#[test]
191	fn file_header_extra_space() {
192	let options = ["--- file", "--- file", "---\tfile"];
193
194	for option in options {
195	pretty_assertions::assert_eq!(
196	DiffLine::parse(option),
197	DiffLine::OldPath {
198	path: "file".into()
199	},
200	"{option}",
201	);
202	}
203	}
204
205	#[test]
206	fn hunk_header_extra_space() {
207	let options = [
208	"@@ -1,2 +1,3 @@",
209	"@@ -1,2 +1,3 @@",
210	"@@\t-1,2\t+1,3\t@@",
211	"@@ -1,2 +1,3 @@",
212	"@@ -1,2 +1,3 @@",
213	"@@ -1,2 +1,3 @@",
214	"@@ -1,2 +1,3 @@ garbage",
215	];
216
217	for option in options {
218	pretty_assertions::assert_eq!(
219	DiffLine::parse(option),
220	DiffLine::HunkHeader(Some(HunkLocation {
221	start_line_old: 0,
222	count_old: 2,
223	start_line_new: 0,
224	count_new: 3
225	})),
226	"{option}",
227	);
228	}
229	}
230
231	#[test]
232	fn hunk_header_without_location() {
233	pretty_assertions::assert_eq!(DiffLine::parse("@@ ... @@"), DiffLine::HunkHeader(None));
234	}
235
236	#[test]
237	fn test_parse_path() {
238	assert_eq!(parse_header_path("a/", "foo.txt"), "foo.txt");
239	assert_eq!(
240	parse_header_path("a/", "foo/bar/baz.txt"),
241	"foo/bar/baz.txt"
242	);
243	assert_eq!(parse_header_path("a/", "a/foo.txt"), "foo.txt");
244	assert_eq!(
245	parse_header_path("a/", "a/foo/bar/baz.txt"),
246	"foo/bar/baz.txt"
247	);
248
249	// Extra
250	assert_eq!(
251	parse_header_path("a/", "a/foo/bar/baz.txt 2025"),
252	"foo/bar/baz.txt"
253	);
254	assert_eq!(
255	parse_header_path("a/", "a/foo/bar/baz.txt\t2025"),
256	"foo/bar/baz.txt"
257	);
258	assert_eq!(
259	parse_header_path("a/", "a/foo/bar/baz.txt \""),
260	"foo/bar/baz.txt"
261	);
262
263	// Quoted
264	assert_eq!(
265	parse_header_path("a/", "a/foo/bar/\"baz quox.txt\""),
266	"foo/bar/baz quox.txt"
267	);
268	assert_eq!(
269	parse_header_path("a/", "\"a/foo/bar/baz quox.txt\""),
270	"foo/bar/baz quox.txt"
271	);
272	assert_eq!(
273	parse_header_path("a/", "\"foo/bar/baz quox.txt\""),
274	"foo/bar/baz quox.txt"
275	);
276	assert_eq!(parse_header_path("a/", "\"whatever 🤷\""), "whatever 🤷");
277	assert_eq!(
278	parse_header_path("a/", "\"foo/bar/baz quox.txt\" 2025"),
279	"foo/bar/baz quox.txt"
280	);
281	// unescaped quotes are dropped
282	assert_eq!(parse_header_path("a/", "foo/\"bar\""), "foo/bar");
283
284	// Escaped
285	assert_eq!(
286	parse_header_path("a/", "\"foo/\\\"bar\\\"/baz.txt\""),
287	"foo/\"bar\"/baz.txt"
288	);
289	assert_eq!(
290	parse_header_path("a/", "\"C:\\\\Projects\\\\My App\\\\old file.txt\""),
291	"C:\\Projects\\My App\\old file.txt"
292	);
293	}
294	}

crates/cloud_zeta2_prompt/Cargo.toml+2−0

@@ -17,5 +17,7 @@ cloud_llm_client.workspace = true · 16 unmodified lines
17	indoc.workspace = true	17	indoc.workspace = true
18	ordered-float.workspace = true	18	ordered-float.workspace = true
19	rustc-hash.workspace = true	19	rustc-hash.workspace = true
		20	schemars.workspace = true
20	serde.workspace = true	21	serde.workspace = true
		22	serde_json.workspace = true
21	strum.workspace = true	23	strum.workspace = true

crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs+4−3

 //! Zeta2 prompt planning and generation code shared with cloud.
+pub mod retrieval_prompt;
 use anyhow::{Context as _, Result, anyhow};
 use cloud_llm_client::predict_edits_v3::{
-    self, Excerpt, Line, Point, PromptFormat, ReferencedDeclaration,
+    self, DiffPathFmt, Excerpt, Line, Point, PromptFormat, ReferencedDeclaration,
 };
 use indoc::indoc;
 use ordered_float::OrderedFloat;
     include_line_numbers: bool,
     output: &'a mut String,
 ) {
-    writeln!(output, "`````{}", path.display()).unwrap();
+    writeln!(output, "`````{}", DiffPathFmt(path)).unwrap();
     write_excerpts(
         excerpts,
         sorted_insertions,
     }
 }
-fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
+pub fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
     if events.is_empty() {
         return;
     };

crates/cloud_zeta2_prompt/src/retrieval_prompt.rs+92−0

@@ -0,0 +1,92 @@
	1	use anyhow::Result;
	2	use cloud_llm_client::predict_edits_v3::{self, Excerpt};
	3	use indoc::indoc;
	4	use schemars::JsonSchema;
	5	use serde::{Deserialize, Serialize};
	6	use std::{fmt::Write, sync::LazyLock};
	7
	8	use crate::{push_events, write_codeblock};
	9
	10	pub fn build_prompt(request: predict_edits_v3::PlanContextRetrievalRequest) -> Result<String> {
	11	let mut prompt = SEARCH_INSTRUCTIONS.to_string();
	12
	13	if !request.events.is_empty() {
	14	writeln!(&mut prompt, "## User Edits\n")?;
	15	push_events(&mut prompt, &request.events);
	16	}
	17
	18	writeln!(&mut prompt, "## Excerpt around the cursor\n")?;
	19	write_codeblock(
	20	&request.excerpt_path,
	21	&[Excerpt {
	22	start_line: request.excerpt_line_range.start,
	23	text: request.excerpt.into(),
	24	}],
	25	&[],
	26	request.cursor_file_max_row,
	27	true,
	28	&mut prompt,
	29	);
	30
	31	writeln!(&mut prompt, "{TOOL_USE_REMINDER}")?;
	32
	33	Ok(prompt)
	34	}
	35
	36	/// Search for relevant code
	37	///
	38	/// For the best results, run multiple queries at once with a single invocation of this tool.
	39	#[derive(Clone, Deserialize, Serialize, JsonSchema)]
	40	pub struct SearchToolInput {
	41	/// An array of queries to run for gathering context relevant to the next prediction
	42	#[schemars(length(max = 5))]
	43	pub queries: Box<[SearchToolQuery]>,
	44	}
	45
	46	#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
	47	pub struct SearchToolQuery {
	48	/// A glob pattern to match file paths in the codebase
	49	pub glob: String,
	50	/// A regular expression to match content within the files matched by the glob pattern
	51	pub regex: String,
	52	}
	53
	54	pub static TOOL_SCHEMA: LazyLock<(serde_json::Value, String)> = LazyLock::new(\|\| {
	55	let schema = schemars::schema_for!(SearchToolInput);
	56
	57	let description = schema
	58	.get("description")
	59	.and_then(\|description\| description.as_str())
	60	.unwrap()
	61	.to_string();
	62
	63	(schema.into(), description)
	64	});
	65
	66	pub const TOOL_NAME: &str = "search";
	67
	68	const SEARCH_INSTRUCTIONS: &str = indoc! {r#"
	69	## Task
	70
	71	You are part of an edit prediction system in a code editor. Your role is to identify relevant code locations
	72	that will serve as context for predicting the next required edit.
	73
	74	Your task:
	75	- Analyze the user's recent edits and current cursor context
	76	- Use the `search` tool to find code that may be relevant for predicting the next edit
	77	- Focus on finding:
	78	- Code patterns that might need similar changes based on the recent edits
	79	- Functions, variables, types, and constants referenced in the current cursor context
	80	- Related implementations, usages, or dependencies that may require consistent updates
	81
	82	Important constraints:
	83	- This conversation has exactly 2 turns
	84	- You must make ALL search queries in your first response via the `search` tool
	85	- All queries will be executed in parallel and results returned together
	86	- In the second turn, you will select the most relevant results via the `select` tool.
	87	"#};
	88
	89	const TOOL_USE_REMINDER: &str = indoc! {"
	90	--
	91	Use the `search` tool now
	92	"};

crates/codestral/src/codestral.rs+4−4

@@ -34,7 +34,7 @@ struct CurrentCompletion { · 33 unmodified lines
34	snapshot: BufferSnapshot,	34	snapshot: BufferSnapshot,
35	/// The edits that should be applied to transform the original text into the predicted text.	35	/// The edits that should be applied to transform the original text into the predicted text.
36	/// Each edit is a range in the buffer and the text to replace it with.	36	/// Each edit is a range in the buffer and the text to replace it with.
37	edits: Arc<[(Range<Anchor>, String)]>,	37	edits: Arc<[(Range<Anchor>, Arc<str>)]>,
38	/// Preview of how the buffer will look after applying the edits.	38	/// Preview of how the buffer will look after applying the edits.
39	edit_preview: EditPreview,	39	edit_preview: EditPreview,
40	}	40	}
@@ -42,7 +42,7 @@ struct CurrentCompletion { · 1 unmodified line
42	impl CurrentCompletion {	42	impl CurrentCompletion {
43	/// Attempts to adjust the edits based on changes made to the buffer since the completion was generated.	43	/// Attempts to adjust the edits based on changes made to the buffer since the completion was generated.
44	/// Returns None if the user's edits conflict with the predicted edits.	44	/// Returns None if the user's edits conflict with the predicted edits.
45	fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, String)>> {	45	fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
46	edit_prediction::interpolate_edits(&self.snapshot, new_snapshot, &self.edits)	46	edit_prediction::interpolate_edits(&self.snapshot, new_snapshot, &self.edits)
47	}	47	}
48	}	48	}
@@ -281,8 +281,8 @@ impl EditPredictionProvider for CodestralCompletionProvider { · 232 unmodified lines
281	return Ok(());	281	return Ok(());
282	}	282	}
283		283
284	let edits: Arc<[(Range<Anchor>, String)]> =	284	let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
285	vec![(cursor_position..cursor_position, completion_text)].into();	285	vec![(cursor_position..cursor_position, completion_text.into())].into();
286	let edit_preview = buffer	286	let edit_preview = buffer
287	.read_with(cx, \|buffer, cx\| buffer.preview_edits(edits.clone(), cx))?	287	.read_with(cx, \|buffer, cx\| buffer.preview_edits(edits.clone(), cx))?
288	.await;	288	.await;

crates/edit_prediction/src/edit_prediction.rs+5−5

@@ -1,4 +1,4 @@
1	use std::ops::Range;	1	use std::{ops::Range, sync::Arc};
2		2
3	use client::EditPredictionUsage;	3	use client::EditPredictionUsage;
4	use gpui::{App, Context, Entity, SharedString};	4	use gpui::{App, Context, Entity, SharedString};
@@ -19,7 +19,7 @@ pub enum EditPrediction { · 14 unmodified lines
19	/// Edits within the buffer that requested the prediction	19	/// Edits within the buffer that requested the prediction
20	Local {	20	Local {
21	id: Option<SharedString>,	21	id: Option<SharedString>,
22	edits: Vec<(Range<language::Anchor>, String)>,	22	edits: Vec<(Range<language::Anchor>, Arc<str>)>,
23	edit_preview: Option<language::EditPreview>,	23	edit_preview: Option<language::EditPreview>,
24	},	24	},
25	/// Jump to a different file from the one that requested the prediction	25	/// Jump to a different file from the one that requested the prediction
@@ -248,8 +248,8 @@ where · 222 unmodified lines
248	pub fn interpolate_edits(	248	pub fn interpolate_edits(
249	old_snapshot: &BufferSnapshot,	249	old_snapshot: &BufferSnapshot,
250	new_snapshot: &BufferSnapshot,	250	new_snapshot: &BufferSnapshot,
251	current_edits: &[(Range<Anchor>, String)],	251	current_edits: &[(Range<Anchor>, Arc<str>)],
252	) -> Option<Vec<(Range<Anchor>, String)>> {	252	) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
253	let mut edits = Vec::new();	253	let mut edits = Vec::new();
254		254
255	let mut model_edits = current_edits.iter().peekable();	255	let mut model_edits = current_edits.iter().peekable();
@@ -274,7 +274,7 @@ pub fn interpolate_edits( · 18 unmodified lines
274	if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {	274	if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
275	if !model_suffix.is_empty() {	275	if !model_suffix.is_empty() {
276	let anchor = old_snapshot.anchor_after(user_edit.old.end);	276	let anchor = old_snapshot.anchor_after(user_edit.old.end);
277	edits.push((anchor..anchor, model_suffix.to_string()));	277	edits.push((anchor..anchor, model_suffix.into()));
278	}	278	}
279		279
280	model_edits.next();	280	model_edits.next();

crates/editor/src/edit_prediction_tests.rs+4−4

@@ -2,7 +2,7 @@ use edit_prediction::EditPredictionProvider; · 1 unmodified line
2	use gpui::{Entity, KeyBinding, Modifiers, prelude::*};	2	use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
3	use indoc::indoc;	3	use indoc::indoc;
4	use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};	4	use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
5	use std::ops::Range;	5	use std::{ops::Range, sync::Arc};
6	use text::{Point, ToOffset};	6	use text::{Point, ToOffset};
7		7
8	use crate::{	8	use crate::{
@@ -24,7 +24,7 @@ async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) { · 15 unmodified lines
24		24
25	assert_editor_active_edit_completion(&mut cx, \|_, edits\| {	25	assert_editor_active_edit_completion(&mut cx, \|_, edits\| {
26	assert_eq!(edits.len(), 1);	26	assert_eq!(edits.len(), 1);
27	assert_eq!(edits[0].1.as_str(), "-273.15");	27	assert_eq!(edits[0].1.as_ref(), "-273.15");
28	});	28	});
29		29
30	accept_completion(&mut cx);	30	accept_completion(&mut cx);
@@ -46,7 +46,7 @@ async fn test_edit_prediction_modification(cx: &mut gpui::TestAppContext) { · 15 unmodified lines
46		46
47	assert_editor_active_edit_completion(&mut cx, \|_, edits\| {	47	assert_editor_active_edit_completion(&mut cx, \|_, edits\| {
48	assert_eq!(edits.len(), 1);	48	assert_eq!(edits.len(), 1);
49	assert_eq!(edits[0].1.as_str(), "3.14159");	49	assert_eq!(edits[0].1.as_ref(), "3.14159");
50	});	50	});
51		51
52	accept_completion(&mut cx);	52	accept_completion(&mut cx);
@@ -330,7 +330,7 @@ async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestA · 277 unmodified lines
330		330
331	fn assert_editor_active_edit_completion(	331	fn assert_editor_active_edit_completion(
332	cx: &mut EditorTestContext,	332	cx: &mut EditorTestContext,
333	assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, String)>),	333	assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, Arc<str>)>),
334	) {	334	) {
335	cx.editor(\|editor, _, cx\| {	335	cx.editor(\|editor, _, cx\| {
336	let completion_state = editor	336	let completion_state = editor

crates/editor/src/editor.rs+7−12

 enum EditPrediction {
     Edit {
-        edits: Vec<(Range<Anchor>, String)>,
+        edits: Vec<(Range<Anchor>, Arc<str>)>,
         edit_preview: Option<EditPreview>,
         display_mode: EditDisplayMode,
         snapshot: BufferSnapshot,
                         let inlay = Inlay::edit_prediction(
                             post_inc(&mut self.next_inlay_id),
                             range.start,
-                            new_text.as_str(),
+                            new_text.as_ref(),
                         );
                         inlay_ids.push(inlay.id);
                         inlays.push(inlay);
         newest_selection_head: Option<DisplayPoint>,
         editor_width: Pixels,
         style: &EditorStyle,
-        edits: &Vec<(Range<Anchor>, String)>,
+        edits: &Vec<(Range<Anchor>, Arc<str>)>,
         edit_preview: &Option<language::EditPreview>,
         snapshot: &language::BufferSnapshot,
         window: &mut Window,
 fn edit_prediction_edit_text(
     current_snapshot: &BufferSnapshot,
-    edits: &[(Range<Anchor>, String)],
+    edits: &[(Range<Anchor>, impl AsRef<str>)],
     edit_preview: &EditPreview,
     include_deletions: bool,
     cx: &App,
 ) -> HighlightedText {
     let edits = edits
         .iter()
-        .map(|(anchor, text)| {
+        .map(|(anchor, text)| (anchor.start.text_anchor..anchor.end.text_anchor, text))
-            (
-                anchor.start.text_anchor..anchor.end.text_anchor,
-                text.clone(),
-            )
-        })
         .collect::<Vec<_>>();
     edit_preview.highlight_edits(current_snapshot, &edits, include_deletions, cx)
 }
-fn edit_prediction_fallback_text(edits: &[(Range<Anchor>, String)], cx: &App) -> HighlightedText {
+fn edit_prediction_fallback_text(edits: &[(Range<Anchor>, Arc<str>)], cx: &App) -> HighlightedText {
     // Fallback for providers that don't provide edit_preview (like Copilot/Supermaven)
     // Just show the raw edit text with basic styling
     let mut text = String::new();
 }
 fn all_edits_insertions_or_deletions(
-    edits: &Vec<(Range<Anchor>, String)>,
+    edits: &Vec<(Range<Anchor>, Arc<str>)>,
     snapshot: &MultiBufferSnapshot,
 ) -> bool {
     let mut all_insertions = true;

crates/editor/src/editor_tests.rs+1−1

@@ -22915,7 +22915,7 @@ async fn assert_highlighted_edits( · 22914 unmodified lines
22915	let text_anchor_edits = edits	22915	let text_anchor_edits = edits
22916	.clone()	22916	.clone()
22917	.into_iter()	22917	.into_iter()
22918	.map(\|(range, edit)\| (range.start.text_anchor..range.end.text_anchor, edit))	22918	.map(\|(range, edit)\| (range.start.text_anchor..range.end.text_anchor, edit.into()))
22919	.collect::<Vec<_>>();	22919	.collect::<Vec<_>>();
22920		22920
22921	let edit_preview = window	22921	let edit_preview = window

crates/language/src/buffer.rs+6−5

     pub fn highlight_edits(
         &self,
         current_snapshot: &BufferSnapshot,
-        edits: &[(Range<Anchor>, String)],
+        edits: &[(Range<Anchor>, impl AsRef<str>)],
         include_deletions: bool,
         cx: &App,
     ) -> HighlightedText {
                 .end
                 .bias_right(&self.old_snapshot)
                 .to_offset(&self.applied_edits_snapshot);
-            let edit_start_in_preview_snapshot = edit_new_end_in_preview_snapshot - edit_text.len();
+            let edit_start_in_preview_snapshot =
+                edit_new_end_in_preview_snapshot - edit_text.as_ref().len();
             let unchanged_range_in_preview_snapshot =
                 offset_in_preview_snapshot..edit_start_in_preview_snapshot;
                 );
             }
-            if !edit_text.is_empty() {
+            if !edit_text.as_ref().is_empty() {
                 highlighted_text.add_text_from_buffer_range(
                     edit_start_in_preview_snapshot..edit_new_end_in_preview_snapshot,
                     &self.applied_edits_snapshot,
         highlighted_text.build()
     }
-    fn compute_visible_range(&self, edits: &[(Range<Anchor>, String)]) -> Option<Range<usize>> {
+    fn compute_visible_range<T>(&self, edits: &[(Range<Anchor>, T)]) -> Option<Range<usize>> {
         let (first, _) = edits.first()?;
         let (last, _) = edits.last()?;
     pub fn preview_edits(
         &self,
-        edits: Arc<[(Range<Anchor>, String)]>,
+        edits: Arc<[(Range<Anchor>, Arc<str>)]>,
         cx: &App,
     ) -> Task<EditPreview> {
         let registry = self.language_registry();

crates/language/src/buffer_tests.rs+3−5

                 .map(|(range, text)| {
                     (
                         buffer.anchor_before(range.start)..buffer.anchor_after(range.end),
-                        text.to_string(),
+                        text.into(),
                     )
                 })
-                .collect::<Vec<_>>()
+                .collect::<Arc<[_]>>()
         });
         let edit_preview = buffer
-            .read_with(cx, |buffer, cx| {
+            .read_with(cx, |buffer, cx| buffer.preview_edits(edits.clone(), cx))
-                buffer.preview_edits(edits.clone().into(), cx)
-            })
             .await;
         let highlighted_edits = cx.read(|cx| {
             edit_preview.highlight_edits(&buffer.read(cx).snapshot(), &edits, include_deletions, cx)

crates/open_ai/src/open_ai.rs+22−5

     pub parameters: Option<Value>,
 }
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
 #[serde(tag = "role", rename_all = "lowercase")]
 pub enum RequestMessage {
     Assistant {
     pub detail: Option<String>,
 }
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
 pub struct ToolCall {
     pub id: String,
     #[serde(flatten)]
     pub content: ToolCallContent,
 }
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
 #[serde(tag = "type", rename_all = "lowercase")]
 pub enum ToolCallContent {
     Function { function: FunctionContent },
 }
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
 pub struct FunctionContent {
     pub name: String,
     pub arguments: String,
 }
+#[derive(Clone, Serialize, Deserialize, Debug)]
+pub struct Response {
+    pub id: String,
+    pub object: String,
+    pub created: u64,
+    pub model: String,
+    pub choices: Vec<Choice>,
+    pub usage: Usage,
+}
+#[derive(Clone, Serialize, Deserialize, Debug)]
+pub struct Choice {
+    pub index: u32,
+    pub message: RequestMessage,
+    pub finish_reason: Option<String>,
+}
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
 pub struct ResponseMessageDelta {
     pub role: Option<Role>,
     pub arguments: Option<String>,
 }
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Clone, Serialize, Deserialize, Debug)]
 pub struct Usage {
     pub prompt_tokens: u64,
     pub completion_tokens: u64,

crates/supermaven/src/supermaven_completion_provider.rs+7−3

 use std::{
     ops::{AddAssign, Range},
     path::Path,
+    sync::Arc,
     time::Duration,
 };
 use text::{ToOffset, ToPoint};
 ) -> EditPrediction {
     let buffer_text = snapshot.text_for_range(delete_range).collect::<String>();
-    let mut edits: Vec<(Range<language::Anchor>, String)> = Vec::new();
+    let mut edits: Vec<(Range<language::Anchor>, Arc<str>)> = Vec::new();
     let completion_graphemes: Vec<&str> = completion_text.graphemes(true).collect();
     let buffer_graphemes: Vec<&str> = buffer_text.graphemes(true).collect();
                 if k != 0 {
                     let offset = snapshot.anchor_after(offset);
                     // the range from the current position to item is an inlay.
-                    let edit = (offset..offset, completion_graphemes[i..i + k].join(""));
+                    let edit = (
+                        offset..offset,
+                        completion_graphemes[i..i + k].join("").into(),
+                    );
                     edits.push(edit);
                 }
                 i += k + 1;
         // there is leftover completion text, so drop it as an inlay.
         let edit_range = offset..offset;
         let edit_text = completion_graphemes[i..].join("");
-        edits.push((edit_range, edit_text));
+        edits.push((edit_range, edit_text.into()));
     }
     EditPrediction::Local {

crates/zeta/src/zeta.rs+22−27

     path: Arc<Path>,
     excerpt_range: Range<usize>,
     cursor_offset: usize,
-    edits: Arc<[(Range<Anchor>, String)]>,
+    edits: Arc<[(Range<Anchor>, Arc<str>)]>,
     snapshot: BufferSnapshot,
     edit_preview: EditPreview,
     input_outline: Arc<str>,
             .duration_since(self.buffer_snapshotted_at)
     }
-    fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, String)>> {
+    fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
         edit_prediction::interpolate_edits(&self.snapshot, new_snapshot, &self.edits)
     }
 }
         cx.spawn(async move |cx| {
             let output_excerpt: Arc<str> = output_excerpt.into();
-            let edits: Arc<[(Range<Anchor>, String)]> = cx
+            let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx
                 .background_spawn({
                     let output_excerpt = output_excerpt.clone();
                     let editable_range = editable_range.clone();
                 let edits = edits.clone();
                 move |buffer, cx| {
                     let new_snapshot = buffer.snapshot();
-                    let edits: Arc<[(Range<Anchor>, String)]> =
+                    let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
                         edit_prediction::interpolate_edits(&snapshot, &new_snapshot, &edits)?
                             .into();
                     Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
         output_excerpt: Arc<str>,
         editable_range: Range<usize>,
         snapshot: &BufferSnapshot,
-    ) -> Result<Vec<(Range<Anchor>, String)>> {
+    ) -> Result<Vec<(Range<Anchor>, Arc<str>)>> {
         let content = output_excerpt.replace(CURSOR_MARKER, "");
         let start_markers = content
         new_text: &str,
         offset: usize,
         snapshot: &BufferSnapshot,
-    ) -> Vec<(Range<Anchor>, String)> {
+    ) -> Vec<(Range<Anchor>, Arc<str>)> {
         text_diff(&old_text, new_text)
             .into_iter()
             .map(|(mut old_range, new_text)| {
                 );
                 old_range.end = old_range.end.saturating_sub(suffix_len);
-                let new_text = new_text[prefix_len..new_text.len() - suffix_len].to_string();
+                let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
                 let range = if old_range.is_empty() {
                     let anchor = snapshot.anchor_after(old_range.start);
                     anchor..anchor
         if old_edits.len() == 1 && new_edits.len() == 1 {
             let (old_range, old_text) = &old_edits[0];
             let (new_range, new_text) = &new_edits[0];
-            new_range == old_range && new_text.starts_with(old_text)
+            new_range == old_range && new_text.starts_with(old_text.as_ref())
         } else {
             true
         }
     #[gpui::test]
     async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
         let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
-        let edits: Arc<[(Range<Anchor>, String)]> = cx.update(|cx| {
+        let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
-            to_completion_edits(
+            to_completion_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
-                [(2..5, "REM".to_string()), (9..11, "".to_string())],
-                &buffer,
-                cx,
-            )
-            .into()
         });
         let edit_preview = cx
                     &buffer,
                     cx
                 ),
-                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
+                vec![(2..5, "REM".into()), (9..11, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(2..2, "REM".to_string()), (6..8, "".to_string())]
+                vec![(2..2, "REM".into()), (6..8, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.undo(cx));
                     &buffer,
                     cx
                 ),
-                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
+                vec![(2..5, "REM".into()), (9..11, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(3..3, "EM".to_string()), (7..9, "".to_string())]
+                vec![(3..3, "EM".into()), (7..9, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
+                vec![(4..4, "M".into()), (8..10, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(9..11, "".to_string())]
+                vec![(9..11, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
+                vec![(4..4, "M".into()), (8..10, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string())]
+                vec![(4..4, "M".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
     }
     fn to_completion_edits(
-        iterator: impl IntoIterator<Item = (Range<usize>, String)>,
+        iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
         buffer: &Entity<Buffer>,
         cx: &App,
-    ) -> Vec<(Range<Anchor>, String)> {
+    ) -> Vec<(Range<Anchor>, Arc<str>)> {
         let buffer = buffer.read(cx);
         iterator
             .into_iter()
     }
     fn from_completion_edits(
-        editor_edits: &[(Range<Anchor>, String)],
+        editor_edits: &[(Range<Anchor>, Arc<str>)],
         buffer: &Entity<Buffer>,
         cx: &App,
-    ) -> Vec<(Range<usize>, String)> {
+    ) -> Vec<(Range<usize>, Arc<str>)> {
         let buffer = buffer.read(cx);
         editor_edits
             .iter()

crates/zeta2/Cargo.toml+2−1

 language.workspace = true
 language_model.workspace = true
 log.workspace = true
+open_ai.workspace = true
 project.workspace = true
 release_channel.workspace = true
-schemars.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
 pretty_assertions.workspace = true
 project = { workspace = true, features = ["test-support"] }
 settings = { workspace = true, features = ["test-support"] }
+zlog.workspace = true

crates/zeta2/src/prediction.rs+41−219

-use std::{borrow::Cow, ops::Range, path::Path, sync::Arc};
+use std::{ops::Range, sync::Arc};
-use anyhow::Context as _;
+use gpui::{AsyncApp, Entity};
-use cloud_llm_client::predict_edits_v3;
+use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, TextBufferSnapshot};
-use gpui::{App, AsyncApp, Entity};
-use language::{
-    Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, TextBufferSnapshot, text_diff,
-};
-use project::Project;
-use util::ResultExt;
 use uuid::Uuid;
 #[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
-pub struct EditPredictionId(Uuid);
+pub struct EditPredictionId(pub Uuid);
 impl Into<Uuid> for EditPredictionId {
     fn into(self) -> Uuid {
 #[derive(Clone)]
 pub struct EditPrediction {
     pub id: EditPredictionId,
-    pub path: Arc<Path>,
+    pub edits: Arc<[(Range<Anchor>, Arc<str>)]>,
-    pub edits: Arc<[(Range<Anchor>, String)]>,
     pub snapshot: BufferSnapshot,
     pub edit_preview: EditPreview,
     // We keep a reference to the buffer so that we do not need to reload it from disk when applying the prediction.
 }
 impl EditPrediction {
-    pub async fn from_response(
+    pub async fn new(
-        response: predict_edits_v3::PredictEditsResponse,
+        id: EditPredictionId,
-        active_buffer_old_snapshot: &TextBufferSnapshot,
+        edited_buffer: &Entity<Buffer>,
-        active_buffer: &Entity<Buffer>,
+        edited_buffer_snapshot: &BufferSnapshot,
-        project: &Entity<Project>,
+        edits: Vec<(Range<Anchor>, Arc<str>)>,
         cx: &mut AsyncApp,
     ) -> Option<Self> {
-        // TODO only allow cloud to return one path
+        let (edits, snapshot, edit_preview_task) = edited_buffer
-        let Some(path) = response.edits.first().map(|e| e.path.clone()) else {
+            .read_with(cx, |buffer, cx| {
-            return None;
+                let new_snapshot = buffer.snapshot();
-        };
+                let edits: Arc<[_]> =
+                    interpolate_edits(&edited_buffer_snapshot, &new_snapshot, edits.into())?.into();
-        let is_same_path = active_buffer
+                Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
-            .read_with(cx, |buffer, cx| buffer_path_eq(buffer, &path, cx))
+            })
-            .ok()?;
+            .ok()??;
-        let (buffer, edits, snapshot, edit_preview_task) = if is_same_path {
-            active_buffer
-                .read_with(cx, |buffer, cx| {
-                    let new_snapshot = buffer.snapshot();
-                    let edits = edits_from_response(&response.edits, &active_buffer_old_snapshot);
-                    let edits: Arc<[_]> =
-                        interpolate_edits(active_buffer_old_snapshot, &new_snapshot, edits)?.into();
-                    Some((
-                        active_buffer.clone(),
-                        edits.clone(),
-                        new_snapshot,
-                        buffer.preview_edits(edits, cx),
-                    ))
-                })
-                .ok()??
-        } else {
-            let buffer_handle = project
-                .update(cx, |project, cx| {
-                    let project_path = project
-                        .find_project_path(&path, cx)
-                        .context("Failed to find project path for zeta edit")?;
-                    anyhow::Ok(project.open_buffer(project_path, cx))
-                })
-                .ok()?
-                .log_err()?
-                .await
-                .context("Failed to open buffer for zeta edit")
-                .log_err()?;
-            buffer_handle
-                .read_with(cx, |buffer, cx| {
-                    let snapshot = buffer.snapshot();
-                    let edits = edits_from_response(&response.edits, &snapshot);
-                    if edits.is_empty() {
-                        return None;
-                    }
-                    Some((
-                        buffer_handle.clone(),
-                        edits.clone(),
-                        snapshot,
-                        buffer.preview_edits(edits, cx),
-                    ))
-                })
-                .ok()??
-        };
         let edit_preview = edit_preview_task.await;
         Some(EditPrediction {
-            id: EditPredictionId(response.request_id),
+            id,
-            path,
             edits,
             snapshot,
             edit_preview,
-            buffer,
+            buffer: edited_buffer.clone(),
         })
     }
     pub fn interpolate(
         &self,
         new_snapshot: &TextBufferSnapshot,
-    ) -> Option<Vec<(Range<Anchor>, String)>> {
+    ) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
         interpolate_edits(&self.snapshot, new_snapshot, self.edits.clone())
     }
-    pub fn targets_buffer(&self, buffer: &Buffer, cx: &App) -> bool {
+    pub fn targets_buffer(&self, buffer: &Buffer) -> bool {
-        buffer_path_eq(buffer, &self.path, cx)
+        self.snapshot.remote_id() == buffer.remote_id()
     }
 }
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("EditPrediction")
             .field("id", &self.id)
-            .field("path", &self.path)
             .field("edits", &self.edits)
             .finish()
     }
 }
-pub fn buffer_path_eq(buffer: &Buffer, path: &Path, cx: &App) -> bool {
-    buffer.file().map(|p| p.full_path(cx)).as_deref() == Some(path)
-}
 pub fn interpolate_edits(
     old_snapshot: &TextBufferSnapshot,
     new_snapshot: &TextBufferSnapshot,
-    current_edits: Arc<[(Range<Anchor>, String)]>,
+    current_edits: Arc<[(Range<Anchor>, Arc<str>)]>,
-) -> Option<Vec<(Range<Anchor>, String)>> {
+) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
     let mut edits = Vec::new();
     let mut model_edits = current_edits.iter().peekable();
                 if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
                     if !model_suffix.is_empty() {
                         let anchor = old_snapshot.anchor_after(user_edit.old.end);
-                        edits.push((anchor..anchor, model_suffix.to_string()));
+                        edits.push((anchor..anchor, model_suffix.into()));
                     }
                     model_edits.next();
     if edits.is_empty() { None } else { Some(edits) }
 }
-pub fn line_range_to_point_range(range: Range<predict_edits_v3::Line>) -> Range<language::Point> {
-    language::Point::new(range.start.0, 0)..language::Point::new(range.end.0, 0)
-}
-fn edits_from_response(
-    edits: &[predict_edits_v3::Edit],
-    snapshot: &TextBufferSnapshot,
-) -> Arc<[(Range<Anchor>, String)]> {
-    edits
-        .iter()
-        .flat_map(|edit| {
-            let point_range = line_range_to_point_range(edit.range.clone());
-            let offset = point_range.to_offset(snapshot).start;
-            let old_text = snapshot.text_for_range(point_range);
-            excerpt_edits_from_response(
-                old_text.collect::<Cow<str>>(),
-                &edit.content,
-                offset,
-                &snapshot,
-            )
-        })
-        .collect::<Vec<_>>()
-        .into()
-}
-fn excerpt_edits_from_response(
-    old_text: Cow<str>,
-    new_text: &str,
-    offset: usize,
-    snapshot: &TextBufferSnapshot,
-) -> impl Iterator<Item = (Range<Anchor>, String)> {
-    text_diff(&old_text, new_text)
-        .into_iter()
-        .map(move |(mut old_range, new_text)| {
-            old_range.start += offset;
-            old_range.end += offset;
-            let prefix_len = common_prefix(
-                snapshot.chars_for_range(old_range.clone()),
-                new_text.chars(),
-            );
-            old_range.start += prefix_len;
-            let suffix_len = common_prefix(
-                snapshot.reversed_chars_for_range(old_range.clone()),
-                new_text[prefix_len..].chars().rev(),
-            );
-            old_range.end = old_range.end.saturating_sub(suffix_len);
-            let new_text = new_text[prefix_len..new_text.len() - suffix_len].to_string();
-            let range = if old_range.is_empty() {
-                let anchor = snapshot.anchor_after(old_range.start);
-                anchor..anchor
-            } else {
-                snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
-            };
-            (range, new_text)
-        })
-}
-fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
-    a.zip(b)
-        .take_while(|(a, b)| a == b)
-        .map(|(a, _)| a.len_utf8())
-        .sum()
-}
 #[cfg(test)]
 mod tests {
-    use std::path::PathBuf;
     use super::*;
-    use cloud_llm_client::predict_edits_v3;
-    use edit_prediction_context::Line;
     use gpui::{App, Entity, TestAppContext, prelude::*};
-    use indoc::indoc;
     use language::{Buffer, ToOffset as _};
-    #[gpui::test]
-    async fn test_compute_edits(cx: &mut TestAppContext) {
-        let old = indoc! {r#"
-            fn main() {
-                let args =
-                println!("{}", args[1])
-            }
-        "#};
-        let new = indoc! {r#"
-            fn main() {
-                let args = std::env::args();
-                println!("{}", args[1]);
-            }
-        "#};
-        let buffer = cx.new(|cx| Buffer::local(old, cx));
-        let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
-        // TODO cover more cases when multi-file is supported
-        let big_edits = vec![predict_edits_v3::Edit {
-            path: PathBuf::from("test.txt").into(),
-            range: Line(0)..Line(old.lines().count() as u32),
-            content: new.into(),
-        }];
-        let edits = edits_from_response(&big_edits, &snapshot);
-        assert_eq!(edits.len(), 2);
-        assert_eq!(
-            edits[0].0.to_point(&snapshot).start,
-            language::Point::new(1, 14)
-        );
-        assert_eq!(edits[0].1, " std::env::args();");
-        assert_eq!(
-            edits[1].0.to_point(&snapshot).start,
-            language::Point::new(2, 27)
-        );
-        assert_eq!(edits[1].1, ";");
-    }
     #[gpui::test]
     async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
         let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
-        let edits: Arc<[(Range<Anchor>, String)]> = cx.update(|cx| {
+        let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
-            to_prediction_edits(
+            to_prediction_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
-                [(2..5, "REM".to_string()), (9..11, "".to_string())],
-                &buffer,
-                cx,
-            )
-            .into()
         });
         let edit_preview = cx
             id: EditPredictionId(Uuid::new_v4()),
             edits,
             snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
-            path: Path::new("test.txt").into(),
             buffer: buffer.clone(),
             edit_preview,
         };
                     &buffer,
                     cx
                 ),
-                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
+                vec![(2..5, "REM".into()), (9..11, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(2..2, "REM".to_string()), (6..8, "".to_string())]
+                vec![(2..2, "REM".into()), (6..8, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.undo(cx));
                     &buffer,
                     cx
                 ),
-                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
+                vec![(2..5, "REM".into()), (9..11, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(3..3, "EM".to_string()), (7..9, "".to_string())]
+                vec![(3..3, "EM".into()), (7..9, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
+                vec![(4..4, "M".into()), (8..10, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(9..11, "".to_string())]
+                vec![(9..11, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
+                vec![(4..4, "M".into()), (8..10, "".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string())]
+                vec![(4..4, "M".into())]
             );
             buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
     }
     fn to_prediction_edits(
-        iterator: impl IntoIterator<Item = (Range<usize>, String)>,
+        iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
         buffer: &Entity<Buffer>,
         cx: &App,
-    ) -> Vec<(Range<Anchor>, String)> {
+    ) -> Vec<(Range<Anchor>, Arc<str>)> {
         let buffer = buffer.read(cx);
         iterator
             .into_iter()
     }
     fn from_prediction_edits(
-        editor_edits: &[(Range<Anchor>, String)],
+        editor_edits: &[(Range<Anchor>, Arc<str>)],
         buffer: &Entity<Buffer>,
         cx: &App,
-    ) -> Vec<(Range<usize>, String)> {
+    ) -> Vec<(Range<usize>, Arc<str>)> {
         let buffer = buffer.read(cx);
         editor_edits
             .iter()

crates/zeta2/src/related_excerpts.rs+0−717

@@ -1,717 +0,0 @@
1	use std::{
2	cmp::Reverse, collections::hash_map::Entry, ops::Range, path::PathBuf, sync::Arc, time::Instant,
3	};
4
5	use crate::{
6	ZetaContextRetrievalDebugInfo, ZetaContextRetrievalStartedDebugInfo, ZetaDebugInfo,
7	ZetaSearchQueryDebugInfo, merge_excerpts::merge_excerpts,
8	};
9	use anyhow::{Result, anyhow};
10	use cloud_zeta2_prompt::write_codeblock;
11	use collections::HashMap;
12	use edit_prediction_context::{EditPredictionExcerpt, EditPredictionExcerptOptions, Line};
13	use futures::{
14	StreamExt,
15	channel::mpsc::{self, UnboundedSender},
16	stream::BoxStream,
17	};
18	use gpui::{App, AppContext, AsyncApp, Entity, Task};
19	use indoc::indoc;
20	use language::{
21	Anchor, Bias, Buffer, BufferSnapshot, OffsetRangeExt, Point, TextBufferSnapshot, ToPoint as _,
22	};
23	use language_model::{
24	LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId,
25	LanguageModelProviderId, LanguageModelRegistry, LanguageModelRequest,
26	LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
27	LanguageModelToolUse, MessageContent, Role,
28	};
29	use project::{
30	Project, WorktreeSettings,
31	search::{SearchQuery, SearchResult},
32	};
33	use schemars::JsonSchema;
34	use serde::{Deserialize, Serialize};
35	use util::{
36	ResultExt as _,
37	paths::{PathMatcher, PathStyle},
38	};
39	use workspace::item::Settings as _;
40
41	const SEARCH_PROMPT: &str = indoc! {r#"
42	## Task
43
44	You are part of an edit prediction system in a code editor. Your role is to identify relevant code locations
45	that will serve as context for predicting the next required edit.
46
47	Your task:
48	- Analyze the user's recent edits and current cursor context
49	- Use the `search` tool to find code that may be relevant for predicting the next edit
50	- Focus on finding:
51	- Code patterns that might need similar changes based on the recent edits
52	- Functions, variables, types, and constants referenced in the current cursor context
53	- Related implementations, usages, or dependencies that may require consistent updates
54
55	Important constraints:
56	- This conversation has exactly 2 turns
57	- You must make ALL search queries in your first response via the `search` tool
58	- All queries will be executed in parallel and results returned together
59	- In the second turn, you will select the most relevant results via the `select` tool.
60
61	## User Edits
62
63	{edits}
64
65	## Current cursor context
66
67	`````{current_file_path}
68	{cursor_excerpt}
69	`````
70
71	--
72	Use the `search` tool now
73	"#};
74
75	const SEARCH_TOOL_NAME: &str = "search";
76
77	/// Search for relevant code
78	///
79	/// For the best results, run multiple queries at once with a single invocation of this tool.
80	#[derive(Clone, Deserialize, Serialize, JsonSchema)]
81	pub struct SearchToolInput {
82	/// An array of queries to run for gathering context relevant to the next prediction
83	#[schemars(length(max = 5))]
84	pub queries: Box<[SearchToolQuery]>,
85	}
86
87	#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
88	pub struct SearchToolQuery {
89	/// A glob pattern to match file paths in the codebase
90	pub glob: String,
91	/// A regular expression to match content within the files matched by the glob pattern
92	pub regex: String,
93	}
94
95	const RESULTS_MESSAGE: &str = indoc! {"
96	Here are the results of your queries combined and grouped by file:
97
98	"};
99
100	const SELECT_TOOL_NAME: &str = "select";
101
102	const SELECT_PROMPT: &str = indoc! {"
103	Use the `select` tool now to pick the most relevant line ranges according to the user state provided in the first message.
104	Make sure to include enough lines of context so that the edit prediction model can suggest accurate edits.
105	Include up to 200 lines in total.
106	"};
107
108	/// Select line ranges from search results
109	#[derive(Deserialize, JsonSchema)]
110	struct SelectToolInput {
111	/// The line ranges to select from search results.
112	ranges: Vec<SelectLineRange>,
113	}
114
115	/// A specific line range to select from a file
116	#[derive(Debug, Deserialize, JsonSchema)]
117	struct SelectLineRange {
118	/// The file path containing the lines to select
119	/// Exactly as it appears in the search result codeblocks.
120	path: PathBuf,
121	/// The starting line number (1-based)
122	#[schemars(range(min = 1))]
123	start_line: u32,
124	/// The ending line number (1-based, inclusive)
125	#[schemars(range(min = 1))]
126	end_line: u32,
127	}
128
129	#[derive(Debug, Clone, PartialEq)]
130	pub struct LlmContextOptions {
131	pub excerpt: EditPredictionExcerptOptions,
132	}
133
134	pub const MODEL_PROVIDER_ID: LanguageModelProviderId = language_model::ANTHROPIC_PROVIDER_ID;
135
136	pub fn find_related_excerpts(
137	buffer: Entity<language::Buffer>,
138	cursor_position: Anchor,
139	project: &Entity<Project>,
140	mut edit_history_unified_diff: String,
141	options: &LlmContextOptions,
142	debug_tx: Option<mpsc::UnboundedSender<ZetaDebugInfo>>,
143	cx: &App,
144	) -> Task<Result<HashMap<Entity<Buffer>, Vec<Range<Anchor>>>>> {
145	let language_model_registry = LanguageModelRegistry::global(cx);
146	let Some(model) = language_model_registry
147	.read(cx)
148	.available_models(cx)
149	.find(\|model\| {
150	model.provider_id() == MODEL_PROVIDER_ID
151	&& model.id() == LanguageModelId("claude-haiku-4-5-latest".into())
152	// model.provider_id() == LanguageModelProviderId::new("zeta-ctx-qwen-30b")
153	// model.provider_id() == LanguageModelProviderId::new("ollama")
154	// && model.id() == LanguageModelId("gpt-oss:20b".into())
155	})
156	else {
157	return Task::ready(Err(anyhow!("could not find context model")));
158	};
159
160	if edit_history_unified_diff.is_empty() {
161	edit_history_unified_diff.push_str("(No user edits yet)");
162	}
163
164	// TODO [zeta2] include breadcrumbs?
165	let snapshot = buffer.read(cx).snapshot();
166	let cursor_point = cursor_position.to_point(&snapshot);
167	let Some(cursor_excerpt) =
168	EditPredictionExcerpt::select_from_buffer(cursor_point, &snapshot, &options.excerpt, None)
169	else {
170	return Task::ready(Ok(HashMap::default()));
171	};
172
173	let current_file_path = snapshot
174	.file()
175	.map(\|f\| f.full_path(cx).display().to_string())
176	.unwrap_or_else(\|\| "untitled".to_string());
177
178	let prompt = SEARCH_PROMPT
179	.replace("{edits}", &edit_history_unified_diff)
180	.replace("{current_file_path}", &current_file_path)
181	.replace("{cursor_excerpt}", &cursor_excerpt.text(&snapshot).body);
182
183	if let Some(debug_tx) = &debug_tx {
184	debug_tx
185	.unbounded_send(ZetaDebugInfo::ContextRetrievalStarted(
186	ZetaContextRetrievalStartedDebugInfo {
187	project: project.clone(),
188	timestamp: Instant::now(),
189	search_prompt: prompt.clone(),
190	},
191	))
192	.ok();
193	}
194
195	let path_style = project.read(cx).path_style(cx);
196
197	let exclude_matcher = {
198	let global_settings = WorktreeSettings::get_global(cx);
199	let exclude_patterns = global_settings
200	.file_scan_exclusions
201	.sources()
202	.iter()
203	.chain(global_settings.private_files.sources().iter());
204
205	match PathMatcher::new(exclude_patterns, path_style) {
206	Ok(matcher) => matcher,
207	Err(err) => {
208	return Task::ready(Err(anyhow!(err)));
209	}
210	}
211	};
212
213	let project = project.clone();
214	cx.spawn(async move \|cx\| {
215	let initial_prompt_message = LanguageModelRequestMessage {
216	role: Role::User,
217	content: vec![prompt.into()],
218	cache: false,
219	};
220
221	let mut search_stream = request_tool_call::<SearchToolInput>(
222	vec![initial_prompt_message.clone()],
223	SEARCH_TOOL_NAME,
224	&model,
225	cx,
226	)
227	.await?;
228
229	let mut select_request_messages = Vec::with_capacity(5); // initial prompt, LLM response/thinking, tool use, tool result, select prompt
230	select_request_messages.push(initial_prompt_message);
231
232	let mut regex_by_glob: HashMap<String, String> = HashMap::default();
233	let mut search_calls = Vec::new();
234
235	while let Some(event) = search_stream.next().await {
236	match event? {
237	LanguageModelCompletionEvent::ToolUse(tool_use) => {
238	if !tool_use.is_input_complete {
239	continue;
240	}
241
242	if tool_use.name.as_ref() == SEARCH_TOOL_NAME {
243	let input =
244	serde_json::from_value::<SearchToolInput>(tool_use.input.clone())?;
245
246	for query in input.queries {
247	let regex = regex_by_glob.entry(query.glob).or_default();
248	if !regex.is_empty() {
249	regex.push('\|');
250	}
251	regex.push_str(&query.regex);
252	}
253
254	search_calls.push(tool_use);
255	} else {
256	log::warn!(
257	"context gathering model tried to use unknown tool: {}",
258	tool_use.name
259	);
260	}
261	}
262	LanguageModelCompletionEvent::Text(txt) => {
263	if let Some(LanguageModelRequestMessage {
264	role: Role::Assistant,
265	content,
266	..
267	}) = select_request_messages.last_mut()
268	{
269	if let Some(MessageContent::Text(existing_text)) = content.last_mut() {
270	existing_text.push_str(&txt);
271	} else {
272	content.push(MessageContent::Text(txt));
273	}
274	} else {
275	select_request_messages.push(LanguageModelRequestMessage {
276	role: Role::Assistant,
277	content: vec![MessageContent::Text(txt)],
278	cache: false,
279	});
280	}
281	}
282	LanguageModelCompletionEvent::Thinking { text, signature } => {
283	if let Some(LanguageModelRequestMessage {
284	role: Role::Assistant,
285	content,
286	..
287	}) = select_request_messages.last_mut()
288	{
289	if let Some(MessageContent::Thinking {
290	text: existing_text,
291	signature: existing_signature,
292	}) = content.last_mut()
293	{
294	existing_text.push_str(&text);
295	*existing_signature = signature;
296	} else {
297	content.push(MessageContent::Thinking { text, signature });
298	}
299	} else {
300	select_request_messages.push(LanguageModelRequestMessage {
301	role: Role::Assistant,
302	content: vec![MessageContent::Thinking { text, signature }],
303	cache: false,
304	});
305	}
306	}
307	LanguageModelCompletionEvent::RedactedThinking { data } => {
308	if let Some(LanguageModelRequestMessage {
309	role: Role::Assistant,
310	content,
311	..
312	}) = select_request_messages.last_mut()
313	{
314	if let Some(MessageContent::RedactedThinking(existing_data)) =
315	content.last_mut()
316	{
317	existing_data.push_str(&data);
318	} else {
319	content.push(MessageContent::RedactedThinking(data));
320	}
321	} else {
322	select_request_messages.push(LanguageModelRequestMessage {
323	role: Role::Assistant,
324	content: vec![MessageContent::RedactedThinking(data)],
325	cache: false,
326	});
327	}
328	}
329	ev @ LanguageModelCompletionEvent::ToolUseJsonParseError { .. } => {
330	log::error!("{ev:?}");
331	}
332	ev => {
333	log::trace!("context search event: {ev:?}")
334	}
335	}
336	}
337
338	let search_tool_use = if search_calls.is_empty() {
339	log::warn!("context model ran 0 searches");
340	return anyhow::Ok(Default::default());
341	} else if search_calls.len() == 1 {
342	search_calls.swap_remove(0)
343	} else {
344	// In theory, the model could perform multiple search calls
345	// Dealing with them separately is not worth it when it doesn't happen in practice.
346	// If it were to happen, here we would combine them into one.
347	// The second request doesn't need to know it was actually two different calls ;)
348	let input = serde_json::to_value(&SearchToolInput {
349	queries: regex_by_glob
350	.iter()
351	.map(\|(glob, regex)\| SearchToolQuery {
352	glob: glob.clone(),
353	regex: regex.clone(),
354	})
355	.collect(),
356	})
357	.unwrap_or_default();
358
359	LanguageModelToolUse {
360	id: search_calls.swap_remove(0).id,
361	name: SELECT_TOOL_NAME.into(),
362	raw_input: serde_json::to_string(&input).unwrap_or_default(),
363	input,
364	is_input_complete: true,
365	}
366	};
367
368	if let Some(debug_tx) = &debug_tx {
369	debug_tx
370	.unbounded_send(ZetaDebugInfo::SearchQueriesGenerated(
371	ZetaSearchQueryDebugInfo {
372	project: project.clone(),
373	timestamp: Instant::now(),
374	queries: regex_by_glob
375	.iter()
376	.map(\|(glob, regex)\| SearchToolQuery {
377	glob: glob.clone(),
378	regex: regex.clone(),
379	})
380	.collect(),
381	},
382	))
383	.ok();
384	}
385
386	let (results_tx, mut results_rx) = mpsc::unbounded();
387
388	for (glob, regex) in regex_by_glob {
389	let exclude_matcher = exclude_matcher.clone();
390	let results_tx = results_tx.clone();
391	let project = project.clone();
392	cx.spawn(async move \|cx\| {
393	run_query(
394	&glob,
395	&regex,
396	results_tx.clone(),
397	path_style,
398	exclude_matcher,
399	&project,
400	cx,
401	)
402	.await
403	.log_err();
404	})
405	.detach()
406	}
407	drop(results_tx);
408
409	struct ResultBuffer {
410	buffer: Entity<Buffer>,
411	snapshot: TextBufferSnapshot,
412	}
413
414	let (result_buffers_by_path, merged_result) = cx
415	.background_spawn(async move {
416	let mut excerpts_by_buffer: HashMap<Entity<Buffer>, MatchedBuffer> =
417	HashMap::default();
418
419	while let Some((buffer, matched)) = results_rx.next().await {
420	match excerpts_by_buffer.entry(buffer) {
421	Entry::Occupied(mut entry) => {
422	let entry = entry.get_mut();
423	entry.full_path = matched.full_path;
424	entry.snapshot = matched.snapshot;
425	entry.line_ranges.extend(matched.line_ranges);
426	}
427	Entry::Vacant(entry) => {
428	entry.insert(matched);
429	}
430	}
431	}
432
433	let mut result_buffers_by_path = HashMap::default();
434	let mut merged_result = RESULTS_MESSAGE.to_string();
435
436	for (buffer, mut matched) in excerpts_by_buffer {
437	matched
438	.line_ranges
439	.sort_unstable_by_key(\|range\| (range.start, Reverse(range.end)));
440
441	write_codeblock(
442	&matched.full_path,
443	merge_excerpts(&matched.snapshot, matched.line_ranges).iter(),
444	&[],
445	Line(matched.snapshot.max_point().row),
446	true,
447	&mut merged_result,
448	);
449
450	result_buffers_by_path.insert(
451	matched.full_path,
452	ResultBuffer {
453	buffer,
454	snapshot: matched.snapshot.text,
455	},
456	);
457	}
458
459	(result_buffers_by_path, merged_result)
460	})
461	.await;
462
463	if let Some(debug_tx) = &debug_tx {
464	debug_tx
465	.unbounded_send(ZetaDebugInfo::SearchQueriesExecuted(
466	ZetaContextRetrievalDebugInfo {
467	project: project.clone(),
468	timestamp: Instant::now(),
469	},
470	))
471	.ok();
472	}
473
474	let tool_result = LanguageModelToolResult {
475	tool_use_id: search_tool_use.id.clone(),
476	tool_name: SEARCH_TOOL_NAME.into(),
477	is_error: false,
478	content: merged_result.into(),
479	output: None,
480	};
481
482	select_request_messages.extend([
483	LanguageModelRequestMessage {
484	role: Role::Assistant,
485	content: vec![MessageContent::ToolUse(search_tool_use)],
486	cache: false,
487	},
488	LanguageModelRequestMessage {
489	role: Role::User,
490	content: vec![MessageContent::ToolResult(tool_result)],
491	cache: false,
492	},
493	]);
494
495	if result_buffers_by_path.is_empty() {
496	log::trace!("context gathering queries produced no results");
497	return anyhow::Ok(HashMap::default());
498	}
499
500	select_request_messages.push(LanguageModelRequestMessage {
501	role: Role::User,
502	content: vec![SELECT_PROMPT.into()],
503	cache: false,
504	});
505
506	let mut select_stream = request_tool_call::<SelectToolInput>(
507	select_request_messages,
508	SELECT_TOOL_NAME,
509	&model,
510	cx,
511	)
512	.await?;
513
514	cx.background_spawn(async move {
515	let mut selected_ranges = Vec::new();
516
517	while let Some(event) = select_stream.next().await {
518	match event? {
519	LanguageModelCompletionEvent::ToolUse(tool_use) => {
520	if !tool_use.is_input_complete {
521	continue;
522	}
523
524	if tool_use.name.as_ref() == SELECT_TOOL_NAME {
525	let call =
526	serde_json::from_value::<SelectToolInput>(tool_use.input.clone())?;
527	selected_ranges.extend(call.ranges);
528	} else {
529	log::warn!(
530	"context gathering model tried to use unknown tool: {}",
531	tool_use.name
532	);
533	}
534	}
535	ev @ LanguageModelCompletionEvent::ToolUseJsonParseError { .. } => {
536	log::error!("{ev:?}");
537	}
538	ev => {
539	log::trace!("context select event: {ev:?}")
540	}
541	}
542	}
543
544	if let Some(debug_tx) = &debug_tx {
545	debug_tx
546	.unbounded_send(ZetaDebugInfo::SearchResultsFiltered(
547	ZetaContextRetrievalDebugInfo {
548	project: project.clone(),
549	timestamp: Instant::now(),
550	},
551	))
552	.ok();
553	}
554
555	if selected_ranges.is_empty() {
556	log::trace!("context gathering selected no ranges")
557	}
558
559	selected_ranges.sort_unstable_by(\|a, b\| {
560	a.start_line
561	.cmp(&b.start_line)
562	.then(b.end_line.cmp(&a.end_line))
563	});
564
565	let mut related_excerpts_by_buffer: HashMap<_, Vec<_>> = HashMap::default();
566
567	for selected_range in selected_ranges {
568	if let Some(ResultBuffer { buffer, snapshot }) =
569	result_buffers_by_path.get(&selected_range.path)
570	{
571	let start_point = Point::new(selected_range.start_line.saturating_sub(1), 0);
572	let end_point =
573	snapshot.clip_point(Point::new(selected_range.end_line, 0), Bias::Left);
574	let range =
575	snapshot.anchor_after(start_point)..snapshot.anchor_before(end_point);
576
577	related_excerpts_by_buffer
578	.entry(buffer.clone())
579	.or_default()
580	.push(range);
581	} else {
582	log::warn!(
583	"selected path that wasn't included in search results: {}",
584	selected_range.path.display()
585	);
586	}
587	}
588
589	anyhow::Ok(related_excerpts_by_buffer)
590	})
591	.await
592	})
593	}
594
595	async fn request_tool_call<T: JsonSchema>(
596	messages: Vec<LanguageModelRequestMessage>,
597	tool_name: &'static str,
598	model: &Arc<dyn LanguageModel>,
599	cx: &mut AsyncApp,
600	) -> Result<BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>>
601	{
602	let schema = schemars::schema_for!(T);
603
604	let request = LanguageModelRequest {
605	messages,
606	tools: vec![LanguageModelRequestTool {
607	name: tool_name.into(),
608	description: schema
609	.get("description")
610	.and_then(\|description\| description.as_str())
611	.unwrap()
612	.to_string(),
613	input_schema: serde_json::to_value(schema).unwrap(),
614	}],
615	..Default::default()
616	};
617
618	Ok(model.stream_completion(request, cx).await?)
619	}
620
621	const MIN_EXCERPT_LEN: usize = 16;
622	const MAX_EXCERPT_LEN: usize = 768;
623	const MAX_RESULT_BYTES_PER_QUERY: usize = MAX_EXCERPT_LEN * 5;
624
625	struct MatchedBuffer {
626	snapshot: BufferSnapshot,
627	line_ranges: Vec<Range<Line>>,
628	full_path: PathBuf,
629	}
630
631	async fn run_query(
632	glob: &str,
633	regex: &str,
634	results_tx: UnboundedSender<(Entity<Buffer>, MatchedBuffer)>,
635	path_style: PathStyle,
636	exclude_matcher: PathMatcher,
637	project: &Entity<Project>,
638	cx: &mut AsyncApp,
639	) -> Result<()> {
640	let include_matcher = PathMatcher::new(vec![glob], path_style)?;
641
642	let query = SearchQuery::regex(
643	regex,
644	false,
645	true,
646	false,
647	true,
648	include_matcher,
649	exclude_matcher,
650	true,
651	None,
652	)?;
653
654	let results = project.update(cx, \|project, cx\| project.search(query, cx))?;
655	futures::pin_mut!(results);
656
657	let mut total_bytes = 0;
658
659	while let Some(SearchResult::Buffer { buffer, ranges }) = results.next().await {
660	if ranges.is_empty() {
661	continue;
662	}
663
664	let Some((snapshot, full_path)) = buffer.read_with(cx, \|buffer, cx\| {
665	Some((buffer.snapshot(), buffer.file()?.full_path(cx)))
666	})?
667	else {
668	continue;
669	};
670
671	let results_tx = results_tx.clone();
672	cx.background_spawn(async move {
673	let mut line_ranges = Vec::with_capacity(ranges.len());
674
675	for range in ranges {
676	let offset_range = range.to_offset(&snapshot);
677	let query_point = (offset_range.start + offset_range.len() / 2).to_point(&snapshot);
678
679	if total_bytes + MIN_EXCERPT_LEN >= MAX_RESULT_BYTES_PER_QUERY {
680	break;
681	}
682
683	let excerpt = EditPredictionExcerpt::select_from_buffer(
684	query_point,
685	&snapshot,
686	&EditPredictionExcerptOptions {
687	max_bytes: MAX_EXCERPT_LEN.min(MAX_RESULT_BYTES_PER_QUERY - total_bytes),
688	min_bytes: MIN_EXCERPT_LEN,
689	target_before_cursor_over_total_bytes: 0.5,
690	},
691	None,
692	);
693
694	if let Some(excerpt) = excerpt {
695	total_bytes += excerpt.range.len();
696	if !excerpt.line_range.is_empty() {
697	line_ranges.push(excerpt.line_range);
698	}
699	}
700	}
701
702	results_tx
703	.unbounded_send((
704	buffer,
705	MatchedBuffer {
706	snapshot,
707	line_ranges,
708	full_path,
709	},
710	))
711	.log_err();
712	})
713	.detach();
714	}
715
716	anyhow::Ok(())
717	}

crates/zeta2/src/retrieval_search.rs+194−0

@@ -0,0 +1,194 @@
	1	use std::ops::Range;
	2
	3	use anyhow::Result;
	4	use collections::HashMap;
	5	use edit_prediction_context::{EditPredictionExcerpt, EditPredictionExcerptOptions};
	6	use futures::{
	7	StreamExt,
	8	channel::mpsc::{self, UnboundedSender},
	9	};
	10	use gpui::{AppContext, AsyncApp, Entity};
	11	use language::{Anchor, Buffer, BufferSnapshot, OffsetRangeExt, ToPoint as _};
	12	use project::{
	13	Project, WorktreeSettings,
	14	search::{SearchQuery, SearchResult},
	15	};
	16	use util::{
	17	ResultExt as _,
	18	paths::{PathMatcher, PathStyle},
	19	};
	20	use workspace::item::Settings as _;
	21
	22	pub async fn run_retrieval_searches(
	23	project: Entity<Project>,
	24	regex_by_glob: HashMap<String, String>,
	25	cx: &mut AsyncApp,
	26	) -> Result<HashMap<Entity<Buffer>, Vec<Range<Anchor>>>> {
	27	let (exclude_matcher, path_style) = project.update(cx, \|project, cx\| {
	28	let global_settings = WorktreeSettings::get_global(cx);
	29	let exclude_patterns = global_settings
	30	.file_scan_exclusions
	31	.sources()
	32	.iter()
	33	.chain(global_settings.private_files.sources().iter());
	34	let path_style = project.path_style(cx);
	35	anyhow::Ok((PathMatcher::new(exclude_patterns, path_style)?, path_style))
	36	})??;
	37
	38	let (results_tx, mut results_rx) = mpsc::unbounded();
	39
	40	for (glob, regex) in regex_by_glob {
	41	let exclude_matcher = exclude_matcher.clone();
	42	let results_tx = results_tx.clone();
	43	let project = project.clone();
	44	cx.spawn(async move \|cx\| {
	45	run_query(
	46	&glob,
	47	&regex,
	48	results_tx.clone(),
	49	path_style,
	50	exclude_matcher,
	51	&project,
	52	cx,
	53	)
	54	.await
	55	.log_err();
	56	})
	57	.detach()
	58	}
	59	drop(results_tx);
	60
	61	cx.background_spawn(async move {
	62	let mut results: HashMap<Entity<Buffer>, Vec<Range<Anchor>>> = HashMap::default();
	63	let mut snapshots = HashMap::default();
	64
	65	let mut total_bytes = 0;
	66	'outer: while let Some((buffer, snapshot, excerpts)) = results_rx.next().await {
	67	snapshots.insert(buffer.entity_id(), snapshot);
	68	let existing = results.entry(buffer).or_default();
	69	existing.reserve(excerpts.len());
	70
	71	for (range, size) in excerpts {
	72	// Blunt trimming of the results until we have a proper algorithmic filtering step
	73	if (total_bytes + size) > MAX_RESULTS_LEN {
	74	log::trace!("Combined results reached limit of {MAX_RESULTS_LEN}B");
	75	break 'outer;
	76	}
	77	total_bytes += size;
	78	existing.push(range);
	79	}
	80	}
	81
	82	for (buffer, ranges) in results.iter_mut() {
	83	if let Some(snapshot) = snapshots.get(&buffer.entity_id()) {
	84	ranges.sort_unstable_by(\|a, b\| {
	85	a.start
	86	.cmp(&b.start, snapshot)
	87	.then(b.end.cmp(&b.end, snapshot))
	88	});
	89
	90	let mut index = 1;
	91	while index < ranges.len() {
	92	if ranges[index - 1]
	93	.end
	94	.cmp(&ranges[index].start, snapshot)
	95	.is_gt()
	96	{
	97	let removed = ranges.remove(index);
	98	ranges[index - 1].end = removed.end;
	99	} else {
	100	index += 1;
	101	}
	102	}
	103	}
	104	}
	105
	106	Ok(results)
	107	})
	108	.await
	109	}
	110
	111	const MIN_EXCERPT_LEN: usize = 16;
	112	const MAX_EXCERPT_LEN: usize = 768;
	113	const MAX_RESULTS_LEN: usize = MAX_EXCERPT_LEN * 5;
	114
	115	async fn run_query(
	116	glob: &str,
	117	regex: &str,
	118	results_tx: UnboundedSender<(Entity<Buffer>, BufferSnapshot, Vec<(Range<Anchor>, usize)>)>,
	119	path_style: PathStyle,
	120	exclude_matcher: PathMatcher,
	121	project: &Entity<Project>,
	122	cx: &mut AsyncApp,
	123	) -> Result<()> {
	124	let include_matcher = PathMatcher::new(vec![glob], path_style)?;
	125
	126	let query = SearchQuery::regex(
	127	regex,
	128	false,
	129	true,
	130	false,
	131	true,
	132	include_matcher,
	133	exclude_matcher,
	134	true,
	135	None,
	136	)?;
	137
	138	let results = project.update(cx, \|project, cx\| project.search(query, cx))?;
	139	futures::pin_mut!(results);
	140
	141	while let Some(SearchResult::Buffer { buffer, ranges }) = results.next().await {
	142	if results_tx.is_closed() {
	143	break;
	144	}
	145
	146	if ranges.is_empty() {
	147	continue;
	148	}
	149
	150	let snapshot = buffer.read_with(cx, \|buffer, _cx\| buffer.snapshot())?;
	151	let results_tx = results_tx.clone();
	152
	153	cx.background_spawn(async move {
	154	let mut excerpts = Vec::with_capacity(ranges.len());
	155
	156	for range in ranges {
	157	let offset_range = range.to_offset(&snapshot);
	158	let query_point = (offset_range.start + offset_range.len() / 2).to_point(&snapshot);
	159
	160	let excerpt = EditPredictionExcerpt::select_from_buffer(
	161	query_point,
	162	&snapshot,
	163	&EditPredictionExcerptOptions {
	164	max_bytes: MAX_EXCERPT_LEN,
	165	min_bytes: MIN_EXCERPT_LEN,
	166	target_before_cursor_over_total_bytes: 0.5,
	167	},
	168	None,
	169	);
	170
	171	if let Some(excerpt) = excerpt
	172	&& !excerpt.line_range.is_empty()
	173	{
	174	excerpts.push((
	175	snapshot.anchor_after(excerpt.range.start)
	176	..snapshot.anchor_before(excerpt.range.end),
	177	excerpt.range.len(),
	178	));
	179	}
	180	}
	181
	182	let send_result = results_tx.unbounded_send((buffer, snapshot, excerpts));
	183
	184	if let Err(err) = send_result
	185	&& !err.is_disconnected()
	186	{
	187	log::error!("{err}");
	188	}
	189	})
	190	.detach();
	191	}
	192
	193	anyhow::Ok(())
	194	}

crates/zeta2/src/udiff.rs+1024−0

@@ -0,0 +1,1024 @@
	1	use std::borrow::Cow;
	2	use std::fmt::Display;
	3	use std::sync::Arc;
	4	use std::{
	5	fmt::{Debug, Write},
	6	mem,
	7	ops::Range,
	8	path::Path,
	9	};
	10
	11	use anyhow::Context as _;
	12	use anyhow::Result;
	13	use anyhow::anyhow;
	14	use collections::HashMap;
	15	use gpui::AsyncApp;
	16	use gpui::Entity;
	17	use language::{Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, TextBufferSnapshot};
	18	use project::Project;
	19
	20	pub async fn parse_diff<'a>(
	21	diff: &'a str,
	22	get_buffer: impl Fn(&Path) -> Option<(&'a BufferSnapshot, &'a [Range<Anchor>])> + Send,
	23	) -> Result<(&'a BufferSnapshot, Vec<(Range<Anchor>, Arc<str>)>)> {
	24	let mut diff = DiffParser::new(diff);
	25	let mut edited_buffer = None;
	26	let mut edits = Vec::new();
	27
	28	while let Some(event) = diff.next()? {
	29	match event {
	30	DiffEvent::Hunk {
	31	path: file_path,
	32	hunk,
	33	} => {
	34	let (buffer, ranges) = match edited_buffer {
	35	None => {
	36	edited_buffer = get_buffer(&Path::new(file_path.as_ref()));
	37	edited_buffer
	38	.as_ref()
	39	.context("Model tried to edit a file that wasn't included")?
	40	}
	41	Some(ref current) => current,
	42	};
	43
	44	edits.extend(resolve_hunk_edits_in_buffer(hunk, &buffer.text, ranges)?);
	45	}
	46	DiffEvent::FileEnd { renamed_to } => {
	47	let (buffer, _) = edited_buffer
	48	.take()
	49	.expect("Got a FileEnd event before an Hunk event");
	50
	51	if renamed_to.is_some() {
	52	anyhow::bail!("edit predictions cannot rename files");
	53	}
	54
	55	if diff.next()?.is_some() {
	56	anyhow::bail!("Edited more than one file");
	57	}
	58
	59	return Ok((buffer, edits));
	60	}
	61	}
	62	}
	63
	64	Err(anyhow::anyhow!("No EOF"))
	65	}
	66
	67	#[derive(Debug)]
	68	pub struct OpenedBuffers<'a>(#[allow(unused)] HashMap<Cow<'a, str>, Entity<Buffer>>);
	69
	70	#[must_use]
	71	pub async fn apply_diff<'a>(
	72	diff: &'a str,
	73	project: &Entity<Project>,
	74	cx: &mut AsyncApp,
	75	) -> Result<OpenedBuffers<'a>> {
	76	let mut included_files = HashMap::default();
	77
	78	for line in diff.lines() {
	79	let diff_line = DiffLine::parse(line);
	80
	81	if let DiffLine::OldPath { path } = diff_line {
	82	let buffer = project
	83	.update(cx, \|project, cx\| {
	84	let project_path =
	85	project
	86	.find_project_path(path.as_ref(), cx)
	87	.with_context(\|\| {
	88	format!("Failed to find worktree for new path: {}", path)
	89	})?;
	90	anyhow::Ok(project.open_buffer(project_path, cx))
	91	})??
	92	.await?;
	93
	94	included_files.insert(path, buffer);
	95	}
	96	}
	97
	98	let ranges = [Anchor::MIN..Anchor::MAX];
	99
	100	let mut diff = DiffParser::new(diff);
	101	let mut current_file = None;
	102	let mut edits = vec![];
	103
	104	while let Some(event) = diff.next()? {
	105	match event {
	106	DiffEvent::Hunk {
	107	path: file_path,
	108	hunk,
	109	} => {
	110	let (buffer, ranges) = match current_file {
	111	None => {
	112	let buffer = included_files
	113	.get_mut(&file_path)
	114	.expect("Opened all files in diff");
	115
	116	current_file = Some((buffer, ranges.as_slice()));
	117	current_file.as_ref().unwrap()
	118	}
	119	Some(ref current) => current,
	120	};
	121
	122	buffer.read_with(cx, \|buffer, _\| {
	123	edits.extend(resolve_hunk_edits_in_buffer(hunk, buffer, ranges)?);
	124	anyhow::Ok(())
	125	})??;
	126	}
	127	DiffEvent::FileEnd { renamed_to } => {
	128	let (buffer, _) = current_file
	129	.take()
	130	.expect("Got a FileEnd event before an Hunk event");
	131
	132	if let Some(renamed_to) = renamed_to {
	133	project
	134	.update(cx, \|project, cx\| {
	135	let new_project_path = project
	136	.find_project_path(Path::new(renamed_to.as_ref()), cx)
	137	.with_context(\|\| {
	138	format!("Failed to find worktree for new path: {}", renamed_to)
	139	})?;
	140
	141	let project_file = project::File::from_dyn(buffer.read(cx).file())
	142	.expect("Wrong file type");
	143
	144	anyhow::Ok(project.rename_entry(
	145	project_file.entry_id.unwrap(),
	146	new_project_path,
	147	cx,
	148	))
	149	})??
	150	.await?;
	151	}
	152
	153	let edits = mem::take(&mut edits);
	154	buffer.update(cx, \|buffer, cx\| {
	155	buffer.edit(edits, None, cx);
	156	})?;
	157	}
	158	}
	159	}
	160
	161	Ok(OpenedBuffers(included_files))
	162	}
	163
	164	struct PatchFile<'a> {
	165	old_path: Cow<'a, str>,
	166	new_path: Cow<'a, str>,
	167	}
	168
	169	struct DiffParser<'a> {
	170	current_file: Option<PatchFile<'a>>,
	171	current_line: Option<(&'a str, DiffLine<'a>)>,
	172	hunk: Hunk,
	173	diff: std::str::Lines<'a>,
	174	}
	175
	176	#[derive(Debug, PartialEq)]
	177	enum DiffEvent<'a> {
	178	Hunk { path: Cow<'a, str>, hunk: Hunk },
	179	FileEnd { renamed_to: Option<Cow<'a, str>> },
	180	}
	181
	182	#[derive(Debug, Default, PartialEq)]
	183	struct Hunk {
	184	context: String,
	185	edits: Vec<Edit>,
	186	}
	187
	188	impl Hunk {
	189	fn is_empty(&self) -> bool {
	190	self.context.is_empty() && self.edits.is_empty()
	191	}
	192	}
	193
	194	#[derive(Debug, PartialEq)]
	195	struct Edit {
	196	range: Range<usize>,
	197	text: String,
	198	}
	199
	200	impl<'a> DiffParser<'a> {
	201	fn new(diff: &'a str) -> Self {
	202	let mut diff = diff.lines();
	203	let current_line = diff.next().map(\|line\| (line, DiffLine::parse(line)));
	204	DiffParser {
	205	current_file: None,
	206	hunk: Hunk::default(),
	207	current_line,
	208	diff,
	209	}
	210	}
	211
	212	fn next(&mut self) -> Result<Option<DiffEvent<'a>>> {
	213	loop {
	214	let (hunk_done, file_done) = match self.current_line.as_ref().map(\|e\| &e.1) {
	215	Some(DiffLine::OldPath { .. }) \| Some(DiffLine::Garbage(_)) \| None => (true, true),
	216	Some(DiffLine::HunkHeader(_)) => (true, false),
	217	_ => (false, false),
	218	};
	219
	220	if hunk_done {
	221	if let Some(file) = &self.current_file
	222	&& !self.hunk.is_empty()
	223	{
	224	return Ok(Some(DiffEvent::Hunk {
	225	path: file.old_path.clone(),
	226	hunk: mem::take(&mut self.hunk),
	227	}));
	228	}
	229	}
	230
	231	if file_done {
	232	if let Some(PatchFile { old_path, new_path }) = self.current_file.take() {
	233	return Ok(Some(DiffEvent::FileEnd {
	234	renamed_to: if old_path != new_path {
	235	Some(new_path)
	236	} else {
	237	None
	238	},
	239	}));
	240	}
	241	}
	242
	243	let Some((line, parsed_line)) = self.current_line.take() else {
	244	break;
	245	};
	246
	247	util::maybe!({
	248	match parsed_line {
	249	DiffLine::OldPath { path } => {
	250	self.current_file = Some(PatchFile {
	251	old_path: path,
	252	new_path: "".into(),
	253	});
	254	}
	255	DiffLine::NewPath { path } => {
	256	if let Some(current_file) = &mut self.current_file {
	257	current_file.new_path = path
	258	}
	259	}
	260	DiffLine::HunkHeader(_) => {}
	261	DiffLine::Context(ctx) => {
	262	if self.current_file.is_some() {
	263	writeln!(&mut self.hunk.context, "{ctx}")?;
	264	}
	265	}
	266	DiffLine::Deletion(del) => {
	267	if self.current_file.is_some() {
	268	let range = self.hunk.context.len()
	269	..self.hunk.context.len() + del.len() + '\n'.len_utf8();
	270	if let Some(last_edit) = self.hunk.edits.last_mut()
	271	&& last_edit.range.end == range.start
	272	{
	273	last_edit.range.end = range.end;
	274	} else {
	275	self.hunk.edits.push(Edit {
	276	range,
	277	text: String::new(),
	278	});
	279	}
	280	writeln!(&mut self.hunk.context, "{del}")?;
	281	}
	282	}
	283	DiffLine::Addition(add) => {
	284	if self.current_file.is_some() {
	285	let range = self.hunk.context.len()..self.hunk.context.len();
	286	if let Some(last_edit) = self.hunk.edits.last_mut()
	287	&& last_edit.range.end == range.start
	288	{
	289	writeln!(&mut last_edit.text, "{add}").unwrap();
	290	} else {
	291	self.hunk.edits.push(Edit {
	292	range,
	293	text: format!("{add}\n"),
	294	});
	295	}
	296	}
	297	}
	298	DiffLine::Garbage(_) => {}
	299	}
	300
	301	anyhow::Ok(())
	302	})
	303	.with_context(\|\| format!("on line:\n\n```\n{}```", line))?;
	304
	305	self.current_line = self.diff.next().map(\|line\| (line, DiffLine::parse(line)));
	306	}
	307
	308	anyhow::Ok(None)
	309	}
	310	}
	311
	312	fn resolve_hunk_edits_in_buffer(
	313	hunk: Hunk,
	314	buffer: &TextBufferSnapshot,
	315	ranges: &[Range<Anchor>],
	316	) -> Result<impl Iterator<Item = (Range<Anchor>, Arc<str>)>, anyhow::Error> {
	317	let context_offset = if hunk.context.is_empty() {
	318	Ok(0)
	319	} else {
	320	let mut offset = None;
	321	for range in ranges {
	322	let range = range.to_offset(buffer);
	323	let text = buffer.text_for_range(range.clone()).collect::<String>();
	324	for (ix, _) in text.match_indices(&hunk.context) {
	325	if offset.is_some() {
	326	anyhow::bail!("Context is not unique enough:\n{}", hunk.context);
	327	}
	328	offset = Some(range.start + ix);
	329	}
	330	}
	331	offset.ok_or_else(\|\| {
	332	anyhow!(
	333	"Failed to match context:\n{}\n\nBuffer:\n{}",
	334	hunk.context,
	335	buffer.text(),
	336	)
	337	})
	338	}?;
	339	let iter = hunk.edits.into_iter().flat_map(move \|edit\| {
	340	let old_text = buffer
	341	.text_for_range(context_offset + edit.range.start..context_offset + edit.range.end)
	342	.collect::<String>();
	343	let edits_within_hunk = language::text_diff(&old_text, &edit.text);
	344	edits_within_hunk
	345	.into_iter()
	346	.map(move \|(inner_range, inner_text)\| {
	347	(
	348	buffer.anchor_after(context_offset + edit.range.start + inner_range.start)
	349	..buffer.anchor_before(context_offset + edit.range.start + inner_range.end),
	350	inner_text,
	351	)
	352	})
	353	});
	354	Ok(iter)
	355	}
	356
	357	#[derive(Debug, PartialEq)]
	358	pub enum DiffLine<'a> {
	359	OldPath { path: Cow<'a, str> },
	360	NewPath { path: Cow<'a, str> },
	361	HunkHeader(Option<HunkLocation>),
	362	Context(&'a str),
	363	Deletion(&'a str),
	364	Addition(&'a str),
	365	Garbage(&'a str),
	366	}
	367
	368	#[derive(Debug, PartialEq)]
	369	pub struct HunkLocation {
	370	start_line_old: u32,
	371	count_old: u32,
	372	start_line_new: u32,
	373	count_new: u32,
	374	}
	375
	376	impl<'a> DiffLine<'a> {
	377	pub fn parse(line: &'a str) -> Self {
	378	Self::try_parse(line).unwrap_or(Self::Garbage(line))
	379	}
	380
	381	fn try_parse(line: &'a str) -> Option<Self> {
	382	if let Some(header) = line.strip_prefix("---").and_then(eat_required_whitespace) {
	383	let path = parse_header_path("a/", header);
	384	Some(Self::OldPath { path })
	385	} else if let Some(header) = line.strip_prefix("+++").and_then(eat_required_whitespace) {
	386	Some(Self::NewPath {
	387	path: parse_header_path("b/", header),
	388	})
	389	} else if let Some(header) = line.strip_prefix("@@").and_then(eat_required_whitespace) {
	390	if header.starts_with("...") {
	391	return Some(Self::HunkHeader(None));
	392	}
	393
	394	let (start_line_old, header) = header.strip_prefix('-')?.split_once(',')?;
	395	let mut parts = header.split_ascii_whitespace();
	396	let count_old = parts.next()?;
	397	let (start_line_new, count_new) = parts.next()?.strip_prefix('+')?.split_once(',')?;
	398
	399	Some(Self::HunkHeader(Some(HunkLocation {
	400	start_line_old: start_line_old.parse::<u32>().ok()?.saturating_sub(1),
	401	count_old: count_old.parse().ok()?,
	402	start_line_new: start_line_new.parse::<u32>().ok()?.saturating_sub(1),
	403	count_new: count_new.parse().ok()?,
	404	})))
	405	} else if let Some(deleted_header) = line.strip_prefix("-") {
	406	Some(Self::Deletion(deleted_header))
	407	} else if line.is_empty() {
	408	Some(Self::Context(""))
	409	} else if let Some(context) = line.strip_prefix(" ") {
	410	Some(Self::Context(context))
	411	} else {
	412	Some(Self::Addition(line.strip_prefix("+")?))
	413	}
	414	}
	415	}
	416
	417	impl<'a> Display for DiffLine<'a> {
	418	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
	419	match self {
	420	DiffLine::OldPath { path } => write!(f, "--- {path}"),
	421	DiffLine::NewPath { path } => write!(f, "+++ {path}"),
	422	DiffLine::HunkHeader(Some(hunk_location)) => {
	423	write!(
	424	f,
	425	"@@ -{},{} +{},{} @@",
	426	hunk_location.start_line_old + 1,
	427	hunk_location.count_old,
	428	hunk_location.start_line_new + 1,
	429	hunk_location.count_new
	430	)
	431	}
	432	DiffLine::HunkHeader(None) => write!(f, "@@ ... @@"),
	433	DiffLine::Context(content) => write!(f, " {content}"),
	434	DiffLine::Deletion(content) => write!(f, "-{content}"),
	435	DiffLine::Addition(content) => write!(f, "+{content}"),
	436	DiffLine::Garbage(line) => write!(f, "{line}"),
	437	}
	438	}
	439	}
	440
	441	fn parse_header_path<'a>(strip_prefix: &'static str, header: &'a str) -> Cow<'a, str> {
	442	if !header.contains(['"', '\\']) {
	443	let path = header.split_ascii_whitespace().next().unwrap_or(header);
	444	return Cow::Borrowed(path.strip_prefix(strip_prefix).unwrap_or(path));
	445	}
	446
	447	let mut path = String::with_capacity(header.len());
	448	let mut in_quote = false;
	449	let mut chars = header.chars().peekable();
	450	let mut strip_prefix = Some(strip_prefix);
	451
	452	while let Some(char) = chars.next() {
	453	if char == '"' {
	454	in_quote = !in_quote;
	455	} else if char == '\\' {
	456	let Some(&next_char) = chars.peek() else {
	457	break;
	458	};
	459	chars.next();
	460	path.push(next_char);
	461	} else if char.is_ascii_whitespace() && !in_quote {
	462	break;
	463	} else {
	464	path.push(char);
	465	}
	466
	467	if let Some(prefix) = strip_prefix
	468	&& path == prefix
	469	{
	470	strip_prefix.take();
	471	path.clear();
	472	}
	473	}
	474
	475	Cow::Owned(path)
	476	}
	477
	478	fn eat_required_whitespace(header: &str) -> Option<&str> {
	479	let trimmed = header.trim_ascii_start();
	480
	481	if trimmed.len() == header.len() {
	482	None
	483	} else {
	484	Some(trimmed)
	485	}
	486	}
	487
	488	#[cfg(test)]
	489	mod tests {
	490	use super::*;
	491	use gpui::TestAppContext;
	492	use indoc::indoc;
	493	use language::Point;
	494	use pretty_assertions::assert_eq;
	495	use project::{FakeFs, Project};
	496	use serde_json::json;
	497	use settings::SettingsStore;
	498	use util::path;
	499
	500	#[test]
	501	fn parse_lines_simple() {
	502	let input = indoc! {"
	503	diff --git a/text.txt b/text.txt
	504	index 86c770d..a1fd855 100644
	505	--- a/file.txt
	506	+++ b/file.txt
	507	@@ -1,2 +1,3 @@
	508	context
	509	-deleted
	510	+inserted
	511	garbage
	512
	513	--- b/file.txt
	514	+++ a/file.txt
	515	"};
	516
	517	let lines = input.lines().map(DiffLine::parse).collect::<Vec<_>>();
	518
	519	pretty_assertions::assert_eq!(
	520	lines,
	521	&[
	522	DiffLine::Garbage("diff --git a/text.txt b/text.txt"),
	523	DiffLine::Garbage("index 86c770d..a1fd855 100644"),
	524	DiffLine::OldPath {
	525	path: "file.txt".into()
	526	},
	527	DiffLine::NewPath {
	528	path: "file.txt".into()
	529	},
	530	DiffLine::HunkHeader(Some(HunkLocation {
	531	start_line_old: 0,
	532	count_old: 2,
	533	start_line_new: 0,
	534	count_new: 3
	535	})),
	536	DiffLine::Context("context"),
	537	DiffLine::Deletion("deleted"),
	538	DiffLine::Addition("inserted"),
	539	DiffLine::Garbage("garbage"),
	540	DiffLine::Context(""),
	541	DiffLine::OldPath {
	542	path: "b/file.txt".into()
	543	},
	544	DiffLine::NewPath {
	545	path: "a/file.txt".into()
	546	},
	547	]
	548	);
	549	}
	550
	551	#[test]
	552	fn file_header_extra_space() {
	553	let options = ["--- file", "--- file", "---\tfile"];
	554
	555	for option in options {
	556	pretty_assertions::assert_eq!(
	557	DiffLine::parse(option),
	558	DiffLine::OldPath {
	559	path: "file".into()
	560	},
	561	"{option}",
	562	);
	563	}
	564	}
	565
	566	#[test]
	567	fn hunk_header_extra_space() {
	568	let options = [
	569	"@@ -1,2 +1,3 @@",
	570	"@@ -1,2 +1,3 @@",
	571	"@@\t-1,2\t+1,3\t@@",
	572	"@@ -1,2 +1,3 @@",
	573	"@@ -1,2 +1,3 @@",
	574	"@@ -1,2 +1,3 @@",
	575	"@@ -1,2 +1,3 @@ garbage",
	576	];
	577
	578	for option in options {
	579	pretty_assertions::assert_eq!(
	580	DiffLine::parse(option),
	581	DiffLine::HunkHeader(Some(HunkLocation {
	582	start_line_old: 0,
	583	count_old: 2,
	584	start_line_new: 0,
	585	count_new: 3
	586	})),
	587	"{option}",
	588	);
	589	}
	590	}
	591
	592	#[test]
	593	fn hunk_header_without_location() {
	594	pretty_assertions::assert_eq!(DiffLine::parse("@@ ... @@"), DiffLine::HunkHeader(None));
	595	}
	596
	597	#[test]
	598	fn test_parse_path() {
	599	assert_eq!(parse_header_path("a/", "foo.txt"), "foo.txt");
	600	assert_eq!(
	601	parse_header_path("a/", "foo/bar/baz.txt"),
	602	"foo/bar/baz.txt"
	603	);
	604	assert_eq!(parse_header_path("a/", "a/foo.txt"), "foo.txt");
	605	assert_eq!(
	606	parse_header_path("a/", "a/foo/bar/baz.txt"),
	607	"foo/bar/baz.txt"
	608	);
	609
	610	// Extra
	611	assert_eq!(
	612	parse_header_path("a/", "a/foo/bar/baz.txt 2025"),
	613	"foo/bar/baz.txt"
	614	);
	615	assert_eq!(
	616	parse_header_path("a/", "a/foo/bar/baz.txt\t2025"),
	617	"foo/bar/baz.txt"
	618	);
	619	assert_eq!(
	620	parse_header_path("a/", "a/foo/bar/baz.txt \""),
	621	"foo/bar/baz.txt"
	622	);
	623
	624	// Quoted
	625	assert_eq!(
	626	parse_header_path("a/", "a/foo/bar/\"baz quox.txt\""),
	627	"foo/bar/baz quox.txt"
	628	);
	629	assert_eq!(
	630	parse_header_path("a/", "\"a/foo/bar/baz quox.txt\""),
	631	"foo/bar/baz quox.txt"
	632	);
	633	assert_eq!(
	634	parse_header_path("a/", "\"foo/bar/baz quox.txt\""),
	635	"foo/bar/baz quox.txt"
	636	);
	637	assert_eq!(parse_header_path("a/", "\"whatever 🤷\""), "whatever 🤷");
	638	assert_eq!(
	639	parse_header_path("a/", "\"foo/bar/baz quox.txt\" 2025"),
	640	"foo/bar/baz quox.txt"
	641	);
	642	// unescaped quotes are dropped
	643	assert_eq!(parse_header_path("a/", "foo/\"bar\""), "foo/bar");
	644
	645	// Escaped
	646	assert_eq!(
	647	parse_header_path("a/", "\"foo/\\\"bar\\\"/baz.txt\""),
	648	"foo/\"bar\"/baz.txt"
	649	);
	650	assert_eq!(
	651	parse_header_path("a/", "\"C:\\\\Projects\\\\My App\\\\old file.txt\""),
	652	"C:\\Projects\\My App\\old file.txt"
	653	);
	654	}
	655
	656	#[test]
	657	fn test_parse_diff_with_leading_and_trailing_garbage() {
	658	let diff = indoc! {"
	659	I need to make some changes.
	660
	661	I'll change the following things:
	662	- one
	663	- two
	664	- three
	665
	666	```
	667	--- a/file.txt
	668	+++ b/file.txt
	669	one
	670	+AND
	671	two
	672	```
	673
	674	Summary of what I did:
	675	- one
	676	- two
	677	- three
	678
	679	That's about it.
	680	"};
	681
	682	let mut events = Vec::new();
	683	let mut parser = DiffParser::new(diff);
	684	while let Some(event) = parser.next().unwrap() {
	685	events.push(event);
	686	}
	687
	688	assert_eq!(
	689	events,
	690	&[
	691	DiffEvent::Hunk {
	692	path: "file.txt".into(),
	693	hunk: Hunk {
	694	context: "one\ntwo\n".into(),
	695	edits: vec![Edit {
	696	range: 4..4,
	697	text: "AND\n".into()
	698	}],
	699	}
	700	},
	701	DiffEvent::FileEnd { renamed_to: None }
	702	],
	703	)
	704	}
	705
	706	#[gpui::test]
	707	async fn test_apply_diff_successful(cx: &mut TestAppContext) {
	708	let fs = init_test(cx);
	709
	710	let buffer_1_text = indoc! {r#"
	711	one
	712	two
	713	three
	714	four
	715	five
	716	"# };
	717
	718	let buffer_1_text_final = indoc! {r#"
	719	3
	720	4
	721	5
	722	"# };
	723
	724	let buffer_2_text = indoc! {r#"
	725	six
	726	seven
	727	eight
	728	nine
	729	ten
	730	"# };
	731
	732	let buffer_2_text_final = indoc! {r#"
	733	5
	734	six
	735	seven
	736	7.5
	737	eight
	738	nine
	739	ten
	740	11
	741	"# };
	742
	743	fs.insert_tree(
	744	path!("/root"),
	745	json!({
	746	"file1": buffer_1_text,
	747	"file2": buffer_2_text,
	748	}),
	749	)
	750	.await;
	751
	752	let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
	753
	754	let diff = indoc! {r#"
	755	--- a/root/file1
	756	+++ b/root/file1
	757	one
	758	two
	759	-three
	760	+3
	761	four
	762	five
	763	--- a/root/file1
	764	+++ b/root/file1
	765	3
	766	-four
	767	-five
	768	+4
	769	+5
	770	--- a/root/file1
	771	+++ b/root/file1
	772	-one
	773	-two
	774	3
	775	4
	776	--- a/root/file2
	777	+++ b/root/file2
	778	+5
	779	six
	780	--- a/root/file2
	781	+++ b/root/file2
	782	seven
	783	+7.5
	784	eight
	785	--- a/root/file2
	786	+++ b/root/file2
	787	ten
	788	+11
	789	"#};
	790
	791	let _buffers = apply_diff(diff, &project, &mut cx.to_async())
	792	.await
	793	.unwrap();
	794	let buffer_1 = project
	795	.update(cx, \|project, cx\| {
	796	let project_path = project.find_project_path(path!("/root/file1"), cx).unwrap();
	797	project.open_buffer(project_path, cx)
	798	})
	799	.await
	800	.unwrap();
	801
	802	buffer_1.read_with(cx, \|buffer, _cx\| {
	803	assert_eq!(buffer.text(), buffer_1_text_final);
	804	});
	805	let buffer_2 = project
	806	.update(cx, \|project, cx\| {
	807	let project_path = project.find_project_path(path!("/root/file2"), cx).unwrap();
	808	project.open_buffer(project_path, cx)
	809	})
	810	.await
	811	.unwrap();
	812
	813	buffer_2.read_with(cx, \|buffer, _cx\| {
	814	assert_eq!(buffer.text(), buffer_2_text_final);
	815	});
	816	}
	817
	818	#[gpui::test]
	819	async fn test_apply_diff_non_unique(cx: &mut TestAppContext) {
	820	let fs = init_test(cx);
	821
	822	let buffer_1_text = indoc! {r#"
	823	one
	824	two
	825	three
	826	four
	827	five
	828	one
	829	two
	830	three
	831	four
	832	five
	833	"# };
	834
	835	fs.insert_tree(
	836	path!("/root"),
	837	json!({
	838	"file1": buffer_1_text,
	839	}),
	840	)
	841	.await;
	842
	843	let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
	844	let buffer = project
	845	.update(cx, \|project, cx\| {
	846	project.open_local_buffer(path!("/root/file1"), cx)
	847	})
	848	.await
	849	.unwrap();
	850	let buffer_snapshot = buffer.read_with(cx, \|buffer, _\| buffer.snapshot());
	851
	852	let diff = indoc! {r#"
	853	--- a/root/file1
	854	+++ b/root/file1
	855	one
	856	two
	857	-three
	858	+3
	859	four
	860	five
	861	"#};
	862
	863	let final_text = indoc! {r#"
	864	one
	865	two
	866	three
	867	four
	868	five
	869	one
	870	two
	871	3
	872	four
	873	five
	874	"#};
	875
	876	apply_diff(diff, &project, &mut cx.to_async())
	877	.await
	878	.expect_err("Non-unique edits should fail");
	879
	880	let ranges = [buffer_snapshot.anchor_before(Point::new(1, 0))
	881	..buffer_snapshot.anchor_after(buffer_snapshot.max_point())];
	882
	883	let (edited_snapshot, edits) = parse_diff(diff, \|_path\| Some((&buffer_snapshot, &ranges)))
	884	.await
	885	.unwrap();
	886
	887	assert_eq!(edited_snapshot.remote_id(), buffer_snapshot.remote_id());
	888	buffer.update(cx, \|buffer, cx\| {
	889	buffer.edit(edits, None, cx);
	890	assert_eq!(buffer.text(), final_text);
	891	});
	892	}
	893
	894	#[gpui::test]
	895	async fn test_parse_diff_with_edits_within_line(cx: &mut TestAppContext) {
	896	let fs = init_test(cx);
	897
	898	let buffer_1_text = indoc! {r#"
	899	one two three four
	900	five six seven eight
	901	nine ten eleven twelve
	902	"# };
	903
	904	fs.insert_tree(
	905	path!("/root"),
	906	json!({
	907	"file1": buffer_1_text,
	908	}),
	909	)
	910	.await;
	911
	912	let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
	913	let buffer = project
	914	.update(cx, \|project, cx\| {
	915	project.open_local_buffer(path!("/root/file1"), cx)
	916	})
	917	.await
	918	.unwrap();
	919	let buffer_snapshot = buffer.read_with(cx, \|buffer, _\| buffer.snapshot());
	920
	921	let diff = indoc! {r#"
	922	--- a/root/file1
	923	+++ b/root/file1
	924	one two three four
	925	-five six seven eight
	926	+five SIX seven eight!
	927	nine ten eleven twelve
	928	"#};
	929
	930	let (buffer, edits) = parse_diff(diff, \|_path\| {
	931	Some((&buffer_snapshot, &[(Anchor::MIN..Anchor::MAX)] as &[_]))
	932	})
	933	.await
	934	.unwrap();
	935
	936	let edits = edits
	937	.into_iter()
	938	.map(\|(range, text)\| (range.to_point(&buffer), text))
	939	.collect::<Vec<_>>();
	940	assert_eq!(
	941	edits,
	942	&[
	943	(Point::new(1, 5)..Point::new(1, 8), "SIX".into()),
	944	(Point::new(1, 20)..Point::new(1, 20), "!".into())
	945	]
	946	);
	947	}
	948
	949	#[gpui::test]
	950	async fn test_apply_diff_unique_via_previous_context(cx: &mut TestAppContext) {
	951	let fs = init_test(cx);
	952
	953	let start = indoc! {r#"
	954	one
	955	two
	956	three
	957	four
	958	five
	959
	960	four
	961	five
	962	"# };
	963
	964	let end = indoc! {r#"
	965	one
	966	two
	967	3
	968	four
	969	5
	970
	971	four
	972	five
	973	"# };
	974
	975	fs.insert_tree(
	976	path!("/root"),
	977	json!({
	978	"file1": start,
	979	}),
	980	)
	981	.await;
	982
	983	let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
	984
	985	let diff = indoc! {r#"
	986	--- a/root/file1
	987	+++ b/root/file1
	988	one
	989	two
	990	-three
	991	+3
	992	four
	993	-five
	994	+5
	995	"#};
	996
	997	let _buffers = apply_diff(diff, &project, &mut cx.to_async())
	998	.await
	999	.unwrap();
	1000
	1001	let buffer_1 = project
	1002	.update(cx, \|project, cx\| {
	1003	let project_path = project.find_project_path(path!("/root/file1"), cx).unwrap();
	1004	project.open_buffer(project_path, cx)
	1005	})
	1006	.await
	1007	.unwrap();
	1008
	1009	buffer_1.read_with(cx, \|buffer, _cx\| {
	1010	assert_eq!(buffer.text(), end);
	1011	});
	1012	}
	1013
	1014	fn init_test(cx: &mut TestAppContext) -> Arc<FakeFs> {
	1015	cx.update(\|cx\| {
	1016	let settings_store = SettingsStore::test(cx);
	1017	cx.set_global(settings_store);
	1018	Project::init_settings(cx);
	1019	language::init(cx);
	1020	});
	1021
	1022	FakeFs::new(cx.background_executor.clone())
	1023	}
	1024	}

crates/zeta2/src/zeta2.rs+557−269

@@ -6,7 +6,8 @@ use cloud_llm_client::{ · 5 unmodified lines
6	AcceptEditPredictionBody, EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME,	6	AcceptEditPredictionBody, EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
7	ZED_VERSION_HEADER_NAME,	7	ZED_VERSION_HEADER_NAME,
8	};	8	};
9	use cloud_zeta2_prompt::{DEFAULT_MAX_PROMPT_BYTES, build_prompt};	9	use cloud_zeta2_prompt::DEFAULT_MAX_PROMPT_BYTES;
		10	use cloud_zeta2_prompt::retrieval_prompt::SearchToolInput;
10	use collections::HashMap;	11	use collections::HashMap;
11	use edit_prediction_context::{	12	use edit_prediction_context::{
12	DeclarationId, DeclarationStyle, EditPredictionContext, EditPredictionContextOptions,	13	DeclarationId, DeclarationStyle, EditPredictionContext, EditPredictionContextOptions,
@@ -24,11 +25,13 @@ use gpui::{ · 11 unmodified lines
24	use language::{Anchor, Buffer, DiagnosticSet, LanguageServerId, ToOffset as _, ToPoint};	25	use language::{Anchor, Buffer, DiagnosticSet, LanguageServerId, ToOffset as _, ToPoint};
25	use language::{BufferSnapshot, OffsetRangeExt};	26	use language::{BufferSnapshot, OffsetRangeExt};
26	use language_model::{LlmApiToken, RefreshLlmTokenListener};	27	use language_model::{LlmApiToken, RefreshLlmTokenListener};
		28	use open_ai::FunctionDefinition;
27	use project::Project;	29	use project::Project;
28	use release_channel::AppVersion;	30	use release_channel::AppVersion;
29	use serde::de::DeserializeOwned;	31	use serde::de::DeserializeOwned;
30	use std::collections::{VecDeque, hash_map};	32	use std::collections::{VecDeque, hash_map};
31	use std::fmt::Write;	33	use uuid::Uuid;
		34
32	use std::ops::Range;	35	use std::ops::Range;
33	use std::path::Path;	36	use std::path::Path;
34	use std::str::FromStr as _;	37	use std::str::FromStr as _;
@@ -42,12 +45,12 @@ use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_noti · 7 unmodified lines
42	pub mod merge_excerpts;	45	pub mod merge_excerpts;
43	mod prediction;	46	mod prediction;
44	mod provider;	47	mod provider;
45	pub mod related_excerpts;	48	pub mod retrieval_search;
		49	pub mod udiff;
46		50
47	use crate::merge_excerpts::merge_excerpts;	51	use crate::merge_excerpts::merge_excerpts;
48	use crate::prediction::EditPrediction;	52	use crate::prediction::EditPrediction;
49	use crate::related_excerpts::find_related_excerpts;	53	pub use crate::prediction::EditPredictionId;
50	pub use crate::related_excerpts::{LlmContextOptions, SearchToolQuery};
51	pub use provider::ZetaEditPredictionProvider;	54	pub use provider::ZetaEditPredictionProvider;
52		55
53	/// Maximum number of events to track.	56	/// Maximum number of events to track.
@@ -59,9 +62,10 @@ pub const DEFAULT_EXCERPT_OPTIONS: EditPredictionExcerptOptions = EditPrediction · 5 unmodified lines
59	target_before_cursor_over_total_bytes: 0.5,	62	target_before_cursor_over_total_bytes: 0.5,
60	};	63	};
61		64
62	pub const DEFAULT_CONTEXT_OPTIONS: ContextMode = ContextMode::Llm(DEFAULT_LLM_CONTEXT_OPTIONS);	65	pub const DEFAULT_CONTEXT_OPTIONS: ContextMode =
		66	ContextMode::Agentic(DEFAULT_AGENTIC_CONTEXT_OPTIONS);
63		67
64	pub const DEFAULT_LLM_CONTEXT_OPTIONS: LlmContextOptions = LlmContextOptions {	68	pub const DEFAULT_AGENTIC_CONTEXT_OPTIONS: AgenticContextOptions = AgenticContextOptions {
65	excerpt: DEFAULT_EXCERPT_OPTIONS,	69	excerpt: DEFAULT_EXCERPT_OPTIONS,
66	};	70	};
67		71
@@ -122,14 +126,19 @@ pub struct ZetaOptions { · 54 unmodified lines
122		126
123	#[derive(Debug, Clone, PartialEq)]	127	#[derive(Debug, Clone, PartialEq)]
124	pub enum ContextMode {	128	pub enum ContextMode {
125	Llm(LlmContextOptions),	129	Agentic(AgenticContextOptions),
126	Syntax(EditPredictionContextOptions),	130	Syntax(EditPredictionContextOptions),
127	}	131	}
128		132
		133	#[derive(Debug, Clone, PartialEq)]
		134	pub struct AgenticContextOptions {
		135	pub excerpt: EditPredictionExcerptOptions,
		136	}
		137
129	impl ContextMode {	138	impl ContextMode {
130	pub fn excerpt(&self) -> &EditPredictionExcerptOptions {	139	pub fn excerpt(&self) -> &EditPredictionExcerptOptions {
131	match self {	140	match self {
132	ContextMode::Llm(options) => &options.excerpt,	141	ContextMode::Agentic(options) => &options.excerpt,
133	ContextMode::Syntax(options) => &options.excerpt,	142	ContextMode::Syntax(options) => &options.excerpt,
134	}	143	}
135	}	144	}
@@ -140,9 +149,8 @@ pub enum ZetaDebugInfo { · 4 unmodified lines
140	ContextRetrievalStarted(ZetaContextRetrievalStartedDebugInfo),	149	ContextRetrievalStarted(ZetaContextRetrievalStartedDebugInfo),
141	SearchQueriesGenerated(ZetaSearchQueryDebugInfo),	150	SearchQueriesGenerated(ZetaSearchQueryDebugInfo),
142	SearchQueriesExecuted(ZetaContextRetrievalDebugInfo),	151	SearchQueriesExecuted(ZetaContextRetrievalDebugInfo),
143	SearchResultsFiltered(ZetaContextRetrievalDebugInfo),
144	ContextRetrievalFinished(ZetaContextRetrievalDebugInfo),	152	ContextRetrievalFinished(ZetaContextRetrievalDebugInfo),
145	EditPredicted(ZetaEditPredictionDebugInfo),	153	EditPredictionRequested(ZetaEditPredictionDebugInfo),
146	}	154	}
147		155
148	#[derive(Debug)]	156	#[derive(Debug)]
@@ -165,14 +173,14 @@ pub struct ZetaEditPredictionDebugInfo { · 16 unmodified lines
165	pub buffer: WeakEntity<Buffer>,	173	pub buffer: WeakEntity<Buffer>,
166	pub position: language::Anchor,	174	pub position: language::Anchor,
167	pub local_prompt: Result<String, String>,	175	pub local_prompt: Result<String, String>,
168	pub response_rx: oneshot::Receiver<Result<predict_edits_v3::PredictEditsResponse, String>>,	176	pub response_rx: oneshot::Receiver<(Result<open_ai::Response, String>, TimeDelta)>,
169	}	177	}
170		178
171	#[derive(Debug)]	179	#[derive(Debug)]
172	pub struct ZetaSearchQueryDebugInfo {	180	pub struct ZetaSearchQueryDebugInfo {
173	pub project: Entity<Project>,	181	pub project: Entity<Project>,
174	pub timestamp: Instant,	182	pub timestamp: Instant,
175	pub queries: Vec<SearchToolQuery>,	183	pub regex_by_glob: HashMap<String, String>,
176	}	184	}
177		185
178	pub type RequestDebugInfo = predict_edits_v3::DebugInfo;	186	pub type RequestDebugInfo = predict_edits_v3::DebugInfo;
@@ -224,7 +232,7 @@ impl CurrentEditPrediction { · 45 unmodified lines
224	{	232	{
225	let (old_range, old_text) = &old_edits[0];	233	let (old_range, old_text) = &old_edits[0];
226	let (new_range, new_text) = &new_edits[0];	234	let (new_range, new_text) = &new_edits[0];
227	new_range == old_range && new_text.starts_with(old_text)	235	new_range == old_range && new_text.starts_with(old_text.as_ref())
228	} else {	236	} else {
229	true	237	true
230	}	238	}
@@ -539,7 +547,7 @@ impl Zeta { · 308 unmodified lines
539	prediction,	547	prediction,
540	} = project_state.current_prediction.as_ref()?;	548	} = project_state.current_prediction.as_ref()?;
541		549
542	if prediction.targets_buffer(buffer.read(cx), cx) {	550	if prediction.targets_buffer(buffer.read(cx)) {
543	Some(BufferEditPrediction::Local { prediction })	551	Some(BufferEditPrediction::Local { prediction })
544	} else if *requested_by_buffer_id == buffer.entity_id() {	552	} else if *requested_by_buffer_id == buffer.entity_id() {
545	Some(BufferEditPrediction::Jump { prediction })	553	Some(BufferEditPrediction::Jump { prediction })
@@ -639,7 +647,7 @@ impl Zeta { · 93 unmodified lines
639	pub fn request_prediction(	647	pub fn request_prediction(
640	&mut self,	648	&mut self,
641	project: &Entity<Project>,	649	project: &Entity<Project>,
642	buffer: &Entity<Buffer>,	650	active_buffer: &Entity<Buffer>,
643	position: language::Anchor,	651	position: language::Anchor,
644	cx: &mut Context<Self>,	652	cx: &mut Context<Self>,
645	) -> Task<Result<Option<EditPrediction>>> {	653	) -> Task<Result<Option<EditPrediction>>> {
@@ -651,8 +659,8 @@ impl Zeta { · 5 unmodified lines
651	.read_with(cx, \|index, _cx\| index.state().clone())	659	.read_with(cx, \|index, _cx\| index.state().clone())
652	});	660	});
653	let options = self.options.clone();	661	let options = self.options.clone();
654	let snapshot = buffer.read(cx).snapshot();	662	let active_snapshot = active_buffer.read(cx).snapshot();
655	let Some(excerpt_path) = snapshot	663	let Some(excerpt_path) = active_snapshot
656	.file()	664	.file()
657	.map(\|path\| -> Arc<Path> { path.full_path(cx).into() })	665	.map(\|path\| -> Arc<Path> { path.full_path(cx).into() })
658	else {	666	else {
@@ -678,12 +686,13 @@ impl Zeta { · 19 unmodified lines
678	})	686	})
679	.unwrap_or_default();	687	.unwrap_or_default();
680		688
681	let diagnostics = snapshot.diagnostic_sets().clone();	689	let diagnostics = active_snapshot.diagnostic_sets().clone();
682		690
683	let parent_abs_path = project::File::from_dyn(buffer.read(cx).file()).and_then(\|f\| {	691	let parent_abs_path =
684	let mut path = f.worktree.read(cx).absolutize(&f.path);	692	project::File::from_dyn(active_buffer.read(cx).file()).and_then(\|f\| {
685	if path.pop() { Some(path) } else { None }	693	let mut path = f.worktree.read(cx).absolutize(&f.path);
686	});	694	if path.pop() { Some(path) } else { None }
		695	});
687		696
688	// TODO data collection	697	// TODO data collection
689	let can_collect_data = cx.is_staff();	698	let can_collect_data = cx.is_staff();
@@ -692,9 +701,10 @@ impl Zeta { · 2 unmodified lines
692	.and_then(\|project_state\| project_state.context.as_ref())	701	.and_then(\|project_state\| project_state.context.as_ref())
693	.unwrap_or(&HashMap::default())	702	.unwrap_or(&HashMap::default())
694	.iter()	703	.iter()
695	.filter_map(\|(buffer, ranges)\| {	704	.filter_map(\|(buffer_entity, ranges)\| {
696	let buffer = buffer.read(cx);	705	let buffer = buffer_entity.read(cx);
697	Some((	706	Some((
		707	buffer_entity.clone(),
698	buffer.snapshot(),	708	buffer.snapshot(),
699	buffer.file()?.full_path(cx).into(),	709	buffer.file()?.full_path(cx).into(),
700	ranges.clone(),	710	ranges.clone(),
@@ -703,8 +713,7 @@ impl Zeta { · 2 unmodified lines
703	.collect::<Vec<_>>();	713	.collect::<Vec<_>>();
704		714
705	let request_task = cx.background_spawn({	715	let request_task = cx.background_spawn({
706	let snapshot = snapshot.clone();	716	let active_buffer = active_buffer.clone();
707	let buffer = buffer.clone();
708	async move {	717	async move {
709	let index_state = if let Some(index_state) = index_state {	718	let index_state = if let Some(index_state) = index_state {
710	Some(index_state.lock_owned().await)	719	Some(index_state.lock_owned().await)
@@ -712,8 +721,8 @@ impl Zeta { · 1 unmodified line
712	None	721	None
713	};	722	};
714		723
715	let cursor_offset = position.to_offset(&snapshot);	724	let cursor_offset = position.to_offset(&active_snapshot);
716	let cursor_point = cursor_offset.to_point(&snapshot);	725	let cursor_point = cursor_offset.to_point(&active_snapshot);
717		726
718	let before_retrieval = chrono::Utc::now();	727	let before_retrieval = chrono::Utc::now();
719		728
@@ -721,29 +730,30 @@ impl Zeta { · 1 unmodified line
721	Self::gather_nearby_diagnostics(	730	Self::gather_nearby_diagnostics(
722	cursor_offset,	731	cursor_offset,
723	&diagnostics,	732	&diagnostics,
724	&snapshot,	733	&active_snapshot,
725	options.max_diagnostic_bytes,	734	options.max_diagnostic_bytes,
726	);	735	);
727		736
728	let request = match options.context {	737	let cloud_request = match options.context {
729	ContextMode::Llm(context_options) => {	738	ContextMode::Agentic(context_options) => {
730	let Some(excerpt) = EditPredictionExcerpt::select_from_buffer(	739	let Some(excerpt) = EditPredictionExcerpt::select_from_buffer(
731	cursor_point,	740	cursor_point,
732	&snapshot,	741	&active_snapshot,
733	&context_options.excerpt,	742	&context_options.excerpt,
734	index_state.as_deref(),	743	index_state.as_deref(),
735	) else {	744	) else {
736	return Ok((None, None));	745	return Ok((None, None));
737	};	746	};
738		747
739	let excerpt_anchor_range = snapshot.anchor_after(excerpt.range.start)	748	let excerpt_anchor_range = active_snapshot.anchor_after(excerpt.range.start)
740	..snapshot.anchor_before(excerpt.range.end);	749	..active_snapshot.anchor_before(excerpt.range.end);
741		750
742	if let Some(buffer_ix) = included_files	751	if let Some(buffer_ix) =
743	.iter()	752	included_files.iter().position(\|(_, snapshot, _, _)\| {
744	.position(\|(buffer, _, _)\| buffer.remote_id() == snapshot.remote_id())	753	snapshot.remote_id() == active_snapshot.remote_id()
		754	})
745	{	755	{
746	let (buffer, _, ranges) = &mut included_files[buffer_ix];	756	let (_, buffer, _, ranges) = &mut included_files[buffer_ix];
747	let range_ix = ranges	757	let range_ix = ranges
748	.binary_search_by(\|probe\| {	758	.binary_search_by(\|probe\| {
749	probe	759	probe
@@ -758,15 +768,16 @@ impl Zeta { · 8 unmodified lines
758	included_files.swap(buffer_ix, last_ix);	768	included_files.swap(buffer_ix, last_ix);
759	} else {	769	} else {
760	included_files.push((	770	included_files.push((
761	snapshot,	771	active_buffer.clone(),
		772	active_snapshot,
762	excerpt_path.clone(),	773	excerpt_path.clone(),
763	vec![excerpt_anchor_range],	774	vec![excerpt_anchor_range],
764	));	775	));
765	}	776	}
766		777
767	let included_files = included_files	778	let included_files = included_files
768	.into_iter()	779	.iter()
769	.map(\|(buffer, path, ranges)\| {	780	.map(\|(_, buffer, path, ranges)\| {
770	let excerpts = merge_excerpts(	781	let excerpts = merge_excerpts(
771	&buffer,	782	&buffer,
772	ranges.iter().map(\|range\| {	783	ranges.iter().map(\|range\| {
@@ -775,7 +786,7 @@ impl Zeta { · 2 unmodified lines
775	}),	786	}),
776	);	787	);
777	predict_edits_v3::IncludedFile {	788	predict_edits_v3::IncludedFile {
778	path,	789	path: path.clone(),
779	max_row: Line(buffer.max_point().row),	790	max_row: Line(buffer.max_point().row),
780	excerpts,	791	excerpts,
781	}	792	}
@@ -809,7 +820,7 @@ impl Zeta { · 27 unmodified lines
809	ContextMode::Syntax(context_options) => {	820	ContextMode::Syntax(context_options) => {
810	let Some(context) = EditPredictionContext::gather_context(	821	let Some(context) = EditPredictionContext::gather_context(
811	cursor_point,	822	cursor_point,
812	&snapshot,	823	&active_snapshot,
813	parent_abs_path.as_deref(),	824	parent_abs_path.as_deref(),
814	&context_options,	825	&context_options,
815	index_state.as_deref(),	826	index_state.as_deref(),
@@ -834,24 +845,27 @@ impl Zeta { · 18 unmodified lines
834	}	845	}
835	};	846	};
836		847
		848	let prompt_result = cloud_zeta2_prompt::build_prompt(&cloud_request);
		849
837	let retrieval_time = chrono::Utc::now() - before_retrieval;	850	let retrieval_time = chrono::Utc::now() - before_retrieval;
838		851
839	let debug_response_tx = if let Some(debug_tx) = &debug_tx {	852	let debug_response_tx = if let Some(debug_tx) = &debug_tx {
840	let (response_tx, response_rx) = oneshot::channel();	853	let (response_tx, response_rx) = oneshot::channel();
841		854
842	let local_prompt = build_prompt(&request)
843	.map(\|(prompt, _)\| prompt)
844	.map_err(\|err\| err.to_string());
845
846	debug_tx	855	debug_tx
847	.unbounded_send(ZetaDebugInfo::EditPredicted(ZetaEditPredictionDebugInfo {	856	.unbounded_send(ZetaDebugInfo::EditPredictionRequested(
848	request: request.clone(),	857	ZetaEditPredictionDebugInfo {
849	retrieval_time,	858	request: cloud_request.clone(),
850	buffer: buffer.downgrade(),	859	retrieval_time,
851	local_prompt,	860	buffer: active_buffer.downgrade(),
852	position,	861	local_prompt: match prompt_result.as_ref() {
853	response_rx,	862	Ok((prompt, _)) => Ok(prompt.clone()),
854	}))	863	Err(err) => Err(err.to_string()),
		864	},
		865	position,
		866	response_rx,
		867	},
		868	))
855	.ok();	869	.ok();
856	Some(response_tx)	870	Some(response_tx)
857	} else {	871	} else {
@@ -861,61 +875,144 @@ impl Zeta { · 3 unmodified lines
861	if cfg!(debug_assertions) && std::env::var("ZED_ZETA2_SKIP_REQUEST").is_ok() {	875	if cfg!(debug_assertions) && std::env::var("ZED_ZETA2_SKIP_REQUEST").is_ok() {
862	if let Some(debug_response_tx) = debug_response_tx {	876	if let Some(debug_response_tx) = debug_response_tx {
863	debug_response_tx	877	debug_response_tx
864	.send(Err("Request skipped".to_string()))	878	.send((Err("Request skipped".to_string()), TimeDelta::zero()))
865	.ok();	879	.ok();
866	}	880	}
867	anyhow::bail!("Skipping request because ZED_ZETA2_SKIP_REQUEST is set")	881	anyhow::bail!("Skipping request because ZED_ZETA2_SKIP_REQUEST is set")
868	}	882	}
869		883
		884	let (prompt, _) = prompt_result?;
		885	let request = open_ai::Request {
		886	model: std::env::var("ZED_ZETA2_MODEL").unwrap_or("yqvev8r3".to_string()),
		887	messages: vec![open_ai::RequestMessage::User {
		888	content: open_ai::MessageContent::Plain(prompt),
		889	}],
		890	stream: false,
		891	max_completion_tokens: None,
		892	stop: Default::default(),
		893	temperature: 0.7,
		894	tool_choice: None,
		895	parallel_tool_calls: None,
		896	tools: vec![],
		897	prompt_cache_key: None,
		898	reasoning_effort: None,
		899	};
		900
		901	log::trace!("Sending edit prediction request");
		902
		903	let before_request = chrono::Utc::now();
870	let response =	904	let response =
871	Self::send_prediction_request(client, llm_token, app_version, request).await;	905	Self::send_raw_llm_request(client, llm_token, app_version, request).await;
		906	let request_time = chrono::Utc::now() - before_request;
		907
		908	log::trace!("Got edit prediction response");
872		909
873	if let Some(debug_response_tx) = debug_response_tx {	910	if let Some(debug_response_tx) = debug_response_tx {
874	debug_response_tx	911	debug_response_tx
875	.send(	912	.send((
876	response	913	response
877	.as_ref()	914	.as_ref()
878	.map_err(\|err\| err.to_string())	915	.map_err(\|err\| err.to_string())
879	.map(\|response\| response.0.clone()),	916	.map(\|response\| response.0.clone()),
880	)	917	request_time,
		918	))
881	.ok();	919	.ok();
882	}	920	}
883		921
884	response.map(\|(res, usage)\| (Some(res), usage))	922	let (mut res, usage) = response?;
		923
		924	let request_id = EditPredictionId(Uuid::from_str(&res.id)?);
		925
		926	let Some(choice) = res.choices.pop() else {
		927	return Ok((None, usage));
		928	};
		929
		930	let output_text = match choice.message {
		931	open_ai::RequestMessage::Assistant {
		932	content: Some(open_ai::MessageContent::Plain(content)),
		933	..
		934	} => content,
		935	open_ai::RequestMessage::Assistant {
		936	content: Some(open_ai::MessageContent::Multipart(mut content)),
		937	..
		938	} => {
		939	if content.is_empty() {
		940	log::error!("No output from Baseten completion response");
		941	return Ok((None, usage));
		942	}
		943
		944	match content.remove(0) {
		945	open_ai::MessagePart::Text { text } => text,
		946	open_ai::MessagePart::Image { .. } => {
		947	log::error!("Expected text, got an image");
		948	return Ok((None, usage));
		949	}
		950	}
		951	}
		952	_ => {
		953	log::error!("Invalid response message: {:?}", choice.message);
		954	return Ok((None, usage));
		955	}
		956	};
		957
		958	let (edited_buffer_snapshot, edits) =
		959	crate::udiff::parse_diff(&output_text, \|path\| {
		960	included_files
		961	.iter()
		962	.find_map(\|(_, buffer, probe_path, ranges)\| {
		963	if probe_path.as_ref() == path {
		964	Some((buffer, ranges.as_slice()))
		965	} else {
		966	None
		967	}
		968	})
		969	})
		970	.await?;
		971
		972	let edited_buffer = included_files
		973	.iter()
		974	.find_map(\|(buffer, snapshot, _, _)\| {
		975	if snapshot.remote_id() == edited_buffer_snapshot.remote_id() {
		976	Some(buffer.clone())
		977	} else {
		978	None
		979	}
		980	})
		981	.context("Failed to find buffer in included_buffers, even though we just found the snapshot")?;
		982
		983	anyhow::Ok((Some((request_id, edited_buffer, edited_buffer_snapshot.clone(), edits)), usage))
885	}	984	}
886	});	985	});
887		986
888	let buffer = buffer.clone();
889
890	cx.spawn({	987	cx.spawn({
891	let project = project.clone();
892	async move \|this, cx\| {	988	async move \|this, cx\| {
893	let Some(response) = Self::handle_api_response(&this, request_task.await, cx)?	989	let Some((id, edited_buffer, edited_buffer_snapshot, edits)) =
		990	Self::handle_api_response(&this, request_task.await, cx)?
894	else {	991	else {
895	return Ok(None);	992	return Ok(None);
896	};	993	};
897		994
898	// TODO telemetry: duration, etc	995	// TODO telemetry: duration, etc
899	Ok(EditPrediction::from_response(response, &snapshot, &buffer, &project, cx).await)	996	Ok(
		997	EditPrediction::new(id, &edited_buffer, &edited_buffer_snapshot, edits, cx)
		998	.await,
		999	)
900	}	1000	}
901	})	1001	})
902	}	1002	}
903		1003
904	async fn send_prediction_request(	1004	async fn send_raw_llm_request(
905	client: Arc<Client>,	1005	client: Arc<Client>,
906	llm_token: LlmApiToken,	1006	llm_token: LlmApiToken,
907	app_version: SemanticVersion,	1007	app_version: SemanticVersion,
908	request: predict_edits_v3::PredictEditsRequest,	1008	request: open_ai::Request,
909	) -> Result<(	1009	) -> Result<(open_ai::Response, Option<EditPredictionUsage>)> {
910	predict_edits_v3::PredictEditsResponse,
911	Option<EditPredictionUsage>,
912	)> {
913	let url = if let Ok(predict_edits_url) = std::env::var("ZED_PREDICT_EDITS_URL") {	1010	let url = if let Ok(predict_edits_url) = std::env::var("ZED_PREDICT_EDITS_URL") {
914	http_client::Url::parse(&predict_edits_url)?	1011	http_client::Url::parse(&predict_edits_url)?
915	} else {	1012	} else {
916	client	1013	client
917	.http_client()	1014	.http_client()
918	.build_zed_llm_url("/predict_edits/v3", &[])?	1015	.build_zed_llm_url("/predict_edits/raw", &[])?
919	};	1016	};
920		1017
921	Self::send_api_request(	1018	Self::send_api_request(
@@ -1052,7 +1149,7 @@ impl Zeta { · 130 unmodified lines
1052	cursor_position: language::Anchor,	1149	cursor_position: language::Anchor,
1053	cx: &mut Context<Self>,	1150	cx: &mut Context<Self>,
1054	) {	1151	) {
1055	if !matches!(&self.options().context, ContextMode::Llm { .. }) {	1152	if !matches!(&self.options().context, ContextMode::Agentic { .. }) {
1056	return;	1153	return;
1057	}	1154	}
1058		1155
@@ -1100,36 +1197,149 @@ impl Zeta { · 41 unmodified lines
1100	cursor_position: language::Anchor,	1197	cursor_position: language::Anchor,
1101	cx: &mut Context<Self>,	1198	cx: &mut Context<Self>,
1102	) -> Task<Result<()>> {	1199	) -> Task<Result<()>> {
		1200	let Some(zeta_project) = self.projects.get(&project.entity_id()) else {
		1201	return Task::ready(anyhow::Ok(()));
		1202	};
		1203
		1204	let ContextMode::Agentic(options) = &self.options().context else {
		1205	return Task::ready(anyhow::Ok(()));
		1206	};
		1207
		1208	let snapshot = buffer.read(cx).snapshot();
		1209	let cursor_point = cursor_position.to_point(&snapshot);
		1210	let Some(cursor_excerpt) = EditPredictionExcerpt::select_from_buffer(
		1211	cursor_point,
		1212	&snapshot,
		1213	&options.excerpt,
		1214	None,
		1215	) else {
		1216	return Task::ready(Ok(()));
		1217	};
		1218
		1219	let current_file_path: Arc<Path> = snapshot
		1220	.file()
		1221	.map(\|f\| f.full_path(cx).into())
		1222	.unwrap_or_else(\|\| Path::new("untitled").into());
		1223
		1224	let prompt = match cloud_zeta2_prompt::retrieval_prompt::build_prompt(
		1225	predict_edits_v3::PlanContextRetrievalRequest {
		1226	excerpt: cursor_excerpt.text(&snapshot).body,
		1227	excerpt_path: current_file_path,
		1228	excerpt_line_range: cursor_excerpt.line_range,
		1229	cursor_file_max_row: Line(snapshot.max_point().row),
		1230	events: zeta_project
		1231	.events
		1232	.iter()
		1233	.filter_map(\|ev\| ev.to_request_event(cx))
		1234	.collect(),
		1235	},
		1236	) {
		1237	Ok(prompt) => prompt,
		1238	Err(err) => {
		1239	return Task::ready(Err(err));
		1240	}
		1241	};
		1242
		1243	let app_version = AppVersion::global(cx);
		1244	let client = self.client.clone();
		1245	let llm_token = self.llm_token.clone();
		1246	let debug_tx = self.debug_tx.clone();
		1247
		1248	let (tool_schema, tool_description) = &*cloud_zeta2_prompt::retrieval_prompt::TOOL_SCHEMA;
		1249
		1250	let request = open_ai::Request {
		1251	model: std::env::var("ZED_ZETA2_MODEL").unwrap_or("2327jz9q".to_string()),
		1252	messages: vec![open_ai::RequestMessage::User {
		1253	content: open_ai::MessageContent::Plain(prompt),
		1254	}],
		1255	stream: false,
		1256	max_completion_tokens: None,
		1257	stop: Default::default(),
		1258	temperature: 0.7,
		1259	tool_choice: None,
		1260	parallel_tool_calls: None,
		1261	tools: vec![open_ai::ToolDefinition::Function {
		1262	function: FunctionDefinition {
		1263	name: cloud_zeta2_prompt::retrieval_prompt::TOOL_NAME.to_string(),
		1264	description: Some(tool_description.clone()),
		1265	parameters: Some(tool_schema.clone()),
		1266	},
		1267	}],
		1268	prompt_cache_key: None,
		1269	reasoning_effort: None,
		1270	};
		1271
1103	cx.spawn(async move \|this, cx\| {	1272	cx.spawn(async move \|this, cx\| {
1104	let related_excerpts_result = this	1273	log::trace!("Sending search planning request");
1105	.update(cx, \|this, cx\| {	1274	let response =
1106	let Some(zeta_project) = this.projects.get(&project.entity_id()) else {	1275	Self::send_raw_llm_request(client, llm_token, app_version, request).await;
1107	return Task::ready(anyhow::Ok(HashMap::default()));	1276	let mut response = Self::handle_api_response(&this, response, cx)?;
1108	};	1277
		1278	log::trace!("Got search planning response");
		1279
		1280	let choice = response
		1281	.choices
		1282	.pop()
		1283	.context("No choices in retrieval response")?;
		1284	let open_ai::RequestMessage::Assistant {
		1285	content: _,
		1286	tool_calls,
		1287	} = choice.message
		1288	else {
		1289	anyhow::bail!("Retrieval response didn't include an assistant message");
		1290	};
1109		1291
1110	let ContextMode::Llm(options) = &this.options().context else {	1292	let mut regex_by_glob: HashMap<String, String> = HashMap::default();
1111	return Task::ready(anyhow::Ok(HashMap::default()));	1293	for tool_call in tool_calls {
1112	};	1294	let open_ai::ToolCallContent::Function { function } = tool_call.content;
		1295	if function.name != cloud_zeta2_prompt::retrieval_prompt::TOOL_NAME {
		1296	log::warn!(
		1297	"Context retrieval response tried to call an unknown tool: {}",
		1298	function.name
		1299	);
1113		1300
1114	let mut edit_history_unified_diff = String::new();	1301	continue;
		1302	}
1115		1303
1116	for event in zeta_project.events.iter() {	1304	let input: SearchToolInput = serde_json::from_str(&function.arguments)?;
1117	if let Some(event) = event.to_request_event(cx) {	1305	for query in input.queries {
1118	writeln!(&mut edit_history_unified_diff, "{event}").ok();	1306	let regex = regex_by_glob.entry(query.glob).or_default();
1119	}	1307	if !regex.is_empty() {
		1308	regex.push('\|');
1120	}	1309	}
		1310	regex.push_str(&query.regex);
		1311	}
		1312	}
1121		1313
1122	find_related_excerpts(	1314	if let Some(debug_tx) = &debug_tx {
1123	buffer.clone(),	1315	debug_tx
1124	cursor_position,	1316	.unbounded_send(ZetaDebugInfo::SearchQueriesGenerated(
1125	&project,	1317	ZetaSearchQueryDebugInfo {
1126	edit_history_unified_diff,	1318	project: project.clone(),
1127	options,	1319	timestamp: Instant::now(),
1128	this.debug_tx.clone(),	1320	regex_by_glob: regex_by_glob.clone(),
1129	cx,	1321	},
1130	)	1322	))
1131	})?	1323	.ok();
1132	.await;	1324	}
		1325
		1326	log::trace!("Running retrieval search: {regex_by_glob:#?}");
		1327
		1328	let related_excerpts_result =
		1329	retrieval_search::run_retrieval_searches(project.clone(), regex_by_glob, cx).await;
		1330
		1331	log::trace!("Search queries executed");
		1332
		1333	if let Some(debug_tx) = &debug_tx {
		1334	debug_tx
		1335	.unbounded_send(ZetaDebugInfo::SearchQueriesExecuted(
		1336	ZetaContextRetrievalDebugInfo {
		1337	project: project.clone(),
		1338	timestamp: Instant::now(),
		1339	},
		1340	))
		1341	.ok();
		1342	}
1133		1343
1134	this.update(cx, \|this, _cx\| {	1344	this.update(cx, \|this, _cx\| {
1135	let Some(zeta_project) = this.projects.get_mut(&project.entity_id()) else {	1345	let Some(zeta_project) = this.projects.get_mut(&project.entity_id()) else {
@@ -1249,7 +1459,7 @@ impl Zeta { · 113 unmodified lines
1249	&snapshot,	1459	&snapshot,
1250	parent_abs_path.as_deref(),	1460	parent_abs_path.as_deref(),
1251	match &options.context {	1461	match &options.context {
1252	ContextMode::Llm(_) => {	1462	ContextMode::Agentic(_) => {
1253	// TODO	1463	// TODO
1254	panic!("Llm mode not supported in zeta cli yet");	1464	panic!("Llm mode not supported in zeta cli yet");
1255	}	1465	}
@@ -1426,15 +1636,11 @@ fn add_signature( · 170 unmodified lines
1426		1636
1427	#[cfg(test)]	1637	#[cfg(test)]
1428	mod tests {	1638	mod tests {
1429	use std::{	1639	use std::{path::Path, sync::Arc};
1430	path::{Path, PathBuf},
1431	sync::Arc,
1432	};
1433		1640
1434	use client::UserStore;	1641	use client::UserStore;
1435	use clock::FakeSystemClock;	1642	use clock::FakeSystemClock;
1436	use cloud_llm_client::predict_edits_v3::{self, Point};	1643	use cloud_zeta2_prompt::retrieval_prompt::{SearchToolInput, SearchToolQuery};
1437	use edit_prediction_context::Line;
1438	use futures::{	1644	use futures::{
1439	AsyncReadExt, StreamExt,	1645	AsyncReadExt, StreamExt,
1440	channel::{mpsc, oneshot},	1646	channel::{mpsc, oneshot},
@@ -1445,7 +1651,8 @@ mod tests { · 4 unmodified lines
1445	prelude::*,	1651	prelude::*,
1446	};	1652	};
1447	use indoc::indoc;	1653	use indoc::indoc;
1448	use language::{LanguageServerId, OffsetRangeExt as _};	1654	use language::OffsetRangeExt as _;
		1655	use open_ai::Usage;
1449	use pretty_assertions::{assert_eq, assert_matches};	1656	use pretty_assertions::{assert_eq, assert_matches};
1450	use project::{FakeFs, Project};	1657	use project::{FakeFs, Project};
1451	use serde_json::json;	1658	use serde_json::json;
@@ -1462,8 +1669,8 @@ mod tests { · 10 unmodified lines
1462	fs.insert_tree(	1669	fs.insert_tree(
1463	"/root",	1670	"/root",
1464	json!({	1671	json!({
1465	"1.txt": "Hello!\nHow\nBye",	1672	"1.txt": "Hello!\nHow\nBye\n",
1466	"2.txt": "Hola!\nComo\nAdios"	1673	"2.txt": "Hola!\nComo\nAdios\n"
1467	}),	1674	}),
1468	)	1675	)
1469	.await;	1676	.await;
@@ -1489,16 +1696,17 @@ mod tests { · 19 unmodified lines
1489	zeta.refresh_prediction(&project, &buffer1, position, cx)	1696	zeta.refresh_prediction(&project, &buffer1, position, cx)
1490	});	1697	});
1491	let (_request, respond_tx) = req_rx.next().await.unwrap();	1698	let (_request, respond_tx) = req_rx.next().await.unwrap();
		1699
1492	respond_tx	1700	respond_tx
1493	.send(predict_edits_v3::PredictEditsResponse {	1701	.send(model_response(indoc! {r"
1494	request_id: Uuid::new_v4(),	1702	--- a/root/1.txt
1495	edits: vec![predict_edits_v3::Edit {	1703	+++ b/root/1.txt
1496	path: Path::new(path!("root/1.txt")).into(),	1704	@@ ... @@
1497	range: Line(0)..Line(snapshot1.max_point().row + 1),	1705	Hello!
1498	content: "Hello!\nHow are you?\nBye".into(),	1706	-How
1499	}],	1707	+How are you?
1500	debug_info: None,	1708	Bye
1501	})	1709	"}))
1502	.unwrap();	1710	.unwrap();
1503	prediction_task.await.unwrap();	1711	prediction_task.await.unwrap();
1504		1712
@@ -1509,21 +1717,67 @@ mod tests { · 4 unmodified lines
1509	assert_matches!(prediction, BufferEditPrediction::Local { .. });	1717	assert_matches!(prediction, BufferEditPrediction::Local { .. });
1510	});	1718	});
1511		1719
		1720	// Context refresh
		1721	let refresh_task = zeta.update(cx, \|zeta, cx\| {
		1722	zeta.refresh_context(project.clone(), buffer1.clone(), position, cx)
		1723	});
		1724	let (_request, respond_tx) = req_rx.next().await.unwrap();
		1725	respond_tx
		1726	.send(open_ai::Response {
		1727	id: Uuid::new_v4().to_string(),
		1728	object: "response".into(),
		1729	created: 0,
		1730	model: "model".into(),
		1731	choices: vec![open_ai::Choice {
		1732	index: 0,
		1733	message: open_ai::RequestMessage::Assistant {
		1734	content: None,
		1735	tool_calls: vec![open_ai::ToolCall {
		1736	id: "search".into(),
		1737	content: open_ai::ToolCallContent::Function {
		1738	function: open_ai::FunctionContent {
		1739	name: cloud_zeta2_prompt::retrieval_prompt::TOOL_NAME
		1740	.to_string(),
		1741	arguments: serde_json::to_string(&SearchToolInput {
		1742	queries: Box::new([SearchToolQuery {
		1743	glob: "root/2.txt".to_string(),
		1744	regex: ".".to_string(),
		1745	}]),
		1746	})
		1747	.unwrap(),
		1748	},
		1749	},
		1750	}],
		1751	},
		1752	finish_reason: None,
		1753	}],
		1754	usage: Usage {
		1755	prompt_tokens: 0,
		1756	completion_tokens: 0,
		1757	total_tokens: 0,
		1758	},
		1759	})
		1760	.unwrap();
		1761	refresh_task.await.unwrap();
		1762
		1763	zeta.update(cx, \|zeta, _cx\| {
		1764	zeta.discard_current_prediction(&project);
		1765	});
		1766
1512	// Prediction for another file	1767	// Prediction for another file
1513	let prediction_task = zeta.update(cx, \|zeta, cx\| {	1768	let prediction_task = zeta.update(cx, \|zeta, cx\| {
1514	zeta.refresh_prediction(&project, &buffer1, position, cx)	1769	zeta.refresh_prediction(&project, &buffer1, position, cx)
1515	});	1770	});
1516	let (_request, respond_tx) = req_rx.next().await.unwrap();	1771	let (_request, respond_tx) = req_rx.next().await.unwrap();
1517	respond_tx	1772	respond_tx
1518	.send(predict_edits_v3::PredictEditsResponse {	1773	.send(model_response(indoc! {r#"
1519	request_id: Uuid::new_v4(),	1774	--- a/root/2.txt
1520	edits: vec![predict_edits_v3::Edit {	1775	+++ b/root/2.txt
1521	path: Path::new(path!("root/2.txt")).into(),	1776	Hola!
1522	range: Line(0)..Line(snapshot1.max_point().row + 1),	1777	-Como
1523	content: "Hola!\nComo estas?\nAdios".into(),	1778	+Como estas?
1524	}],	1779	Adios
1525	debug_info: None,	1780	"#}))
1526	})
1527	.unwrap();	1781	.unwrap();
1528	prediction_task.await.unwrap();	1782	prediction_task.await.unwrap();
1529	zeta.read_with(cx, \|zeta, cx\| {	1783	zeta.read_with(cx, \|zeta, cx\| {
@@ -1532,7 +1786,7 @@ mod tests { · 2 unmodified lines
1532	.unwrap();	1786	.unwrap();
1533	assert_matches!(	1787	assert_matches!(
1534	prediction,	1788	prediction,
1535	BufferEditPrediction::Jump { prediction } if prediction.path.as_ref() == Path::new(path!("root/2.txt"))	1789	BufferEditPrediction::Jump { prediction } if prediction.snapshot.file().unwrap().full_path(cx) == Path::new(path!("root/2.txt"))
1536	);	1790	);
1537	});	1791	});
1538		1792
@@ -1559,7 +1813,7 @@ mod tests { · 20 unmodified lines
1559	fs.insert_tree(	1813	fs.insert_tree(
1560	"/root",	1814	"/root",
1561	json!({	1815	json!({
1562	"foo.md": "Hello!\nHow\nBye"	1816	"foo.md": "Hello!\nHow\nBye\n"
1563	}),	1817	}),
1564	)	1818	)
1565	.await;	1819	.await;
@@ -1579,29 +1833,31 @@ mod tests { · 13 unmodified lines
1579	zeta.request_prediction(&project, &buffer, position, cx)	1833	zeta.request_prediction(&project, &buffer, position, cx)
1580	});	1834	});
1581		1835
1582	let (request, respond_tx) = req_rx.next().await.unwrap();	1836	let (_, respond_tx) = req_rx.next().await.unwrap();
1583	assert_eq!(	1837
1584	request.excerpt_path.as_ref(),	1838	// TODO Put back when we have a structured request again
1585	Path::new(path!("root/foo.md"))	1839	// assert_eq!(
1586	);	1840	// request.excerpt_path.as_ref(),
1587	assert_eq!(	1841	// Path::new(path!("root/foo.md"))
1588	request.cursor_point,	1842	// );
1589	Point {	1843	// assert_eq!(
1590	line: Line(1),	1844	// request.cursor_point,
1591	column: 3	1845	// Point {
1592	}	1846	// line: Line(1),
1593	);	1847	// column: 3
		1848	// }
		1849	// );
1594		1850
1595	respond_tx	1851	respond_tx
1596	.send(predict_edits_v3::PredictEditsResponse {	1852	.send(model_response(indoc! { r"
1597	request_id: Uuid::new_v4(),	1853	--- a/root/foo.md
1598	edits: vec![predict_edits_v3::Edit {	1854	+++ b/root/foo.md
1599	path: Path::new(path!("root/foo.md")).into(),	1855	@@ ... @@
1600	range: Line(0)..Line(snapshot.max_point().row + 1),	1856	Hello!
1601	content: "Hello!\nHow are you?\nBye".into(),	1857	-How
1602	}],	1858	+How are you?
1603	debug_info: None,	1859	Bye
1604	})	1860	"}))
1605	.unwrap();	1861	.unwrap();
1606		1862
1607	let prediction = prediction_task.await.unwrap().unwrap();	1863	let prediction = prediction_task.await.unwrap().unwrap();
@@ -1611,7 +1867,7 @@ mod tests { · 3 unmodified lines
1611	prediction.edits[0].0.to_point(&snapshot).start,	1867	prediction.edits[0].0.to_point(&snapshot).start,
1612	language::Point::new(1, 3)	1868	language::Point::new(1, 3)
1613	);	1869	);
1614	assert_eq!(prediction.edits[0].1, " are you?");	1870	assert_eq!(prediction.edits[0].1.as_ref(), " are you?");
1615	}	1871	}
1616		1872
1617	#[gpui::test]	1873	#[gpui::test]
@@ -1621,7 +1877,7 @@ mod tests { · 3 unmodified lines
1621	fs.insert_tree(	1877	fs.insert_tree(
1622	"/root",	1878	"/root",
1623	json!({	1879	json!({
1624	"foo.md": "Hello!\n\nBye"	1880	"foo.md": "Hello!\n\nBye\n"
1625	}),	1881	}),
1626	)	1882	)
1627	.await;	1883	.await;
@@ -1652,34 +1908,30 @@ mod tests { · 24 unmodified lines
1652		1908
1653	let (request, respond_tx) = req_rx.next().await.unwrap();	1909	let (request, respond_tx) = req_rx.next().await.unwrap();
1654		1910
1655	assert_eq!(request.events.len(), 1);	1911	let prompt = prompt_from_request(&request);
1656	assert_eq!(	1912	assert!(
1657	request.events[0],	1913	prompt.contains(indoc! {"
1658	predict_edits_v3::Event::BufferChange {	1914	--- a/root/foo.md
1659	path: Some(PathBuf::from(path!("root/foo.md"))),	1915	+++ b/root/foo.md
1660	old_path: None,	1916	@@ -1,3 +1,3 @@
1661	diff: indoc! {"	1917	Hello!
1662	@@ -1,3 +1,3 @@	1918	-
1663	Hello!	1919	+How
1664	-	1920	Bye
1665	+How	1921	"}),
1666	Bye	1922	"{prompt}"
1667	"}
1668	.to_string(),
1669	predicted: false
1670	}
1671	);	1923	);
1672		1924
1673	respond_tx	1925	respond_tx
1674	.send(predict_edits_v3::PredictEditsResponse {	1926	.send(model_response(indoc! {r#"
1675	request_id: Uuid::new_v4(),	1927	--- a/root/foo.md
1676	edits: vec![predict_edits_v3::Edit {	1928	+++ b/root/foo.md
1677	path: Path::new(path!("root/foo.md")).into(),	1929	@@ ... @@
1678	range: Line(0)..Line(snapshot.max_point().row + 1),	1930	Hello!
1679	content: "Hello!\nHow are you?\nBye".into(),	1931	-How
1680	}],	1932	+How are you?
1681	debug_info: None,	1933	Bye
1682	})	1934	"#}))
1683	.unwrap();	1935	.unwrap();
1684		1936
1685	let prediction = prediction_task.await.unwrap().unwrap();	1937	let prediction = prediction_task.await.unwrap().unwrap();
@@ -1689,114 +1941,150 @@ mod tests { · 3 unmodified lines
1689	prediction.edits[0].0.to_point(&snapshot).start,	1941	prediction.edits[0].0.to_point(&snapshot).start,
1690	language::Point::new(1, 3)	1942	language::Point::new(1, 3)
1691	);	1943	);
1692	assert_eq!(prediction.edits[0].1, " are you?");	1944	assert_eq!(prediction.edits[0].1.as_ref(), " are you?");
1693	}	1945	}
1694		1946
1695	#[gpui::test]	1947	// Skipped until we start including diagnostics in prompt
1696	async fn test_request_diagnostics(cx: &mut TestAppContext) {	1948	// #[gpui::test]
1697	let (zeta, mut req_rx) = init_test(cx);	1949	// async fn test_request_diagnostics(cx: &mut TestAppContext) {
1698	let fs = FakeFs::new(cx.executor());	1950	// let (zeta, mut req_rx) = init_test(cx);
1699	fs.insert_tree(	1951	// let fs = FakeFs::new(cx.executor());
1700	"/root",	1952	// fs.insert_tree(
1701	json!({	1953	// "/root",
1702	"foo.md": "Hello!\nBye"	1954	// json!({
1703	}),	1955	// "foo.md": "Hello!\nBye"
1704	)	1956	// }),
1705	.await;	1957	// )
1706	let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;	1958	// .await;
		1959	// let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
		1960
		1961	// let path_to_buffer_uri = lsp::Uri::from_file_path(path!("/root/foo.md")).unwrap();
		1962	// let diagnostic = lsp::Diagnostic {
		1963	// range: lsp::Range::new(lsp::Position::new(1, 1), lsp::Position::new(1, 5)),
		1964	// severity: Some(lsp::DiagnosticSeverity::ERROR),
		1965	// message: "\"Hello\" deprecated. Use \"Hi\" instead".to_string(),
		1966	// ..Default::default()
		1967	// };
		1968
		1969	// project.update(cx, \|project, cx\| {
		1970	// project.lsp_store().update(cx, \|lsp_store, cx\| {
		1971	// // Create some diagnostics
		1972	// lsp_store
		1973	// .update_diagnostics(
		1974	// LanguageServerId(0),
		1975	// lsp::PublishDiagnosticsParams {
		1976	// uri: path_to_buffer_uri.clone(),
		1977	// diagnostics: vec![diagnostic],
		1978	// version: None,
		1979	// },
		1980	// None,
		1981	// language::DiagnosticSourceKind::Pushed,
		1982	// &[],
		1983	// cx,
		1984	// )
		1985	// .unwrap();
		1986	// });
		1987	// });
		1988
		1989	// let buffer = project
		1990	// .update(cx, \|project, cx\| {
		1991	// let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
		1992	// project.open_buffer(path, cx)
		1993	// })
		1994	// .await
		1995	// .unwrap();
		1996
		1997	// let snapshot = buffer.read_with(cx, \|buffer, _cx\| buffer.snapshot());
		1998	// let position = snapshot.anchor_before(language::Point::new(0, 0));
		1999
		2000	// let _prediction_task = zeta.update(cx, \|zeta, cx\| {
		2001	// zeta.request_prediction(&project, &buffer, position, cx)
		2002	// });
		2003
		2004	// let (request, _respond_tx) = req_rx.next().await.unwrap();
		2005
		2006	// assert_eq!(request.diagnostic_groups.len(), 1);
		2007	// let value = serde_json::from_str::<serde_json::Value>(request.diagnostic_groups[0].0.get())
		2008	// .unwrap();
		2009	// // We probably don't need all of this. TODO define a specific diagnostic type in predict_edits_v3
		2010	// assert_eq!(
		2011	// value,
		2012	// json!({
		2013	// "entries": [{
		2014	// "range": {
		2015	// "start": 8,
		2016	// "end": 10
		2017	// },
		2018	// "diagnostic": {
		2019	// "source": null,
		2020	// "code": null,
		2021	// "code_description": null,
		2022	// "severity": 1,
		2023	// "message": "\"Hello\" deprecated. Use \"Hi\" instead",
		2024	// "markdown": null,
		2025	// "group_id": 0,
		2026	// "is_primary": true,
		2027	// "is_disk_based": false,
		2028	// "is_unnecessary": false,
		2029	// "source_kind": "Pushed",
		2030	// "data": null,
		2031	// "underline": true
		2032	// }
		2033	// }],
		2034	// "primary_ix": 0
		2035	// })
		2036	// );
		2037	// }
		2038
		2039	fn model_response(text: &str) -> open_ai::Response {
		2040	open_ai::Response {
		2041	id: Uuid::new_v4().to_string(),
		2042	object: "response".into(),
		2043	created: 0,
		2044	model: "model".into(),
		2045	choices: vec![open_ai::Choice {
		2046	index: 0,
		2047	message: open_ai::RequestMessage::Assistant {
		2048	content: Some(open_ai::MessageContent::Plain(text.to_string())),
		2049	tool_calls: vec![],
		2050	},
		2051	finish_reason: None,
		2052	}],
		2053	usage: Usage {
		2054	prompt_tokens: 0,
		2055	completion_tokens: 0,
		2056	total_tokens: 0,
		2057	},
		2058	}
		2059	}
1707		2060
1708	let path_to_buffer_uri = lsp::Uri::from_file_path(path!("/root/foo.md")).unwrap();	2061	fn prompt_from_request(request: &open_ai::Request) -> &str {
1709	let diagnostic = lsp::Diagnostic {	2062	assert_eq!(request.messages.len(), 1);
1710	range: lsp::Range::new(lsp::Position::new(1, 1), lsp::Position::new(1, 5)),	2063	let open_ai::RequestMessage::User {
1711	severity: Some(lsp::DiagnosticSeverity::ERROR),	2064	content: open_ai::MessageContent::Plain(content),
1712	message: "\"Hello\" deprecated. Use \"Hi\" instead".to_string(),	2065	..
1713	..Default::default()	2066	} = &request.messages[0]
		2067	else {
		2068	panic!(
		2069	"Request does not have single user message of type Plain. {:#?}",
		2070	request
		2071	);
1714	};	2072	};
1715		2073	content
1716	project.update(cx, \|project, cx\| {
1717	project.lsp_store().update(cx, \|lsp_store, cx\| {
1718	// Create some diagnostics
1719	lsp_store
1720	.update_diagnostics(
1721	LanguageServerId(0),
1722	lsp::PublishDiagnosticsParams {
1723	uri: path_to_buffer_uri.clone(),
1724	diagnostics: vec![diagnostic],
1725	version: None,
1726	},
1727	None,
1728	language::DiagnosticSourceKind::Pushed,
1729	&[],
1730	cx,
1731	)
1732	.unwrap();
1733	});
1734	});
1735
1736	let buffer = project
1737	.update(cx, \|project, cx\| {
1738	let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1739	project.open_buffer(path, cx)
1740	})
1741	.await
1742	.unwrap();
1743
1744	let snapshot = buffer.read_with(cx, \|buffer, _cx\| buffer.snapshot());
1745	let position = snapshot.anchor_before(language::Point::new(0, 0));
1746
1747	let _prediction_task = zeta.update(cx, \|zeta, cx\| {
1748	zeta.request_prediction(&project, &buffer, position, cx)
1749	});
1750
1751	let (request, _respond_tx) = req_rx.next().await.unwrap();
1752
1753	assert_eq!(request.diagnostic_groups.len(), 1);
1754	let value = serde_json::from_str::<serde_json::Value>(request.diagnostic_groups[0].0.get())
1755	.unwrap();
1756	// We probably don't need all of this. TODO define a specific diagnostic type in predict_edits_v3
1757	assert_eq!(
1758	value,
1759	json!({
1760	"entries": [{
1761	"range": {
1762	"start": 8,
1763	"end": 10
1764	},
1765	"diagnostic": {
1766	"source": null,
1767	"code": null,
1768	"code_description": null,
1769	"severity": 1,
1770	"message": "\"Hello\" deprecated. Use \"Hi\" instead",
1771	"markdown": null,
1772	"group_id": 0,
1773	"is_primary": true,
1774	"is_disk_based": false,
1775	"is_unnecessary": false,
1776	"source_kind": "Pushed",
1777	"data": null,
1778	"underline": true
1779	}
1780	}],
1781	"primary_ix": 0
1782	})
1783	);
1784	}	2074	}
1785		2075
1786	fn init_test(	2076	fn init_test(
1787	cx: &mut TestAppContext,	2077	cx: &mut TestAppContext,
1788	) -> (	2078	) -> (
1789	Entity<Zeta>,	2079	Entity<Zeta>,
1790	mpsc::UnboundedReceiver<(	2080	mpsc::UnboundedReceiver<(open_ai::Request, oneshot::Sender<open_ai::Response>)>,
1791	predict_edits_v3::PredictEditsRequest,
1792	oneshot::Sender<predict_edits_v3::PredictEditsResponse>,
1793	)>,
1794	) {	2081	) {
1795	cx.update(move \|cx\| {	2082	cx.update(move \|cx\| {
1796	let settings_store = SettingsStore::test(cx);	2083	let settings_store = SettingsStore::test(cx);
1797	cx.set_global(settings_store);	2084	cx.set_global(settings_store);
1798	language::init(cx);	2085	language::init(cx);
1799	Project::init_settings(cx);	2086	Project::init_settings(cx);
		2087	zlog::init_test();
1800		2088
1801	let (req_tx, req_rx) = mpsc::unbounded();	2089	let (req_tx, req_rx) = mpsc::unbounded();
1802		2090
@@ -1811,7 +2099,7 @@ mod tests { · 8 unmodified lines
1811	"token": "test"	2099	"token": "test"
1812	}))	2100	}))
1813	.unwrap(),	2101	.unwrap(),
1814	"/predict_edits/v3" => {	2102	"/predict_edits/raw" => {
1815	let mut buf = Vec::new();	2103	let mut buf = Vec::new();
1816	body.read_to_end(&mut buf).await.ok();	2104	body.read_to_end(&mut buf).await.ok();
1817	let req = serde_json::from_slice(&buf).unwrap();	2105	let req = serde_json::from_slice(&buf).unwrap();

crates/zeta2_tools/Cargo.toml+2−1

 multi_buffer.workspace = true
 ordered-float.workspace = true
 project.workspace = true
+regex-syntax = "0.8.8"
 serde.workspace = true
+serde_json.workspace = true
 telemetry.workspace = true
 text.workspace = true
-regex-syntax = "0.8.8"
 ui.workspace = true
 ui_input.workspace = true
 util.workspace = true

crates/zeta2_tools/src/zeta2_context_view.rs+10−40

     started_at: Instant,
     search_results_generated_at: Option<Instant>,
     search_results_executed_at: Option<Instant>,
-    search_results_filtered_at: Option<Instant>,
     finished_at: Option<Instant>,
 }
                     self.handle_search_queries_executed(info, window, cx);
                 }
             }
-            ZetaDebugInfo::SearchResultsFiltered(info) => {
-                if info.project == self.project {
-                    self.handle_search_results_filtered(info, window, cx);
-                }
-            }
             ZetaDebugInfo::ContextRetrievalFinished(info) => {
                 if info.project == self.project {
                     self.handle_context_retrieval_finished(info, window, cx);
                 }
             }
-            ZetaDebugInfo::EditPredicted(_) => {}
+            ZetaDebugInfo::EditPredictionRequested(_) => {}
         }
     }
             started_at: info.timestamp,
             search_results_generated_at: None,
             search_results_executed_at: None,
-            search_results_filtered_at: None,
             finished_at: None,
         });
         run.search_results_generated_at = Some(info.timestamp);
         run.search_queries = info
-            .queries
+            .regex_by_glob
             .into_iter()
-            .map(|query| {
+            .map(|(glob, regex)| {
                 let mut regex_parser = regex_syntax::ast::parse::Parser::new();
                 GlobQueries {
-                    glob: query.glob,
+                    glob,
-                    alternations: match regex_parser.parse(&query.regex) {
+                    alternations: match regex_parser.parse(&regex) {
                         Ok(regex_syntax::ast::Ast::Alternation(ref alt)) => {
                             alt.asts.iter().map(|ast| ast.to_string()).collect()
                         }
-                        _ => vec![query.regex],
+                        _ => vec![regex],
                     },
                 }
             })
         cx.notify();
     }
-    fn handle_search_results_filtered(
-        &mut self,
-        info: ZetaContextRetrievalDebugInfo,
-        _window: &mut Window,
-        cx: &mut Context<Self>,
-    ) {
-        let Some(run) = self.runs.back_mut() else {
-            return;
-        };
-        run.search_results_filtered_at = Some(info.timestamp);
-        cx.notify();
-    }
     fn handle_go_back(
         &mut self,
         _: &Zeta2ContextGoBack,
                         };
                         div = div.child(format!("Ran search: {:>5} ms", (t2 - t1).as_millis()));
-                        let Some(t3) = run.search_results_filtered_at else {
+                        div.child(format!(
-                            return pending_message(div, "Filtering results...");
+                            "Total: {:>5} ms",
-                        };
+                            (run.finished_at.unwrap_or(t0) - t0).as_millis()
-                        div =
+                        ))
-                            div.child(format!("Filtered results: {:>5} ms", (t3 - t2).as_millis()));
-                        let Some(t4) = run.finished_at else {
-                            return pending_message(div, "Building excerpts");
-                        };
-                        div = div
-                            .child(format!("Build excerpts: {:>5} µs", (t4 - t3).as_micros()))
-                            .child(format!("Total: {:>5} ms", (t4 - t0).as_millis()));
-                        div
                     }),
             )
     }

crates/zeta2_tools/src/zeta2_tools.rs+30−63

 use chrono::TimeDelta;
 use client::{Client, UserStore};
 use cloud_llm_client::predict_edits_v3::{
-    self, DeclarationScoreComponents, PredictEditsRequest, PredictEditsResponse, PromptFormat,
+    DeclarationScoreComponents, PredictEditsRequest, PromptFormat,
 };
 use collections::HashMap;
 use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
 use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
 use workspace::{Item, SplitDirection, Workspace};
 use zeta2::{
-    ContextMode, DEFAULT_SYNTAX_CONTEXT_OPTIONS, LlmContextOptions, Zeta, Zeta2FeatureFlag,
+    AgenticContextOptions, ContextMode, DEFAULT_SYNTAX_CONTEXT_OPTIONS, Zeta, Zeta2FeatureFlag,
     ZetaDebugInfo, ZetaEditPredictionDebugInfo, ZetaOptions,
 };
     context_editor: Entity<Editor>,
     prompt_editor: Entity<Editor>,
     retrieval_time: TimeDelta,
+    request_time: Option<TimeDelta>,
     buffer: WeakEntity<Buffer>,
     position: language::Anchor,
     state: LastPredictionState,
         model_response_editor: Entity<Editor>,
         feedback_editor: Entity<Editor>,
         feedback: Option<Feedback>,
-        response: predict_edits_v3::PredictEditsResponse,
+        request_id: String,
     },
     Failed {
         message: String,
         });
         match &options.context {
-            ContextMode::Llm(_) => {
+            ContextMode::Agentic(_) => {
                 self.context_mode = ContextModeState::Llm;
             }
             ContextMode::Syntax(_) => {
                 };
                 let context = match zeta_options.context {
-                    ContextMode::Llm(_context_options) => ContextMode::Llm(LlmContextOptions {
+                    ContextMode::Agentic(_context_options) => {
-                        excerpt: excerpt_options,
+                        ContextMode::Agentic(AgenticContextOptions {
-                    }),
+                            excerpt: excerpt_options,
+                        })
+                    }
                     ContextMode::Syntax(context_options) => {
                         let max_retrieved_declarations = match &this.context_mode {
                             ContextModeState::Llm => {
             let language_registry = self.project.read(cx).languages().clone();
             async move |this, cx| {
                 let mut languages = HashMap::default();
-                let ZetaDebugInfo::EditPredicted(prediction) = prediction else {
+                let ZetaDebugInfo::EditPredictionRequested(prediction) = prediction else {
                     return;
                 };
                 for ext in prediction
                     .await
                     .log_err();
+                let json_language = language_registry.language_for_name("Json").await.log_err();
                 this.update_in(cx, |this, window, cx| {
                     let context_editor = cx.new(|cx| {
                         let mut excerpt_score_components = HashMap::default();
                     let task = cx.spawn_in(window, {
                         let markdown_language = markdown_language.clone();
+                        let json_language = json_language.clone();
                         async move |this, cx| {
                             let response = response_rx.await;
                             this.update_in(cx, |this, window, cx| {
                                 if let Some(prediction) = this.last_prediction.as_mut() {
                                     prediction.state = match response {
-                                        Ok(Ok(response)) => {
+                                        Ok((Ok(response), request_time)) => {
-                                            if let Some(debug_info) = &response.debug_info {
+                                            prediction.request_time = Some(request_time);
-                                                prediction.prompt_editor.update(
-                                                    cx,
-                                                    |prompt_editor, cx| {
-                                                        prompt_editor.set_text(
-                                                            debug_info.prompt.as_str(),
-                                                            window,
-                                                            cx,
-                                                        );
-                                                    },
-                                                );
-                                            }
                                             let feedback_editor = cx.new(|cx| {
                                                 let buffer = cx.new(|cx| {
                                                 model_response_editor: cx.new(|cx| {
                                                     let buffer = cx.new(|cx| {
                                                         let mut buffer = Buffer::local(
-                                                            response
+                                                            serde_json::to_string_pretty(&response)
-                                                                .debug_info
+                                                                .unwrap_or_default(),
-                                                                .as_ref()
-                                                                .map(|p| p.model_response.as_str())
-                                                                .unwrap_or(
-                                                                    "(Debug info not available)",
-                                                                ),
                                                             cx,
                                                         );
-                                                        buffer.set_language(markdown_language, cx);
+                                                        buffer.set_language(json_language, cx);
                                                         buffer
                                                     });
                                                     let buffer = cx.new(|cx| {
                                                 }),
                                                 feedback_editor,
                                                 feedback: None,
-                                                response,
+                                                request_id: response.id.clone(),
                                             }
                                         }
-                                        Ok(Err(err)) => {
+                                        Ok((Err(err), request_time)) => {
+                                            prediction.request_time = Some(request_time);
                                             LastPredictionState::Failed { message: err }
                                         }
                                         Err(oneshot::Canceled) => LastPredictionState::Failed {
                             editor
                         }),
                         retrieval_time,
+                        request_time: None,
                         buffer,
                         position,
                         state: LastPredictionState::Requested,
                     feedback: feedback_state,
                     feedback_editor,
                     model_response_editor,
-                    response,
+                    request_id,
                     ..
                 } = &mut last_prediction.state
                 else {
                 telemetry::event!(
                     "Zeta2 Prediction Rated",
-                    id = response.request_id,
+                    id = request_id,
                     kind = kind,
                     text = text,
                     request = last_prediction.request,
-                    response = response,
                     project_snapshot = project_snapshot,
                 );
             })
                                             let current_options =
                                                 this.zeta.read(cx).options().clone();
                                             match current_options.context.clone() {
-                                                ContextMode::Llm(_) => {}
+                                                ContextMode::Agentic(_) => {}
                                                 ContextMode::Syntax(context_options) => {
                                                     let options = ZetaOptions {
-                                                        context: ContextMode::Llm(
+                                                        context: ContextMode::Agentic(
-                                                            LlmContextOptions {
+                                                            AgenticContextOptions {
                                                                 excerpt: context_options.excerpt,
                                                             },
                                                         ),
                                             let current_options =
                                                 this.zeta.read(cx).options().clone();
                                             match current_options.context.clone() {
-                                                ContextMode::Llm(context_options) => {
+                                                ContextMode::Agentic(context_options) => {
                                                     let options = ZetaOptions {
                                                         context: ContextMode::Syntax(
                                                             EditPredictionContextOptions {
             return None;
         };
-        let (prompt_planning_time, inference_time, parsing_time) =
-            if let LastPredictionState::Success {
-                response:
-                    PredictEditsResponse {
-                        debug_info: Some(debug_info),
-                        ..
-                    },
-                ..
-            } = &prediction.state
-            {
-                (
-                    Some(debug_info.prompt_planning_time),
-                    Some(debug_info.inference_time),
-                    Some(debug_info.parsing_time),
-                )
-            } else {
-                (None, None, None)
-            };
         Some(
             v_flex()
                 .p_4()
                     "Context retrieval",
                     Some(prediction.retrieval_time),
                 ))
-                .child(Self::render_duration(
+                .child(Self::render_duration("Request", prediction.request_time)),
-                    "Prompt planning",
-                    prompt_planning_time,
-                ))
-                .child(Self::render_duration("Inference", inference_time))
-                .child(Self::render_duration("Parsing", parsing_time)),
         )
     }

crates/zeta_cli/src/evaluate.rs+4−5

 use anyhow::Result;
 use clap::Args;
-use cloud_llm_client::udiff::DiffLine;
 use collections::HashSet;
 use gpui::AsyncApp;
+use zeta2::udiff::DiffLine;
 use crate::{
     example::{Example, NamedExample},
     headless::ZetaCliAppState,
+    paths::CACHE_DIR,
     predict::{PredictionDetails, zeta2_predict},
 };
     app_state: Arc<ZetaCliAppState>,
     cx: &mut AsyncApp,
 ) -> Result<EvaluationResult> {
-    let cache_dir = Path::new(&std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default())
-        .join("../../target/zeta-prediction-cache");
     let example = NamedExample::load(&example_path).unwrap();
-    let example_cache_path = cache_dir.join(&example_path.file_name().unwrap());
+    let example_cache_path = CACHE_DIR.join(&example_path.file_name().unwrap());
     let predictions = if !re_run && example_cache_path.exists() {
         let file_contents = fs::read_to_string(&example_cache_path)?;
     };
     if !example_cache_path.exists() {
-        fs::create_dir_all(&cache_dir).unwrap();
+        fs::create_dir_all(&*CACHE_DIR).unwrap();
         fs::write(
             example_cache_path,
             serde_json::to_string(&predictions).unwrap(),

crates/zeta_cli/src/example.rs+64−415

@@ -1,28 +1,31 @@
1	use std::{	1	use std::{
2	borrow::Cow,	2	borrow::Cow,
3	cell::RefCell,	3	cell::RefCell,
4	env,
5	fmt::{self, Display},	4	fmt::{self, Display},
6	fs,	5	fs,
7	io::Write,	6	io::Write,
8	mem,	7	mem,
9	ops::Range,
10	path::{Path, PathBuf},	8	path::{Path, PathBuf},
11	sync::Arc,	9	sync::Arc,
12	};	10	};
13		11
14	use anyhow::{Context as _, Result};	12	use anyhow::{Context as _, Result, anyhow};
15	use clap::ValueEnum;	13	use clap::ValueEnum;
16	use collections::{HashMap, HashSet};	14	use cloud_zeta2_prompt::CURSOR_MARKER;
		15	use collections::HashMap;
17	use futures::{	16	use futures::{
18	AsyncWriteExt as _,	17	AsyncWriteExt as _,
19	lock::{Mutex, OwnedMutexGuard},	18	lock::{Mutex, OwnedMutexGuard},
20	};	19	};
21	use gpui::{AsyncApp, Entity, http_client::Url};	20	use gpui::{AsyncApp, Entity, http_client::Url};
22	use language::Buffer;	21	use language::{Anchor, Buffer};
23	use project::{Project, ProjectPath};	22	use project::{Project, ProjectPath};
24	use pulldown_cmark::CowStr;	23	use pulldown_cmark::CowStr;
25	use serde::{Deserialize, Serialize};	24	use serde::{Deserialize, Serialize};
		25	use util::{paths::PathStyle, rel_path::RelPath};
		26	use zeta2::udiff::OpenedBuffers;
		27
		28	use crate::paths::{REPOS_DIR, WORKTREES_DIR};
26		29
27	const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";	30	const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
28	const EDIT_HISTORY_HEADING: &str = "Edit History";	31	const EDIT_HISTORY_HEADING: &str = "Edit History";
@@ -215,12 +218,10 @@ impl NamedExample { · 186 unmodified lines
215	let (repo_owner, repo_name) = self.repo_name()?;	218	let (repo_owner, repo_name) = self.repo_name()?;
216	let file_name = self.file_name();	219	let file_name = self.file_name();
217		220
218	let worktrees_dir = env::current_dir()?.join("target").join("zeta-worktrees");	221	fs::create_dir_all(&*REPOS_DIR)?;
219	let repos_dir = env::current_dir()?.join("target").join("zeta-repos");	222	fs::create_dir_all(&*WORKTREES_DIR)?;
220	fs::create_dir_all(&repos_dir)?;
221	fs::create_dir_all(&worktrees_dir)?;
222		223
223	let repo_dir = repos_dir.join(repo_owner.as_ref()).join(repo_name.as_ref());	224	let repo_dir = REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref());
224	let repo_lock = lock_repo(&repo_dir).await;	225	let repo_lock = lock_repo(&repo_dir).await;
225		226
226	if !repo_dir.is_dir() {	227	if !repo_dir.is_dir() {
@@ -251,7 +252,7 @@ impl NamedExample { · 24 unmodified lines
251	};	252	};
252		253
253	// Create the worktree for this example if needed.	254	// Create the worktree for this example if needed.
254	let worktree_path = worktrees_dir.join(&file_name);	255	let worktree_path = WORKTREES_DIR.join(&file_name);
255	if worktree_path.is_dir() {	256	if worktree_path.is_dir() {
256	run_git(&worktree_path, &["clean", "--force", "-d"]).await?;	257	run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
257	run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;	258	run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
@@ -309,7 +310,6 @@ impl NamedExample { · 51 unmodified lines
309	.collect()	310	.collect()
310	}	311	}
311		312
312	#[allow(unused)]
313	fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> {	313	fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> {
314	// git@github.com:owner/repo.git	314	// git@github.com:owner/repo.git
315	if self.example.repository_url.contains('@') {	315	if self.example.repository_url.contains('@') {
@@ -344,13 +344,63 @@ impl NamedExample { · 28 unmodified lines
344	}	344	}
345	}	345	}
346		346
		347	pub async fn cursor_position(
		348	&self,
		349	project: &Entity<Project>,
		350	cx: &mut AsyncApp,
		351	) -> Result<(Entity<Buffer>, Anchor)> {
		352	let worktree = project.read_with(cx, \|project, cx\| {
		353	project.visible_worktrees(cx).next().unwrap()
		354	})?;
		355	let cursor_path = RelPath::new(&self.example.cursor_path, PathStyle::Posix)?.into_arc();
		356	let cursor_buffer = project
		357	.update(cx, \|project, cx\| {
		358	project.open_buffer(
		359	ProjectPath {
		360	worktree_id: worktree.read(cx).id(),
		361	path: cursor_path,
		362	},
		363	cx,
		364	)
		365	})?
		366	.await?;
		367	let cursor_offset_within_excerpt = self
		368	.example
		369	.cursor_position
		370	.find(CURSOR_MARKER)
		371	.ok_or_else(\|\| anyhow!("missing cursor marker"))?;
		372	let mut cursor_excerpt = self.example.cursor_position.clone();
		373	cursor_excerpt.replace_range(
		374	cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
		375	"",
		376	);
		377	let excerpt_offset = cursor_buffer.read_with(cx, \|buffer, _cx\| {
		378	let text = buffer.text();
		379
		380	let mut matches = text.match_indices(&cursor_excerpt);
		381	let Some((excerpt_offset, _)) = matches.next() else {
		382	anyhow::bail!(
		383	"Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
		384	);
		385	};
		386	assert!(matches.next().is_none());
		387
		388	Ok(excerpt_offset)
		389	})??;
		390
		391	let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
		392	let cursor_anchor =
		393	cursor_buffer.read_with(cx, \|buffer, _\| buffer.anchor_after(cursor_offset))?;
		394	Ok((cursor_buffer, cursor_anchor))
		395	}
		396
347	#[must_use]	397	#[must_use]
348	pub async fn apply_edit_history(	398	pub async fn apply_edit_history(
349	&self,	399	&self,
350	project: &Entity<Project>,	400	project: &Entity<Project>,
351	cx: &mut AsyncApp,	401	cx: &mut AsyncApp,
352	) -> Result<HashSet<Entity<Buffer>>> {	402	) -> Result<OpenedBuffers<'_>> {
353	apply_diff(&self.example.edit_history, project, cx).await	403	zeta2::udiff::apply_diff(&self.example.edit_history, project, cx).await
354	}	404	}
355	}	405	}
356		406
@@ -446,404 +496,3 @@ pub async fn lock_repo(path: impl AsRef<Path>) -> OwnedMutexGuard<()> { · 89 unmodified lines
446	.lock_owned()	496	.lock_owned()
447	.await	497	.await
448	}	498	}
449
450	#[must_use]
451	pub async fn apply_diff(
452	diff: &str,
453	project: &Entity<Project>,
454	cx: &mut AsyncApp,
455	) -> Result<HashSet<Entity<Buffer>>> {
456	use cloud_llm_client::udiff::DiffLine;
457	use std::fmt::Write;
458
459	#[derive(Debug, Default)]
460	struct HunkState {
461	context: String,
462	edits: Vec<Edit>,
463	}
464
465	#[derive(Debug)]
466	struct Edit {
467	range: Range<usize>,
468	text: String,
469	}
470
471	let mut old_path = None;
472	let mut new_path = None;
473	let mut hunk = HunkState::default();
474	let mut diff_lines = diff.lines().map(DiffLine::parse).peekable();
475	let mut open_buffers = HashSet::default();
476
477	while let Some(diff_line) = diff_lines.next() {
478	match diff_line {
479	DiffLine::OldPath { path } => old_path = Some(path),
480	DiffLine::NewPath { path } => {
481	if old_path.is_none() {
482	anyhow::bail!(
483	"Found a new path header (`+++`) before an (`---`) old path header"
484	);
485	}
486	new_path = Some(path)
487	}
488	DiffLine::Context(ctx) => {
489	writeln!(&mut hunk.context, "{ctx}")?;
490	}
491	DiffLine::Deletion(del) => {
492	let range = hunk.context.len()..hunk.context.len() + del.len() + '\n'.len_utf8();
493	if let Some(last_edit) = hunk.edits.last_mut()
494	&& last_edit.range.end == range.start
495	{
496	last_edit.range.end = range.end;
497	} else {
498	hunk.edits.push(Edit {
499	range,
500	text: String::new(),
501	});
502	}
503	writeln!(&mut hunk.context, "{del}")?;
504	}
505	DiffLine::Addition(add) => {
506	let range = hunk.context.len()..hunk.context.len();
507	if let Some(last_edit) = hunk.edits.last_mut()
508	&& last_edit.range.end == range.start
509	{
510	writeln!(&mut last_edit.text, "{add}").unwrap();
511	} else {
512	hunk.edits.push(Edit {
513	range,
514	text: format!("{add}\n"),
515	});
516	}
517	}
518	DiffLine::HunkHeader(_) \| DiffLine::Garbage(_) => {}
519	}
520
521	let at_hunk_end = match diff_lines.peek() {
522	Some(DiffLine::OldPath { .. }) \| Some(DiffLine::HunkHeader(_)) \| None => true,
523	_ => false,
524	};
525
526	if at_hunk_end {
527	let hunk = mem::take(&mut hunk);
528
529	let Some(old_path) = old_path.as_deref() else {
530	anyhow::bail!("Missing old path (`---`) header")
531	};
532
533	let Some(new_path) = new_path.as_deref() else {
534	anyhow::bail!("Missing new path (`+++`) header")
535	};
536
537	let buffer = project
538	.update(cx, \|project, cx\| {
539	let project_path = project
540	.find_project_path(old_path, cx)
541	.context("Failed to find old_path in project")?;
542
543	anyhow::Ok(project.open_buffer(project_path, cx))
544	})??
545	.await?;
546	open_buffers.insert(buffer.clone());
547
548	if old_path != new_path {
549	project
550	.update(cx, \|project, cx\| {
551	let project_file = project::File::from_dyn(buffer.read(cx).file()).unwrap();
552	let new_path = ProjectPath {
553	worktree_id: project_file.worktree_id(cx),
554	path: project_file.path.clone(),
555	};
556	project.rename_entry(project_file.entry_id.unwrap(), new_path, cx)
557	})?
558	.await?;
559	}
560
561	// TODO is it worth using project search?
562	buffer.update(cx, \|buffer, cx\| {
563	let context_offset = if hunk.context.is_empty() {
564	0
565	} else {
566	let text = buffer.text();
567	if let Some(offset) = text.find(&hunk.context) {
568	if text[offset + 1..].contains(&hunk.context) {
569	anyhow::bail!("Context is not unique enough:\n{}", hunk.context);
570	}
571	offset
572	} else {
573	anyhow::bail!(
574	"Failed to match context:\n{}\n\nBuffer:\n{}",
575	hunk.context,
576	text
577	);
578	}
579	};
580
581	buffer.edit(
582	hunk.edits.into_iter().map(\|edit\| {
583	(
584	context_offset + edit.range.start..context_offset + edit.range.end,
585	edit.text,
586	)
587	}),
588	None,
589	cx,
590	);
591
592	anyhow::Ok(())
593	})??;
594	}
595	}
596
597	anyhow::Ok(open_buffers)
598	}
599
600	#[cfg(test)]
601	mod tests {
602	use super::*;
603	use ::fs::FakeFs;
604	use gpui::TestAppContext;
605	use indoc::indoc;
606	use pretty_assertions::assert_eq;
607	use project::Project;
608	use serde_json::json;
609	use settings::SettingsStore;
610	use util::path;
611
612	#[gpui::test]
613	async fn test_apply_diff_successful(cx: &mut TestAppContext) {
614	let buffer_1_text = indoc! {r#"
615	one
616	two
617	three
618	four
619	five
620	"# };
621
622	let buffer_1_text_final = indoc! {r#"
623	3
624	4
625	5
626	"# };
627
628	let buffer_2_text = indoc! {r#"
629	six
630	seven
631	eight
632	nine
633	ten
634	"# };
635
636	let buffer_2_text_final = indoc! {r#"
637	5
638	six
639	seven
640	7.5
641	eight
642	nine
643	ten
644	11
645	"# };
646
647	cx.update(\|cx\| {
648	let settings_store = SettingsStore::test(cx);
649	cx.set_global(settings_store);
650	Project::init_settings(cx);
651	language::init(cx);
652	});
653
654	let fs = FakeFs::new(cx.background_executor.clone());
655	fs.insert_tree(
656	path!("/root"),
657	json!({
658	"file1": buffer_1_text,
659	"file2": buffer_2_text,
660	}),
661	)
662	.await;
663
664	let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
665
666	let diff = indoc! {r#"
667	--- a/root/file1
668	+++ b/root/file1
669	one
670	two
671	-three
672	+3
673	four
674	five
675	--- a/root/file1
676	+++ b/root/file1
677	3
678	-four
679	-five
680	+4
681	+5
682	--- a/root/file1
683	+++ b/root/file1
684	-one
685	-two
686	3
687	4
688	--- a/root/file2
689	+++ b/root/file2
690	+5
691	six
692	--- a/root/file2
693	+++ b/root/file2
694	seven
695	+7.5
696	eight
697	--- a/root/file2
698	+++ b/root/file2
699	ten
700	+11
701	"#};
702
703	let _buffers = apply_diff(diff, &project, &mut cx.to_async())
704	.await
705	.unwrap();
706	let buffer_1 = project
707	.update(cx, \|project, cx\| {
708	let project_path = project.find_project_path(path!("/root/file1"), cx).unwrap();
709	project.open_buffer(project_path, cx)
710	})
711	.await
712	.unwrap();
713
714	buffer_1.read_with(cx, \|buffer, _cx\| {
715	assert_eq!(buffer.text(), buffer_1_text_final);
716	});
717	let buffer_2 = project
718	.update(cx, \|project, cx\| {
719	let project_path = project.find_project_path(path!("/root/file2"), cx).unwrap();
720	project.open_buffer(project_path, cx)
721	})
722	.await
723	.unwrap();
724
725	buffer_2.read_with(cx, \|buffer, _cx\| {
726	assert_eq!(buffer.text(), buffer_2_text_final);
727	});
728	}
729
730	#[gpui::test]
731	async fn test_apply_diff_non_unique(cx: &mut TestAppContext) {
732	let buffer_1_text = indoc! {r#"
733	one
734	two
735	three
736	four
737	five
738	one
739	two
740	three
741	four
742	five
743	"# };
744
745	cx.update(\|cx\| {
746	let settings_store = SettingsStore::test(cx);
747	cx.set_global(settings_store);
748	Project::init_settings(cx);
749	language::init(cx);
750	});
751
752	let fs = FakeFs::new(cx.background_executor.clone());
753	fs.insert_tree(
754	path!("/root"),
755	json!({
756	"file1": buffer_1_text,
757	}),
758	)
759	.await;
760
761	let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
762
763	let diff = indoc! {r#"
764	--- a/root/file1
765	+++ b/root/file1
766	one
767	two
768	-three
769	+3
770	four
771	five
772	"#};
773
774	apply_diff(diff, &project, &mut cx.to_async())
775	.await
776	.expect_err("Non-unique edits should fail");
777	}
778
779	#[gpui::test]
780	async fn test_apply_diff_unique_via_previous_context(cx: &mut TestAppContext) {
781	let start = indoc! {r#"
782	one
783	two
784	three
785	four
786	five
787
788	four
789	five
790	"# };
791
792	let end = indoc! {r#"
793	one
794	two
795	3
796	four
797	5
798
799	four
800	five
801	"# };
802
803	cx.update(\|cx\| {
804	let settings_store = SettingsStore::test(cx);
805	cx.set_global(settings_store);
806	Project::init_settings(cx);
807	language::init(cx);
808	});
809
810	let fs = FakeFs::new(cx.background_executor.clone());
811	fs.insert_tree(
812	path!("/root"),
813	json!({
814	"file1": start,
815	}),
816	)
817	.await;
818
819	let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
820
821	let diff = indoc! {r#"
822	--- a/root/file1
823	+++ b/root/file1
824	one
825	two
826	-three
827	+3
828	four
829	-five
830	+5
831	"#};
832
833	let _buffers = apply_diff(diff, &project, &mut cx.to_async())
834	.await
835	.unwrap();
836
837	let buffer_1 = project
838	.update(cx, \|project, cx\| {
839	let project_path = project.find_project_path(path!("/root/file1"), cx).unwrap();
840	project.open_buffer(project_path, cx)
841	})
842	.await
843	.unwrap();
844
845	buffer_1.read_with(cx, \|buffer, _cx\| {
846	assert_eq!(buffer.text(), end);
847	});
848	}
849	}

crates/zeta_cli/src/main.rs+6−223

@@ -1,6 +1,7 @@
1	mod evaluate;	1	mod evaluate;
2	mod example;	2	mod example;
3	mod headless;	3	mod headless;
		4	mod paths;
4	mod predict;	5	mod predict;
5	mod source_location;	6	mod source_location;
6	mod syntax_retrieval_stats;	7	mod syntax_retrieval_stats;
@@ -10,28 +11,22 @@ use crate::evaluate::{EvaluateArguments, run_evaluate}; · 3 unmodified lines
10	use crate::example::{ExampleFormat, NamedExample};	11	use crate::example::{ExampleFormat, NamedExample};
11	use crate::predict::{PredictArguments, run_zeta2_predict};	12	use crate::predict::{PredictArguments, run_zeta2_predict};
12	use crate::syntax_retrieval_stats::retrieval_stats;	13	use crate::syntax_retrieval_stats::retrieval_stats;
13	use ::serde::Serialize;
14	use ::util::paths::PathStyle;	14	use ::util::paths::PathStyle;
15	use anyhow::{Context as _, Result, anyhow};	15	use anyhow::{Result, anyhow};
16	use clap::{Args, Parser, Subcommand};	16	use clap::{Args, Parser, Subcommand};
17	use cloud_llm_client::predict_edits_v3::{self, Excerpt};	17	use cloud_llm_client::predict_edits_v3;
18	use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
19	use edit_prediction_context::{	18	use edit_prediction_context::{
20	EditPredictionContextOptions, EditPredictionExcerpt, EditPredictionExcerptOptions,	19	EditPredictionContextOptions, EditPredictionExcerptOptions, EditPredictionScoreOptions,
21	EditPredictionScoreOptions, Line,
22	};	20	};
23	use futures::StreamExt as _;
24	use futures::channel::mpsc;
25	use gpui::{Application, AsyncApp, Entity, prelude::*};	21	use gpui::{Application, AsyncApp, Entity, prelude::*};
26	use language::{Bias, Buffer, BufferSnapshot, OffsetRangeExt, Point};	22	use language::{Bias, Buffer, BufferSnapshot, Point};
27	use language_model::LanguageModelRegistry;
28	use project::{Project, Worktree};	23	use project::{Project, Worktree};
29	use reqwest_client::ReqwestClient;	24	use reqwest_client::ReqwestClient;
30	use serde_json::json;	25	use serde_json::json;
31	use std::io::{self};	26	use std::io::{self};
32	use std::time::Duration;	27	use std::time::Duration;
33	use std::{collections::HashSet, path::PathBuf, str::FromStr, sync::Arc};	28	use std::{collections::HashSet, path::PathBuf, str::FromStr, sync::Arc};
34	use zeta2::{ContextMode, LlmContextOptions, SearchToolQuery};	29	use zeta2::ContextMode;
35		30
36	use crate::headless::ZetaCliAppState;	31	use crate::headless::ZetaCliAppState;
37	use crate::source_location::SourceLocation;	32	use crate::source_location::SourceLocation;
@@ -79,12 +74,6 @@ enum Zeta2Command { · 41 unmodified lines
79	#[command(subcommand)]	74	#[command(subcommand)]
80	command: Zeta2SyntaxCommand,	75	command: Zeta2SyntaxCommand,
81	},	76	},
82	Llm {
83	#[clap(flatten)]
84	args: Zeta2Args,
85	#[command(subcommand)]
86	command: Zeta2LlmCommand,
87	},
88	Predict(PredictArguments),	77	Predict(PredictArguments),
89	Eval(EvaluateArguments),	78	Eval(EvaluateArguments),
90	}	79	}
@@ -107,14 +96,6 @@ enum Zeta2SyntaxCommand { · 16 unmodified lines
107	},	96	},
108	}	97	}
109		98
110	#[derive(Subcommand, Debug)]
111	enum Zeta2LlmCommand {
112	Context {
113	#[clap(flatten)]
114	context_args: ContextArgs,
115	},
116	}
117
118	#[derive(Debug, Args)]	99	#[derive(Debug, Args)]
119	#[group(requires = "worktree")]	100	#[group(requires = "worktree")]
120	struct ContextArgs {	101	struct ContextArgs {
@@ -388,197 +369,6 @@ async fn zeta2_syntax_context( · 267 unmodified lines
388	Ok(output)	369	Ok(output)
389	}	370	}
390		371
391	async fn zeta2_llm_context(
392	zeta2_args: Zeta2Args,
393	context_args: ContextArgs,
394	app_state: &Arc<ZetaCliAppState>,
395	cx: &mut AsyncApp,
396	) -> Result<String> {
397	let LoadedContext {
398	buffer,
399	clipped_cursor,
400	snapshot: cursor_snapshot,
401	project,
402	..
403	} = load_context(&context_args, app_state, cx).await?;
404
405	let cursor_position = cursor_snapshot.anchor_after(clipped_cursor);
406
407	cx.update(\|cx\| {
408	LanguageModelRegistry::global(cx).update(cx, \|registry, cx\| {
409	registry
410	.provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
411	.unwrap()
412	.authenticate(cx)
413	})
414	})?
415	.await?;
416
417	let edit_history_unified_diff = match context_args.edit_history {
418	Some(events) => events.read_to_string().await?,
419	None => String::new(),
420	};
421
422	let (debug_tx, mut debug_rx) = mpsc::unbounded();
423
424	let excerpt_options = EditPredictionExcerptOptions {
425	max_bytes: zeta2_args.max_excerpt_bytes,
426	min_bytes: zeta2_args.min_excerpt_bytes,
427	target_before_cursor_over_total_bytes: zeta2_args.target_before_cursor_over_total_bytes,
428	};
429
430	let related_excerpts = cx
431	.update(\|cx\| {
432	zeta2::related_excerpts::find_related_excerpts(
433	buffer,
434	cursor_position,
435	&project,
436	edit_history_unified_diff,
437	&LlmContextOptions {
438	excerpt: excerpt_options.clone(),
439	},
440	Some(debug_tx),
441	cx,
442	)
443	})?
444	.await?;
445
446	let cursor_excerpt = EditPredictionExcerpt::select_from_buffer(
447	clipped_cursor,
448	&cursor_snapshot,
449	&excerpt_options,
450	None,
451	)
452	.context("line didn't fit")?;
453
454	#[derive(Serialize)]
455	struct Output {
456	excerpts: Vec<OutputExcerpt>,
457	formatted_excerpts: String,
458	meta: OutputMeta,
459	}
460
461	#[derive(Default, Serialize)]
462	struct OutputMeta {
463	search_prompt: String,
464	search_queries: Vec<SearchToolQuery>,
465	}
466
467	#[derive(Serialize)]
468	struct OutputExcerpt {
469	path: PathBuf,
470	#[serde(flatten)]
471	excerpt: Excerpt,
472	}
473
474	let mut meta = OutputMeta::default();
475
476	while let Some(debug_info) = debug_rx.next().await {
477	match debug_info {
478	zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
479	meta.search_prompt = info.search_prompt;
480	}
481	zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
482	meta.search_queries = info.queries
483	}
484	_ => {}
485	}
486	}
487
488	cx.update(\|cx\| {
489	let mut excerpts = Vec::new();
490	let mut formatted_excerpts = String::new();
491
492	let cursor_insertions = [(
493	predict_edits_v3::Point {
494	line: Line(clipped_cursor.row),
495	column: clipped_cursor.column,
496	},
497	CURSOR_MARKER,
498	)];
499
500	let mut cursor_excerpt_added = false;
501
502	for (buffer, ranges) in related_excerpts {
503	let excerpt_snapshot = buffer.read(cx).snapshot();
504
505	let mut line_ranges = ranges
506	.into_iter()
507	.map(\|range\| {
508	let point_range = range.to_point(&excerpt_snapshot);
509	Line(point_range.start.row)..Line(point_range.end.row)
510	})
511	.collect::<Vec<_>>();
512
513	let Some(file) = excerpt_snapshot.file() else {
514	continue;
515	};
516	let path = file.full_path(cx);
517
518	let is_cursor_file = path == cursor_snapshot.file().unwrap().full_path(cx);
519	if is_cursor_file {
520	let insertion_ix = line_ranges
521	.binary_search_by(\|probe\| {
522	probe
523	.start
524	.cmp(&cursor_excerpt.line_range.start)
525	.then(cursor_excerpt.line_range.end.cmp(&probe.end))
526	})
527	.unwrap_or_else(\|ix\| ix);
528	line_ranges.insert(insertion_ix, cursor_excerpt.line_range.clone());
529	cursor_excerpt_added = true;
530	}
531
532	let merged_excerpts =
533	zeta2::merge_excerpts::merge_excerpts(&excerpt_snapshot, line_ranges)
534	.into_iter()
535	.map(\|excerpt\| OutputExcerpt {
536	path: path.clone(),
537	excerpt,
538	});
539
540	let excerpt_start_ix = excerpts.len();
541	excerpts.extend(merged_excerpts);
542
543	write_codeblock(
544	&path,
545	excerpts[excerpt_start_ix..].iter().map(\|e\| &e.excerpt),
546	if is_cursor_file {
547	&cursor_insertions
548	} else {
549	&[]
550	},
551	Line(excerpt_snapshot.max_point().row),
552	true,
553	&mut formatted_excerpts,
554	);
555	}
556
557	if !cursor_excerpt_added {
558	write_codeblock(
559	&cursor_snapshot.file().unwrap().full_path(cx),
560	&[Excerpt {
561	start_line: cursor_excerpt.line_range.start,
562	text: cursor_excerpt.text(&cursor_snapshot).body.into(),
563	}],
564	&cursor_insertions,
565	Line(cursor_snapshot.max_point().row),
566	true,
567	&mut formatted_excerpts,
568	);
569	}
570
571	let output = Output {
572	excerpts,
573	formatted_excerpts,
574	meta,
575	};
576
577	Ok(serde_json::to_string_pretty(&output)?)
578	})
579	.unwrap()
580	}
581
582	async fn zeta1_context(	372	async fn zeta1_context(
583	args: ContextArgs,	373	args: ContextArgs,
584	app_state: &Arc<ZetaCliAppState>,	374	app_state: &Arc<ZetaCliAppState>,
@@ -670,13 +460,6 @@ fn main() { · 85 unmodified lines
670	};	460	};
671	println!("{}", result.unwrap());	461	println!("{}", result.unwrap());
672	}	462	}
673	Zeta2Command::Llm { args, command } => match command {
674	Zeta2LlmCommand::Context { context_args } => {
675	let result =
676	zeta2_llm_context(args, context_args, &app_state, cx).await;
677	println!("{}", result.unwrap());
678	}
679	},
680	},	463	},
681	Command::ConvertExample {	464	Command::ConvertExample {
682	path,	465	path,

crates/zeta_cli/src/paths.rs+8−0

@@ -0,0 +1,8 @@
	1	use std::{env, path::PathBuf, sync::LazyLock};
	2
	3	static TARGET_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| env::current_dir().unwrap().join("target"));
	4	pub static CACHE_DIR: LazyLock<PathBuf> =
	5	LazyLock::new(\|\| TARGET_DIR.join("zeta-prediction-cache"));
	6	pub static REPOS_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| TARGET_DIR.join("zeta-repos"));
	7	pub static WORKTREES_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| TARGET_DIR.join("zeta-worktrees"));
	8	pub static LOGS_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| TARGET_DIR.join("zeta-logs"));

crates/zeta_cli/src/predict.rs+39−62

 use crate::example::{ActualExcerpt, NamedExample};
 use crate::headless::ZetaCliAppState;
+use crate::paths::LOGS_DIR;
 use ::serde::Serialize;
-use ::util::paths::PathStyle;
 use anyhow::{Context as _, Result, anyhow};
 use clap::Args;
 use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
 use futures::StreamExt as _;
 use gpui::AsyncApp;
-use language_model::LanguageModelRegistry;
+use project::Project;
-use project::{Project, ProjectPath};
 use serde::Deserialize;
 use std::cell::Cell;
+use std::fs;
 use std::io::Write;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
-use util::rel_path::RelPath;
 #[derive(Debug, Args)]
 pub struct PredictArguments {
     app_state: &Arc<ZetaCliAppState>,
     cx: &mut AsyncApp,
 ) -> Result<PredictionDetails> {
+    fs::create_dir_all(&*LOGS_DIR)?;
     let worktree_path = example.setup_worktree().await?;
     if !AUTHENTICATED.get() {
         AUTHENTICATED.set(true);
-        cx.update(|cx| {
-            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
-                registry
-                    .provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
-                    .unwrap()
-                    .authenticate(cx)
-            })
-        })?
-        .await?;
         app_state
             .client
             .sign_in_with_optional_connect(true, cx)
         )
     })?;
+    let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
     let worktree = project
         .update(cx, |project, cx| {
             project.create_worktree(&worktree_path, true, cx)
         })?
         .await;
-    let _edited_buffers = example.apply_edit_history(&project, cx).await?;
+    let zeta = cx.update(|cx| zeta2::Zeta::global(&app_state.client, &app_state.user_store, cx))?;
-    let cursor_path = RelPath::new(&example.example.cursor_path, PathStyle::Posix)?.into_arc();
-    let cursor_buffer = project
-        .update(cx, |project, cx| {
-            project.open_buffer(
-                ProjectPath {
-                    worktree_id: worktree.read(cx).id(),
-                    path: cursor_path,
-                },
-                cx,
-            )
-        })?
-        .await?;
-    let cursor_offset_within_excerpt = example
-        .example
-        .cursor_position
-        .find(CURSOR_MARKER)
-        .ok_or_else(|| anyhow!("missing cursor marker"))?;
-    let mut cursor_excerpt = example.example.cursor_position.clone();
-    cursor_excerpt.replace_range(
-        cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
-        "",
-    );
-    let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
-        let text = buffer.text();
-        let mut matches = text.match_indices(&cursor_excerpt);
-        let Some((excerpt_offset, _)) = matches.next() else {
-            anyhow::bail!(
-                "Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
-            );
-        };
-        assert!(matches.next().is_none());
-        Ok(excerpt_offset)
+    cx.subscribe(&buffer_store, {
-    })??;
+        let project = project.clone();
+        move |_, event, cx| match event {
+            project::buffer_store::BufferStoreEvent::BufferAdded(buffer) => {
+                zeta2::Zeta::try_global(cx)
+                    .unwrap()
+                    .update(cx, |zeta, cx| zeta.register_buffer(&buffer, &project, cx));
+            }
+            _ => {}
+        }
+    })?
+    .detach();
-    let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
+    let _edited_buffers = example.apply_edit_history(&project, cx).await?;
-    let cursor_anchor =
+    let (cursor_buffer, cursor_anchor) = example.cursor_position(&project, cx).await?;
-        cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
-    let zeta = cx.update(|cx| zeta2::Zeta::global(&app_state.client, &app_state.user_store, cx))?;
+    let mut debug_rx = zeta.update(cx, |zeta, _| zeta.debug_info())?;
     let refresh_task = zeta.update(cx, |zeta, cx| {
-        zeta.register_buffer(&cursor_buffer, &project, cx);
         zeta.refresh_context(project.clone(), cursor_buffer.clone(), cursor_anchor, cx)
     })?;
-    let mut debug_rx = zeta.update(cx, |zeta, _| zeta.debug_info())?;
     let mut context_retrieval_started_at = None;
     let mut context_retrieval_finished_at = None;
     let mut search_queries_generated_at = None;
         match event {
             zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
                 context_retrieval_started_at = Some(info.timestamp);
+                fs::write(LOGS_DIR.join("search_prompt.md"), &info.search_prompt)?;
             }
             zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
                 search_queries_generated_at = Some(info.timestamp);
+                fs::write(
+                    LOGS_DIR.join("search_queries.json"),
+                    serde_json::to_string_pretty(&info.regex_by_glob).unwrap(),
+                )?;
             }
             zeta2::ZetaDebugInfo::SearchQueriesExecuted(info) => {
                 search_queries_executed_at = Some(info.timestamp);
                     zeta.request_prediction(&project, &cursor_buffer, cursor_anchor, cx)
                 })?);
             }
-            zeta2::ZetaDebugInfo::EditPredicted(request) => {
+            zeta2::ZetaDebugInfo::EditPredictionRequested(request) => {
                 prediction_started_at = Some(Instant::now());
-                request.response_rx.await?.map_err(|err| anyhow!(err))?;
+                fs::write(
+                    LOGS_DIR.join("prediction_prompt.md"),
+                    &request.local_prompt.unwrap_or_default(),
+                )?;
+                let response = request.response_rx.await?.0.map_err(|err| anyhow!(err))?;
                 prediction_finished_at = Some(Instant::now());
+                fs::write(
+                    LOGS_DIR.join("prediction_response.json"),
+                    &serde_json::to_string_pretty(&response).unwrap(),
+                )?;
                 for included_file in request.request.included_files {
                     let insertions = vec![(request.request.cursor_point, CURSOR_MARKER)];
                     result
                 }
                 break;
             }
-            _ => {}
         }
     }

@@ -0,0 +1,8 @@
	1	use std::{env, path::PathBuf, sync::LazyLock};
	2
	3	static TARGET_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| env::current_dir().unwrap().join("target"));
	4	pub static CACHE_DIR: LazyLock<PathBuf> =
	5	LazyLock::new(\|\| TARGET_DIR.join("zeta-prediction-cache"));
	6	pub static REPOS_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| TARGET_DIR.join("zeta-repos"));
	7	pub static WORKTREES_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| TARGET_DIR.join("zeta-worktrees"));
	8	pub static LOGS_DIR: LazyLock<PathBuf> = LazyLock::new(\|\| TARGET_DIR.join("zeta-logs"));