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

@@ -39,6 +39,7 @@ dependencies = [ · 38 unmodified lines
39	39	"util",
40	40	"uuid",
41	41	"watch",
	42	+ "zlog",
42	43	]
43	44
44	45	[[package]]
@@ -3198,7 +3199,9 @@ dependencies = [ · 3153 unmodified lines
3198	3199	"indoc",
3199	3200	"ordered-float 2.10.1",
3200	3201	"rustc-hash 2.1.1",
	3202	+ "schemars 1.0.4",
3201	3203	"serde",
	3204	+ "serde_json",
3202	3205	"strum 0.27.2",
3203	3206	]
3204	3207
@@ -21675,10 +21678,10 @@ dependencies = [ · 18470 unmodified lines
21675	21678	"language_model",
21676	21679	"log",
21677	21680	"lsp",
	21681	+ "open_ai",
21678	21682	"pretty_assertions",
21679	21683	"project",
21680	21684	"release_channel",
21681		− "schemars 1.0.4",
21682	21685	"serde",
21683	21686	"serde_json",
21684	21687	"settings",
@@ -21687,6 +21690,7 @@ dependencies = [ · 2 unmodified lines
21687	21690	"uuid",
21688	21691	"workspace",
21689	21692	"worktree",
	21693	+ "zlog",
21690	21694	]
21691	21695
21692	21696	[[package]]

crates/acp_thread/Cargo.toml+1−0

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

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

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

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

@@ -1,7 +1,7 @@
1	1	use chrono::Duration;
2	2	use serde::{Deserialize, Serialize};
3	3	use std::{
4		− fmt::Display,
	4	+ fmt::{Display, Write as _},
5	5	ops::{Add, Range, Sub},
6	6	path::{Path, PathBuf},
7	7	sync::Arc,
@@ -11,7 +11,14 @@ use uuid::Uuid; · 3 unmodified lines
11	11
12	12	use crate::PredictEditsGitInfo;
13	13
14		−// TODO: snippet ordering within file / relative to excerpt
	14	+#[derive(Debug, Clone, Serialize, Deserialize)]
	15	+pub struct PlanContextRetrievalRequest {
	16	+ pub excerpt: String,
	17	+ pub excerpt_path: Arc<Path>,
	18	+ pub excerpt_line_range: Range<Line>,
	19	+ pub cursor_file_max_row: Line,
	20	+ pub events: Vec<Event>,
	21	+}
15	22
16	23	#[derive(Debug, Clone, Serialize, Deserialize)]
17	24	pub struct PredictEditsRequest {
@@ -125,15 +132,15 @@ impl Display for Event { · 107 unmodified lines
125	132	write!(
126	133	f,
127	134	"// User accepted prediction:\n--- a/{}\n+++ b/{}\n{diff}",
128		− old_path.display(),
129		− new_path.display()
	135	+ DiffPathFmt(old_path),
	136	+ DiffPathFmt(new_path)
130	137	)
131	138	} else {
132	139	write!(
133	140	f,
134	141	"--- a/{}\n+++ b/{}\n{diff}",
135		− old_path.display(),
136		− new_path.display()
	142	+ DiffPathFmt(old_path),
	143	+ DiffPathFmt(new_path)
137	144	)
138	145	}
139	146	}
@@ -141,6 +148,24 @@ impl Display for Event { · 1 unmodified line
141	148	}
142	149	}
143	150
	151	+/// always format the Path as a unix path with `/` as the path sep in Diffs
	152	+pub struct DiffPathFmt<'a>(pub &'a Path);
	153	+
	154	+impl<'a> std::fmt::Display for DiffPathFmt<'a> {
	155	+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
	156	+ let mut is_first = true;
	157	+ for component in self.0.components() {
	158	+ if !is_first {
	159	+ f.write_char('/')?;
	160	+ } else {
	161	+ is_first = false;
	162	+ }
	163	+ write!(f, "{}", component.as_os_str().display())?;
	164	+ }
	165	+ Ok(())
	166	+ }
	167	+}
	168	+
144	169	#[derive(Debug, Clone, Serialize, Deserialize)]
145	170	pub struct Signature {
146	171	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	17	indoc.workspace = true
18	18	ordered-float.workspace = true
19	19	rustc-hash.workspace = true
	20	+schemars.workspace = true
20	21	serde.workspace = true
	22	+serde_json.workspace = true
21	23	strum.workspace = true

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

@@ -1,8 +1,9 @@
1	1	//! Zeta2 prompt planning and generation code shared with cloud.
	2	+pub mod retrieval_prompt;
2	3
3	4	use anyhow::{Context as _, Result, anyhow};
4	5	use cloud_llm_client::predict_edits_v3::{
5		− self, Excerpt, Line, Point, PromptFormat, ReferencedDeclaration,
	6	+ self, DiffPathFmt, Excerpt, Line, Point, PromptFormat, ReferencedDeclaration,
6	7	};
7	8	use indoc::indoc;
8	9	use ordered_float::OrderedFloat;
@@ -212,7 +213,7 @@ pub fn write_codeblock<'a>( · 203 unmodified lines
212	213	include_line_numbers: bool,
213	214	output: &'a mut String,
214	215	) {
215		− writeln!(output, "`````{}", path.display()).unwrap();
	216	+ writeln!(output, "`````{}", DiffPathFmt(path)).unwrap();
216	217	write_excerpts(
217	218	excerpts,
218	219	sorted_insertions,
@@ -275,7 +276,7 @@ pub fn write_excerpts<'a>( · 56 unmodified lines
275	276	}
276	277	}
277	278
278		−fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
	279	+pub fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
279	280	if events.is_empty() {
280	281	return;
281	282	};

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	34	snapshot: BufferSnapshot,
35	35	/// The edits that should be applied to transform the original text into the predicted text.
36	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	38	/// Preview of how the buffer will look after applying the edits.
39	39	edit_preview: EditPreview,
40	40	}
@@ -42,7 +42,7 @@ struct CurrentCompletion { · 1 unmodified line
42	42	impl CurrentCompletion {
43	43	/// Attempts to adjust the edits based on changes made to the buffer since the completion was generated.
44	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	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	281	return Ok(());
282	282	}
283	283
284		− let edits: Arc<[(Range<Anchor>, String)]> =
285		− vec![(cursor_position..cursor_position, completion_text)].into();
	284	+ let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
	285	+ vec![(cursor_position..cursor_position, completion_text.into())].into();
286	286	let edit_preview = buffer
287	287	.read_with(cx, \|buffer, cx\| buffer.preview_edits(edits.clone(), cx))?
288	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	3	use client::EditPredictionUsage;
4	4	use gpui::{App, Context, Entity, SharedString};
@@ -19,7 +19,7 @@ pub enum EditPrediction { · 14 unmodified lines
19	19	/// Edits within the buffer that requested the prediction
20	20	Local {
21	21	id: Option<SharedString>,
22		− edits: Vec<(Range<language::Anchor>, String)>,
	22	+ edits: Vec<(Range<language::Anchor>, Arc<str>)>,
23	23	edit_preview: Option<language::EditPreview>,
24	24	},
25	25	/// Jump to a different file from the one that requested the prediction
@@ -248,8 +248,8 @@ where · 222 unmodified lines
248	248	pub fn interpolate_edits(
249	249	old_snapshot: &BufferSnapshot,
250	250	new_snapshot: &BufferSnapshot,
251		− current_edits: &[(Range<Anchor>, String)],
252		−) -> Option<Vec<(Range<Anchor>, String)>> {
	251	+ current_edits: &[(Range<Anchor>, Arc<str>)],
	252	+) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
253	253	let mut edits = Vec::new();
254	254
255	255	let mut model_edits = current_edits.iter().peekable();
@@ -274,7 +274,7 @@ pub fn interpolate_edits( · 18 unmodified lines
274	274	if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
275	275	if !model_suffix.is_empty() {
276	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	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	2	use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
3	3	use indoc::indoc;
4	4	use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
5		−use std::ops::Range;
	5	+use std::{ops::Range, sync::Arc};
6	6	use text::{Point, ToOffset};
7	7
8	8	use crate::{
@@ -24,7 +24,7 @@ async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) { · 15 unmodified lines
24	24
25	25	assert_editor_active_edit_completion(&mut cx, \|_, edits\| {
26	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	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	47	assert_editor_active_edit_completion(&mut cx, \|_, edits\| {
48	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	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	331	fn assert_editor_active_edit_completion(
332	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	335	cx.editor(\|editor, _, cx\| {
336	336	let completion_state = editor

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

@@ -616,7 +616,7 @@ pub(crate) enum EditDisplayMode { · 615 unmodified lines
616	616
617	617	enum EditPrediction {
618	618	Edit {
619		− edits: Vec<(Range<Anchor>, String)>,
	619	+ edits: Vec<(Range<Anchor>, Arc<str>)>,
620	620	edit_preview: Option<EditPreview>,
621	621	display_mode: EditDisplayMode,
622	622	snapshot: BufferSnapshot,
@@ -7960,7 +7960,7 @@ impl Editor { · 7337 unmodified lines
7960	7960	let inlay = Inlay::edit_prediction(
7961	7961	post_inc(&mut self.next_inlay_id),
7962	7962	range.start,
7963		− new_text.as_str(),
	7963	+ new_text.as_ref(),
7964	7964	);
7965	7965	inlay_ids.push(inlay.id);
7966	7966	inlays.push(inlay);
@@ -8982,7 +8982,7 @@ impl Editor { · 1015 unmodified lines
8982	8982	newest_selection_head: Option<DisplayPoint>,
8983	8983	editor_width: Pixels,
8984	8984	style: &EditorStyle,
8985		− edits: &Vec<(Range<Anchor>, String)>,
	8985	+ edits: &Vec<(Range<Anchor>, Arc<str>)>,
8986	8986	edit_preview: &Option<language::EditPreview>,
8987	8987	snapshot: &language::BufferSnapshot,
8988	8988	window: &mut Window,
@@ -24382,25 +24382,20 @@ impl InvalidationRegion for SnippetState { · 15393 unmodified lines
24382	24382
24383	24383	fn edit_prediction_edit_text(
24384	24384	current_snapshot: &BufferSnapshot,
24385		− edits: &[(Range<Anchor>, String)],
	24385	+ edits: &[(Range<Anchor>, impl AsRef<str>)],
24386	24386	edit_preview: &EditPreview,
24387	24387	include_deletions: bool,
24388	24388	cx: &App,
24389	24389	) -> HighlightedText {
24390	24390	let edits = edits
24391	24391	.iter()
24392		− .map(\|(anchor, text)\| {
24393		− (
24394		− anchor.start.text_anchor..anchor.end.text_anchor,
24395		− text.clone(),
24396		− )
24397		− })
	24392	+ .map(\|(anchor, text)\| (anchor.start.text_anchor..anchor.end.text_anchor, text))
24398	24393	.collect::<Vec<_>>();
24399	24394
24400	24395	edit_preview.highlight_edits(current_snapshot, &edits, include_deletions, cx)
24401	24396	}
24402	24397
24403		−fn edit_prediction_fallback_text(edits: &[(Range<Anchor>, String)], cx: &App) -> HighlightedText {
	24398	+fn edit_prediction_fallback_text(edits: &[(Range<Anchor>, Arc<str>)], cx: &App) -> HighlightedText {
24404	24399	// Fallback for providers that don't provide edit_preview (like Copilot/Supermaven)
24405	24400	// Just show the raw edit text with basic styling
24406	24401	let mut text = String::new();
@@ -24793,7 +24788,7 @@ impl Focusable for BreakpointPromptEditor { · 386 unmodified lines
24793	24788	}
24794	24789
24795	24790	fn all_edits_insertions_or_deletions(
24796		− edits: &Vec<(Range<Anchor>, String)>,
	24791	+ edits: &Vec<(Range<Anchor>, Arc<str>)>,
24797	24792	snapshot: &MultiBufferSnapshot,
24798	24793	) -> bool {
24799	24794	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	22915	let text_anchor_edits = edits
22916	22916	.clone()
22917	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	22919	.collect::<Vec<_>>();
22920	22920
22921	22921	let edit_preview = window

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

@@ -720,7 +720,7 @@ impl EditPreview { · 719 unmodified lines
720	720	pub fn highlight_edits(
721	721	&self,
722	722	current_snapshot: &BufferSnapshot,
723		− edits: &[(Range<Anchor>, String)],
	723	+ edits: &[(Range<Anchor>, impl AsRef<str>)],
724	724	include_deletions: bool,
725	725	cx: &App,
726	726	) -> HighlightedText {
@@ -747,7 +747,8 @@ impl EditPreview { · 20 unmodified lines
747	747	.end
748	748	.bias_right(&self.old_snapshot)
749	749	.to_offset(&self.applied_edits_snapshot);
750		− let edit_start_in_preview_snapshot = edit_new_end_in_preview_snapshot - edit_text.len();
	750	+ let edit_start_in_preview_snapshot =
	751	+ edit_new_end_in_preview_snapshot - edit_text.as_ref().len();
751	752
752	753	let unchanged_range_in_preview_snapshot =
753	754	offset_in_preview_snapshot..edit_start_in_preview_snapshot;
@@ -772,7 +773,7 @@ impl EditPreview { · 18 unmodified lines
772	773	);
773	774	}
774	775
775		− if !edit_text.is_empty() {
	776	+ if !edit_text.as_ref().is_empty() {
776	777	highlighted_text.add_text_from_buffer_range(
777	778	edit_start_in_preview_snapshot..edit_new_end_in_preview_snapshot,
778	779	&self.applied_edits_snapshot,
@@ -796,7 +797,7 @@ impl EditPreview { · 17 unmodified lines
796	797	highlighted_text.build()
797	798	}
798	799
799		− fn compute_visible_range(&self, edits: &[(Range<Anchor>, String)]) -> Option<Range<usize>> {
	800	+ fn compute_visible_range<T>(&self, edits: &[(Range<Anchor>, T)]) -> Option<Range<usize>> {
800	801	let (first, _) = edits.first()?;
801	802	let (last, _) = edits.last()?;
802	803
@@ -1130,7 +1131,7 @@ impl Buffer { · 327 unmodified lines
1130	1131
1131	1132	pub fn preview_edits(
1132	1133	&self,
1133		− edits: Arc<[(Range<Anchor>, String)]>,
	1134	+ edits: Arc<[(Range<Anchor>, Arc<str>)]>,
1134	1135	cx: &App,
1135	1136	) -> Task<EditPreview> {
1136	1137	let registry = self.language_registry();

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

@@ -3120,15 +3120,13 @@ async fn test_preview_edits(cx: &mut TestAppContext) { · 3119 unmodified lines
3120	3120	.map(\|(range, text)\| {
3121	3121	(
3122	3122	buffer.anchor_before(range.start)..buffer.anchor_after(range.end),
3123		− text.to_string(),
	3123	+ text.into(),
3124	3124	)
3125	3125	})
3126		− .collect::<Vec<_>>()
	3126	+ .collect::<Arc<[_]>>()
3127	3127	});
3128	3128	let edit_preview = buffer
3129		− .read_with(cx, \|buffer, cx\| {
3130		− buffer.preview_edits(edits.clone().into(), cx)
3131		− })
	3129	+ .read_with(cx, \|buffer, cx\| buffer.preview_edits(edits.clone(), cx))
3132	3130	.await;
3133	3131	let highlighted_edits = cx.read(\|cx\| {
3134	3132	edit_preview.highlight_edits(&buffer.read(cx).snapshot(), &edits, include_deletions, cx)

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

@@ -293,7 +293,7 @@ pub struct FunctionDefinition { · 292 unmodified lines
293	293	pub parameters: Option<Value>,
294	294	}
295	295
296		−#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
	296	+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
297	297	#[serde(tag = "role", rename_all = "lowercase")]
298	298	pub enum RequestMessage {
299	299	Assistant {
@@ -366,25 +366,42 @@ pub struct ImageUrl { · 66 unmodified lines
366	366	pub detail: Option<String>,
367	367	}
368	368
369		−#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
	369	+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
370	370	pub struct ToolCall {
371	371	pub id: String,
372	372	#[serde(flatten)]
373	373	pub content: ToolCallContent,
374	374	}
375	375
376		−#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
	376	+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
377	377	#[serde(tag = "type", rename_all = "lowercase")]
378	378	pub enum ToolCallContent {
379	379	Function { function: FunctionContent },
380	380	}
381	381
382		−#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
	382	+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
383	383	pub struct FunctionContent {
384	384	pub name: String,
385	385	pub arguments: String,
386	386	}
387	387
	388	+#[derive(Clone, Serialize, Deserialize, Debug)]
	389	+pub struct Response {
	390	+ pub id: String,
	391	+ pub object: String,
	392	+ pub created: u64,
	393	+ pub model: String,
	394	+ pub choices: Vec<Choice>,
	395	+ pub usage: Usage,
	396	+}
	397	+
	398	+#[derive(Clone, Serialize, Deserialize, Debug)]
	399	+pub struct Choice {
	400	+ pub index: u32,
	401	+ pub message: RequestMessage,
	402	+ pub finish_reason: Option<String>,
	403	+}
	404	+
388	405	#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
389	406	pub struct ResponseMessageDelta {
390	407	pub role: Option<Role>,
@@ -410,7 +427,7 @@ pub struct FunctionChunk { · 19 unmodified lines
410	427	pub arguments: Option<String>,
411	428	}
412	429
413		−#[derive(Serialize, Deserialize, Debug)]
	430	+#[derive(Clone, Serialize, Deserialize, Debug)]
414	431	pub struct Usage {
415	432	pub prompt_tokens: u64,
416	433	pub completion_tokens: u64,

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

@@ -7,6 +7,7 @@ use language::{Anchor, Buffer, BufferSnapshot}; · 6 unmodified lines
7	7	use std::{
8	8	ops::{AddAssign, Range},
9	9	path::Path,
	10	+ sync::Arc,
10	11	time::Duration,
11	12	};
12	13	use text::{ToOffset, ToPoint};
@@ -51,7 +52,7 @@ fn completion_from_diff( · 38 unmodified lines
51	52	) -> EditPrediction {
52	53	let buffer_text = snapshot.text_for_range(delete_range).collect::<String>();
53	54
54		− let mut edits: Vec<(Range<language::Anchor>, String)> = Vec::new();
	55	+ let mut edits: Vec<(Range<language::Anchor>, Arc<str>)> = Vec::new();
55	56
56	57	let completion_graphemes: Vec<&str> = completion_text.graphemes(true).collect();
57	58	let buffer_graphemes: Vec<&str> = buffer_text.graphemes(true).collect();
@@ -70,7 +71,10 @@ fn completion_from_diff( · 12 unmodified lines
70	71	if k != 0 {
71	72	let offset = snapshot.anchor_after(offset);
72	73	// the range from the current position to item is an inlay.
73		− let edit = (offset..offset, completion_graphemes[i..i + k].join(""));
	74	+ let edit = (
	75	+ offset..offset,
	76	+ completion_graphemes[i..i + k].join("").into(),
	77	+ );
74	78	edits.push(edit);
75	79	}
76	80	i += k + 1;
@@ -90,7 +94,7 @@ fn completion_from_diff( · 13 unmodified lines
90	94	// there is leftover completion text, so drop it as an inlay.
91	95	let edit_range = offset..offset;
92	96	let edit_text = completion_graphemes[i..].join("");
93		− edits.push((edit_range, edit_text));
	97	+ edits.push((edit_range, edit_text.into()));
94	98	}
95	99
96	100	EditPrediction::Local {

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

@@ -133,7 +133,7 @@ pub struct EditPrediction { · 132 unmodified lines
133	133	path: Arc<Path>,
134	134	excerpt_range: Range<usize>,
135	135	cursor_offset: usize,
136		− edits: Arc<[(Range<Anchor>, String)]>,
	136	+ edits: Arc<[(Range<Anchor>, Arc<str>)]>,
137	137	snapshot: BufferSnapshot,
138	138	edit_preview: EditPreview,
139	139	input_outline: Arc<str>,
@@ -150,7 +150,7 @@ impl EditPrediction { · 10 unmodified lines
150	150	.duration_since(self.buffer_snapshotted_at)
151	151	}
152	152
153		− fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, String)>> {
	153	+ fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
154	154	edit_prediction::interpolate_edits(&self.snapshot, new_snapshot, &self.edits)
155	155	}
156	156	}
@@ -711,7 +711,7 @@ impl Zeta { · 554 unmodified lines
711	711	cx.spawn(async move \|cx\| {
712	712	let output_excerpt: Arc<str> = output_excerpt.into();
713	713
714		− let edits: Arc<[(Range<Anchor>, String)]> = cx
	714	+ let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx
715	715	.background_spawn({
716	716	let output_excerpt = output_excerpt.clone();
717	717	let editable_range = editable_range.clone();
@@ -725,7 +725,7 @@ impl Zeta { · 7 unmodified lines
725	725	let edits = edits.clone();
726	726	move \|buffer, cx\| {
727	727	let new_snapshot = buffer.snapshot();
728		− let edits: Arc<[(Range<Anchor>, String)]> =
	728	+ let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
729	729	edit_prediction::interpolate_edits(&snapshot, &new_snapshot, &edits)?
730	730	.into();
731	731	Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
@@ -759,7 +759,7 @@ impl Zeta { · 27 unmodified lines
759	759	output_excerpt: Arc<str>,
760	760	editable_range: Range<usize>,
761	761	snapshot: &BufferSnapshot,
762		− ) -> Result<Vec<(Range<Anchor>, String)>> {
	762	+ ) -> Result<Vec<(Range<Anchor>, Arc<str>)>> {
763	763	let content = output_excerpt.replace(CURSOR_MARKER, "");
764	764
765	765	let start_markers = content
@@ -817,7 +817,7 @@ impl Zeta { · 51 unmodified lines
817	817	new_text: &str,
818	818	offset: usize,
819	819	snapshot: &BufferSnapshot,
820		− ) -> Vec<(Range<Anchor>, String)> {
	820	+ ) -> Vec<(Range<Anchor>, Arc<str>)> {
821	821	text_diff(&old_text, new_text)
822	822	.into_iter()
823	823	.map(\|(mut old_range, new_text)\| {
@@ -836,7 +836,7 @@ impl Zeta { · 12 unmodified lines
836	836	);
837	837	old_range.end = old_range.end.saturating_sub(suffix_len);
838	838
839		− let new_text = new_text[prefix_len..new_text.len() - suffix_len].to_string();
	839	+ let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
840	840	let range = if old_range.is_empty() {
841	841	let anchor = snapshot.anchor_after(old_range.start);
842	842	anchor..anchor
@@ -1183,7 +1183,7 @@ impl CurrentEditPrediction { · 340 unmodified lines
1183	1183	if old_edits.len() == 1 && new_edits.len() == 1 {
1184	1184	let (old_range, old_text) = &old_edits[0];
1185	1185	let (new_range, new_text) = &new_edits[0];
1186		− new_range == old_range && new_text.starts_with(old_text)
	1186	+ new_range == old_range && new_text.starts_with(old_text.as_ref())
1187	1187	} else {
1188	1188	true
1189	1189	}
@@ -1599,13 +1599,8 @@ mod tests { · 409 unmodified lines
1599	1599	#[gpui::test]
1600	1600	async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
1601	1601	let buffer = cx.new(\|cx\| Buffer::local("Lorem ipsum dolor", cx));
1602		− let edits: Arc<[(Range<Anchor>, String)]> = cx.update(\|cx\| {
1603		− to_completion_edits(
1604		− [(2..5, "REM".to_string()), (9..11, "".to_string())],
1605		− &buffer,
1606		− cx,
1607		− )
1608		− .into()
	1602	+ let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(\|cx\| {
	1603	+ to_completion_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
1609	1604	});
1610	1605
1611	1606	let edit_preview = cx
@@ -1635,7 +1630,7 @@ mod tests { · 23 unmodified lines
1635	1630	&buffer,
1636	1631	cx
1637	1632	),
1638		− vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
	1633	+ vec![(2..5, "REM".into()), (9..11, "".into())]
1639	1634	);
1640	1635
1641	1636	buffer.update(cx, \|buffer, cx\| buffer.edit([(2..5, "")], None, cx));
@@ -1645,7 +1640,7 @@ mod tests { · 3 unmodified lines
1645	1640	&buffer,
1646	1641	cx
1647	1642	),
1648		− vec![(2..2, "REM".to_string()), (6..8, "".to_string())]
	1643	+ vec![(2..2, "REM".into()), (6..8, "".into())]
1649	1644	);
1650	1645
1651	1646	buffer.update(cx, \|buffer, cx\| buffer.undo(cx));
@@ -1655,7 +1650,7 @@ mod tests { · 3 unmodified lines
1655	1650	&buffer,
1656	1651	cx
1657	1652	),
1658		− vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
	1653	+ vec![(2..5, "REM".into()), (9..11, "".into())]
1659	1654	);
1660	1655
1661	1656	buffer.update(cx, \|buffer, cx\| buffer.edit([(2..5, "R")], None, cx));
@@ -1665,7 +1660,7 @@ mod tests { · 3 unmodified lines
1665	1660	&buffer,
1666	1661	cx
1667	1662	),
1668		− vec![(3..3, "EM".to_string()), (7..9, "".to_string())]
	1663	+ vec![(3..3, "EM".into()), (7..9, "".into())]
1669	1664	);
1670	1665
1671	1666	buffer.update(cx, \|buffer, cx\| buffer.edit([(3..3, "E")], None, cx));
@@ -1675,7 +1670,7 @@ mod tests { · 3 unmodified lines
1675	1670	&buffer,
1676	1671	cx
1677	1672	),
1678		− vec![(4..4, "M".to_string()), (8..10, "".to_string())]
	1673	+ vec![(4..4, "M".into()), (8..10, "".into())]
1679	1674	);
1680	1675
1681	1676	buffer.update(cx, \|buffer, cx\| buffer.edit([(4..4, "M")], None, cx));
@@ -1685,7 +1680,7 @@ mod tests { · 3 unmodified lines
1685	1680	&buffer,
1686	1681	cx
1687	1682	),
1688		− vec![(9..11, "".to_string())]
	1683	+ vec![(9..11, "".into())]
1689	1684	);
1690	1685
1691	1686	buffer.update(cx, \|buffer, cx\| buffer.edit([(4..5, "")], None, cx));
@@ -1695,7 +1690,7 @@ mod tests { · 3 unmodified lines
1695	1690	&buffer,
1696	1691	cx
1697	1692	),
1698		− vec![(4..4, "M".to_string()), (8..10, "".to_string())]
	1693	+ vec![(4..4, "M".into()), (8..10, "".into())]
1699	1694	);
1700	1695
1701	1696	buffer.update(cx, \|buffer, cx\| buffer.edit([(8..10, "")], None, cx));
@@ -1705,7 +1700,7 @@ mod tests { · 3 unmodified lines
1705	1700	&buffer,
1706	1701	cx
1707	1702	),
1708		− vec![(4..4, "M".to_string())]
	1703	+ vec![(4..4, "M".into())]
1709	1704	);
1710	1705
1711	1706	buffer.update(cx, \|buffer, cx\| buffer.edit([(4..6, "")], None, cx));
@@ -2211,10 +2206,10 @@ mod tests { · 499 unmodified lines
2211	2206	}
2212	2207
2213	2208	fn to_completion_edits(
2214		− iterator: impl IntoIterator<Item = (Range<usize>, String)>,
	2209	+ iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
2215	2210	buffer: &Entity<Buffer>,
2216	2211	cx: &App,
2217		− ) -> Vec<(Range<Anchor>, String)> {
	2212	+ ) -> Vec<(Range<Anchor>, Arc<str>)> {
2218	2213	let buffer = buffer.read(cx);
2219	2214	iterator
2220	2215	.into_iter()
@@ -2228,10 +2223,10 @@ mod tests { · 7 unmodified lines
2228	2223	}
2229	2224
2230	2225	fn from_completion_edits(
2231		− editor_edits: &[(Range<Anchor>, String)],
	2226	+ editor_edits: &[(Range<Anchor>, Arc<str>)],
2232	2227	buffer: &Entity<Buffer>,
2233	2228	cx: &App,
2234		− ) -> Vec<(Range<usize>, String)> {
	2229	+ ) -> Vec<(Range<usize>, Arc<str>)> {
2235	2230	let buffer = buffer.read(cx);
2236	2231	editor_edits
2237	2232	.iter()

crates/zeta2/Cargo.toml+2−1

@@ -28,9 +28,9 @@ indoc.workspace = true · 27 unmodified lines
28	28	language.workspace = true
29	29	language_model.workspace = true
30	30	log.workspace = true
	31	+open_ai.workspace = true
31	32	project.workspace = true
32	33	release_channel.workspace = true
33		−schemars.workspace = true
34	34	serde.workspace = true
35	35	serde_json.workspace = true
36	36	thiserror.workspace = true
@@ -50,3 +50,4 @@ language_model = { workspace = true, features = ["test-support"] } · 13 unmodified lines
50	50	pretty_assertions.workspace = true
51	51	project = { workspace = true, features = ["test-support"] }
52	52	settings = { workspace = true, features = ["test-support"] }
	53	+zlog.workspace = true

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

@@ -1,17 +1,11 @@
1		−use std::{borrow::Cow, ops::Range, path::Path, sync::Arc};
2		−
3		−use anyhow::Context as _;
4		−use cloud_llm_client::predict_edits_v3;
5		−use gpui::{App, AsyncApp, Entity};
6		−use language::{
7		− Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, TextBufferSnapshot, text_diff,
8		−};
9		−use project::Project;
10		−use util::ResultExt;
	1	+use std::{ops::Range, sync::Arc};
	2	+
	3	+use gpui::{AsyncApp, Entity};
	4	+use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, TextBufferSnapshot};
11	5	use uuid::Uuid;
12	6
13	7	#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
14		−pub struct EditPredictionId(Uuid);
	8	+pub struct EditPredictionId(pub Uuid);
15	9
16	10	impl Into<Uuid> for EditPredictionId {
17	11	fn into(self) -> Uuid {
@@ -34,8 +28,7 @@ impl std::fmt::Display for EditPredictionId { · 16 unmodified lines
34	28	#[derive(Clone)]
35	29	pub struct EditPrediction {
36	30	pub id: EditPredictionId,
37		− pub path: Arc<Path>,
38		− pub edits: Arc<[(Range<Anchor>, String)]>,
	31	+ pub edits: Arc<[(Range<Anchor>, Arc<str>)]>,
39	32	pub snapshot: BufferSnapshot,
40	33	pub edit_preview: EditPreview,
41	34	// We keep a reference to the buffer so that we do not need to reload it from disk when applying the prediction.
@@ -43,90 +36,43 @@ pub struct EditPrediction { · 1 unmodified line
43	36	}
44	37
45	38	impl EditPrediction {
46		− pub async fn from_response(
47		− response: predict_edits_v3::PredictEditsResponse,
48		− active_buffer_old_snapshot: &TextBufferSnapshot,
49		− active_buffer: &Entity<Buffer>,
50		− project: &Entity<Project>,
	39	+ pub async fn new(
	40	+ id: EditPredictionId,
	41	+ edited_buffer: &Entity<Buffer>,
	42	+ edited_buffer_snapshot: &BufferSnapshot,
	43	+ edits: Vec<(Range<Anchor>, Arc<str>)>,
51	44	cx: &mut AsyncApp,
52	45	) -> Option<Self> {
53		− // TODO only allow cloud to return one path
54		− let Some(path) = response.edits.first().map(\|e\| e.path.clone()) else {
55		− return None;
56		− };
	46	+ let (edits, snapshot, edit_preview_task) = edited_buffer
	47	+ .read_with(cx, \|buffer, cx\| {
	48	+ let new_snapshot = buffer.snapshot();
	49	+ let edits: Arc<[_]> =
	50	+ interpolate_edits(&edited_buffer_snapshot, &new_snapshot, edits.into())?.into();
57	51
58		− let is_same_path = active_buffer
59		− .read_with(cx, \|buffer, cx\| buffer_path_eq(buffer, &path, cx))
60		− .ok()?;
61		−
62		− let (buffer, edits, snapshot, edit_preview_task) = if is_same_path {
63		− active_buffer
64		− .read_with(cx, \|buffer, cx\| {
65		− let new_snapshot = buffer.snapshot();
66		− let edits = edits_from_response(&response.edits, &active_buffer_old_snapshot);
67		− let edits: Arc<[_]> =
68		− interpolate_edits(active_buffer_old_snapshot, &new_snapshot, edits)?.into();
69		−
70		− Some((
71		− active_buffer.clone(),
72		− edits.clone(),
73		− new_snapshot,
74		− buffer.preview_edits(edits, cx),
75		− ))
76		− })
77		− .ok()??
78		− } else {
79		− let buffer_handle = project
80		− .update(cx, \|project, cx\| {
81		− let project_path = project
82		− .find_project_path(&path, cx)
83		− .context("Failed to find project path for zeta edit")?;
84		− anyhow::Ok(project.open_buffer(project_path, cx))
85		− })
86		− .ok()?
87		− .log_err()?
88		− .await
89		− .context("Failed to open buffer for zeta edit")
90		− .log_err()?;
91		−
92		− buffer_handle
93		− .read_with(cx, \|buffer, cx\| {
94		− let snapshot = buffer.snapshot();
95		− let edits = edits_from_response(&response.edits, &snapshot);
96		− if edits.is_empty() {
97		− return None;
98		− }
99		− Some((
100		− buffer_handle.clone(),
101		− edits.clone(),
102		− snapshot,
103		− buffer.preview_edits(edits, cx),
104		− ))
105		− })
106		− .ok()??
107		− };
	52	+ Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
	53	+ })
	54	+ .ok()??;
108	55
109	56	let edit_preview = edit_preview_task.await;
110	57
111	58	Some(EditPrediction {
112		− id: EditPredictionId(response.request_id),
113		− path,
	59	+ id,
114	60	edits,
115	61	snapshot,
116	62	edit_preview,
117		− buffer,
	63	+ buffer: edited_buffer.clone(),
118	64	})
119	65	}
120	66
121	67	pub fn interpolate(
122	68	&self,
123	69	new_snapshot: &TextBufferSnapshot,
124		− ) -> Option<Vec<(Range<Anchor>, String)>> {
	70	+ ) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
125	71	interpolate_edits(&self.snapshot, new_snapshot, self.edits.clone())
126	72	}
127	73
128		− pub fn targets_buffer(&self, buffer: &Buffer, cx: &App) -> bool {
129		− buffer_path_eq(buffer, &self.path, cx)
	74	+ pub fn targets_buffer(&self, buffer: &Buffer) -> bool {
	75	+ self.snapshot.remote_id() == buffer.remote_id()
130	76	}
131	77	}
132	78
@@ -134,21 +80,16 @@ impl std::fmt::Debug for EditPrediction { · 1 unmodified line
134	80	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
135	81	f.debug_struct("EditPrediction")
136	82	.field("id", &self.id)
137		− .field("path", &self.path)
138	83	.field("edits", &self.edits)
139	84	.finish()
140	85	}
141	86	}
142	87
143		−pub fn buffer_path_eq(buffer: &Buffer, path: &Path, cx: &App) -> bool {
144		− buffer.file().map(\|p\| p.full_path(cx)).as_deref() == Some(path)
145		−}
146		−
147	88	pub fn interpolate_edits(
148	89	old_snapshot: &TextBufferSnapshot,
149	90	new_snapshot: &TextBufferSnapshot,
150		− current_edits: Arc<[(Range<Anchor>, String)]>,
151		−) -> Option<Vec<(Range<Anchor>, String)>> {
	91	+ current_edits: Arc<[(Range<Anchor>, Arc<str>)]>,
	92	+) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
152	93	let mut edits = Vec::new();
153	94
154	95	let mut model_edits = current_edits.iter().peekable();
@@ -173,7 +114,7 @@ pub fn interpolate_edits( · 18 unmodified lines
173	114	if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
174	115	if !model_suffix.is_empty() {
175	116	let anchor = old_snapshot.anchor_after(user_edit.old.end);
176		− edits.push((anchor..anchor, model_suffix.to_string()));
	117	+ edits.push((anchor..anchor, model_suffix.into()));
177	118	}
178	119
179	120	model_edits.next();
@@ -190,135 +131,17 @@ pub fn interpolate_edits( · 10 unmodified lines
190	131	if edits.is_empty() { None } else { Some(edits) }
191	132	}
192	133
193		−pub fn line_range_to_point_range(range: Range<predict_edits_v3::Line>) -> Range<language::Point> {
194		− language::Point::new(range.start.0, 0)..language::Point::new(range.end.0, 0)
195		−}
196		−
197		−fn edits_from_response(
198		− edits: &[predict_edits_v3::Edit],
199		− snapshot: &TextBufferSnapshot,
200		−) -> Arc<[(Range<Anchor>, String)]> {
201		− edits
202		− .iter()
203		− .flat_map(\|edit\| {
204		− let point_range = line_range_to_point_range(edit.range.clone());
205		− let offset = point_range.to_offset(snapshot).start;
206		− let old_text = snapshot.text_for_range(point_range);
207		−
208		− excerpt_edits_from_response(
209		− old_text.collect::<Cow<str>>(),
210		− &edit.content,
211		− offset,
212		− &snapshot,
213		− )
214		− })
215		− .collect::<Vec<_>>()
216		− .into()
217		−}
218		−
219		−fn excerpt_edits_from_response(
220		− old_text: Cow<str>,
221		− new_text: &str,
222		− offset: usize,
223		− snapshot: &TextBufferSnapshot,
224		−) -> impl Iterator<Item = (Range<Anchor>, String)> {
225		− text_diff(&old_text, new_text)
226		− .into_iter()
227		− .map(move \|(mut old_range, new_text)\| {
228		− old_range.start += offset;
229		− old_range.end += offset;
230		−
231		− let prefix_len = common_prefix(
232		− snapshot.chars_for_range(old_range.clone()),
233		− new_text.chars(),
234		− );
235		− old_range.start += prefix_len;
236		−
237		− let suffix_len = common_prefix(
238		− snapshot.reversed_chars_for_range(old_range.clone()),
239		− new_text[prefix_len..].chars().rev(),
240		− );
241		− old_range.end = old_range.end.saturating_sub(suffix_len);
242		−
243		− let new_text = new_text[prefix_len..new_text.len() - suffix_len].to_string();
244		− let range = if old_range.is_empty() {
245		− let anchor = snapshot.anchor_after(old_range.start);
246		− anchor..anchor
247		− } else {
248		− snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
249		− };
250		− (range, new_text)
251		− })
252		−}
253		−
254		−fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
255		− a.zip(b)
256		− .take_while(\|(a, b)\| a == b)
257		− .map(\|(a, _)\| a.len_utf8())
258		− .sum()
259		−}
260		−
261	134	#[cfg(test)]
262	135	mod tests {
263		− use std::path::PathBuf;
264		−
265	136	use super::*;
266		− use cloud_llm_client::predict_edits_v3;
267		− use edit_prediction_context::Line;
268	137	use gpui::{App, Entity, TestAppContext, prelude::*};
269		− use indoc::indoc;
270	138	use language::{Buffer, ToOffset as _};
271	139
272		− #[gpui::test]
273		− async fn test_compute_edits(cx: &mut TestAppContext) {
274		− let old = indoc! {r#"
275		− fn main() {
276		− let args =
277		− println!("{}", args[1])
278		− }
279		− "#};
280		−
281		− let new = indoc! {r#"
282		− fn main() {
283		− let args = std::env::args();
284		− println!("{}", args[1]);
285		− }
286		− "#};
287		−
288		− let buffer = cx.new(\|cx\| Buffer::local(old, cx));
289		− let snapshot = buffer.read_with(cx, \|buffer, _cx\| buffer.snapshot());
290		−
291		− // TODO cover more cases when multi-file is supported
292		− let big_edits = vec![predict_edits_v3::Edit {
293		− path: PathBuf::from("test.txt").into(),
294		− range: Line(0)..Line(old.lines().count() as u32),
295		− content: new.into(),
296		− }];
297		−
298		− let edits = edits_from_response(&big_edits, &snapshot);
299		− assert_eq!(edits.len(), 2);
300		− assert_eq!(
301		− edits[0].0.to_point(&snapshot).start,
302		− language::Point::new(1, 14)
303		− );
304		− assert_eq!(edits[0].1, " std::env::args();");
305		− assert_eq!(
306		− edits[1].0.to_point(&snapshot).start,
307		− language::Point::new(2, 27)
308		− );
309		− assert_eq!(edits[1].1, ";");
310		− }
311		−
312	140	#[gpui::test]
313	141	async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
314	142	let buffer = cx.new(\|cx\| Buffer::local("Lorem ipsum dolor", cx));
315		− let edits: Arc<[(Range<Anchor>, String)]> = cx.update(\|cx\| {
316		− to_prediction_edits(
317		− [(2..5, "REM".to_string()), (9..11, "".to_string())],
318		− &buffer,
319		− cx,
320		− )
321		− .into()
	143	+ let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(\|cx\| {
	144	+ to_prediction_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
322	145	});
323	146
324	147	let edit_preview = cx
@@ -329,7 +152,6 @@ mod tests { · 4 unmodified lines
329	152	id: EditPredictionId(Uuid::new_v4()),
330	153	edits,
331	154	snapshot: cx.read(\|cx\| buffer.read(cx).snapshot()),
332		− path: Path::new("test.txt").into(),
333	155	buffer: buffer.clone(),
334	156	edit_preview,
335	157	};
@@ -341,7 +163,7 @@ mod tests { · 5 unmodified lines
341	163	&buffer,
342	164	cx
343	165	),
344		− vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
	166	+ vec![(2..5, "REM".into()), (9..11, "".into())]
345	167	);
346	168
347	169	buffer.update(cx, \|buffer, cx\| buffer.edit([(2..5, "")], None, cx));
@@ -351,7 +173,7 @@ mod tests { · 3 unmodified lines
351	173	&buffer,
352	174	cx
353	175	),
354		− vec![(2..2, "REM".to_string()), (6..8, "".to_string())]
	176	+ vec![(2..2, "REM".into()), (6..8, "".into())]
355	177	);
356	178
357	179	buffer.update(cx, \|buffer, cx\| buffer.undo(cx));
@@ -361,7 +183,7 @@ mod tests { · 3 unmodified lines
361	183	&buffer,
362	184	cx
363	185	),
364		− vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
	186	+ vec![(2..5, "REM".into()), (9..11, "".into())]
365	187	);
366	188
367	189	buffer.update(cx, \|buffer, cx\| buffer.edit([(2..5, "R")], None, cx));
@@ -371,7 +193,7 @@ mod tests { · 3 unmodified lines
371	193	&buffer,
372	194	cx
373	195	),
374		− vec![(3..3, "EM".to_string()), (7..9, "".to_string())]
	196	+ vec![(3..3, "EM".into()), (7..9, "".into())]
375	197	);
376	198
377	199	buffer.update(cx, \|buffer, cx\| buffer.edit([(3..3, "E")], None, cx));
@@ -381,7 +203,7 @@ mod tests { · 3 unmodified lines
381	203	&buffer,
382	204	cx
383	205	),
384		− vec![(4..4, "M".to_string()), (8..10, "".to_string())]
	206	+ vec![(4..4, "M".into()), (8..10, "".into())]
385	207	);
386	208
387	209	buffer.update(cx, \|buffer, cx\| buffer.edit([(4..4, "M")], None, cx));
@@ -391,7 +213,7 @@ mod tests { · 3 unmodified lines
391	213	&buffer,
392	214	cx
393	215	),
394		− vec![(9..11, "".to_string())]
	216	+ vec![(9..11, "".into())]
395	217	);
396	218
397	219	buffer.update(cx, \|buffer, cx\| buffer.edit([(4..5, "")], None, cx));
@@ -401,7 +223,7 @@ mod tests { · 3 unmodified lines
401	223	&buffer,
402	224	cx
403	225	),
404		− vec![(4..4, "M".to_string()), (8..10, "".to_string())]
	226	+ vec![(4..4, "M".into()), (8..10, "".into())]
405	227	);
406	228
407	229	buffer.update(cx, \|buffer, cx\| buffer.edit([(8..10, "")], None, cx));
@@ -411,7 +233,7 @@ mod tests { · 3 unmodified lines
411	233	&buffer,
412	234	cx
413	235	),
414		− vec![(4..4, "M".to_string())]
	236	+ vec![(4..4, "M".into())]
415	237	);
416	238
417	239	buffer.update(cx, \|buffer, cx\| buffer.edit([(4..6, "")], None, cx));
@@ -420,10 +242,10 @@ mod tests { · 2 unmodified lines
420	242	}
421	243
422	244	fn to_prediction_edits(
423		− iterator: impl IntoIterator<Item = (Range<usize>, String)>,
	245	+ iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
424	246	buffer: &Entity<Buffer>,
425	247	cx: &App,
426		− ) -> Vec<(Range<Anchor>, String)> {
	248	+ ) -> Vec<(Range<Anchor>, Arc<str>)> {
427	249	let buffer = buffer.read(cx);
428	250	iterator
429	251	.into_iter()
@@ -437,10 +259,10 @@ mod tests { · 7 unmodified lines
437	259	}
438	260
439	261	fn from_prediction_edits(
440		− editor_edits: &[(Range<Anchor>, String)],
	262	+ editor_edits: &[(Range<Anchor>, Arc<str>)],
441	263	buffer: &Entity<Buffer>,
442	264	cx: &App,
443		− ) -> Vec<(Range<usize>, String)> {
	265	+ ) -> Vec<(Range<usize>, Arc<str>)> {
444	266	let buffer = buffer.read(cx);
445	267	editor_edits
446	268	.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	6	AcceptEditPredictionBody, EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
7	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	11	use collections::HashMap;
11	12	use edit_prediction_context::{
12	13	DeclarationId, DeclarationStyle, EditPredictionContext, EditPredictionContextOptions,
@@ -24,11 +25,13 @@ use gpui::{ · 11 unmodified lines
24	25	use language::{Anchor, Buffer, DiagnosticSet, LanguageServerId, ToOffset as _, ToPoint};
25	26	use language::{BufferSnapshot, OffsetRangeExt};
26	27	use language_model::{LlmApiToken, RefreshLlmTokenListener};
	28	+use open_ai::FunctionDefinition;
27	29	use project::Project;
28	30	use release_channel::AppVersion;
29	31	use serde::de::DeserializeOwned;
30	32	use std::collections::{VecDeque, hash_map};
31		−use std::fmt::Write;
	33	+use uuid::Uuid;
	34	+
32	35	use std::ops::Range;
33	36	use std::path::Path;
34	37	use std::str::FromStr as _;
@@ -42,12 +45,12 @@ use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_noti · 7 unmodified lines
42	45	pub mod merge_excerpts;
43	46	mod prediction;
44	47	mod provider;
45		−pub mod related_excerpts;
	48	+pub mod retrieval_search;
	49	+pub mod udiff;
46	50
47	51	use crate::merge_excerpts::merge_excerpts;
48	52	use crate::prediction::EditPrediction;
49		−use crate::related_excerpts::find_related_excerpts;
50		−pub use crate::related_excerpts::{LlmContextOptions, SearchToolQuery};
	53	+pub use crate::prediction::EditPredictionId;
51	54	pub use provider::ZetaEditPredictionProvider;
52	55
53	56	/// Maximum number of events to track.
@@ -59,9 +62,10 @@ pub const DEFAULT_EXCERPT_OPTIONS: EditPredictionExcerptOptions = EditPrediction · 5 unmodified lines
59	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	69	excerpt: DEFAULT_EXCERPT_OPTIONS,
66	70	};
67	71
@@ -122,14 +126,19 @@ pub struct ZetaOptions { · 54 unmodified lines
122	126
123	127	#[derive(Debug, Clone, PartialEq)]
124	128	pub enum ContextMode {
125		− Llm(LlmContextOptions),
	129	+ Agentic(AgenticContextOptions),
126	130	Syntax(EditPredictionContextOptions),
127	131	}
128	132
	133	+#[derive(Debug, Clone, PartialEq)]
	134	+pub struct AgenticContextOptions {
	135	+ pub excerpt: EditPredictionExcerptOptions,
	136	+}
	137	+
129	138	impl ContextMode {
130	139	pub fn excerpt(&self) -> &EditPredictionExcerptOptions {
131	140	match self {
132		− ContextMode::Llm(options) => &options.excerpt,
	141	+ ContextMode::Agentic(options) => &options.excerpt,
133	142	ContextMode::Syntax(options) => &options.excerpt,
134	143	}
135	144	}
@@ -140,9 +149,8 @@ pub enum ZetaDebugInfo { · 4 unmodified lines
140	149	ContextRetrievalStarted(ZetaContextRetrievalStartedDebugInfo),
141	150	SearchQueriesGenerated(ZetaSearchQueryDebugInfo),
142	151	SearchQueriesExecuted(ZetaContextRetrievalDebugInfo),
143		− SearchResultsFiltered(ZetaContextRetrievalDebugInfo),
144	152	ContextRetrievalFinished(ZetaContextRetrievalDebugInfo),
145		− EditPredicted(ZetaEditPredictionDebugInfo),
	153	+ EditPredictionRequested(ZetaEditPredictionDebugInfo),
146	154	}
147	155
148	156	#[derive(Debug)]
@@ -165,14 +173,14 @@ pub struct ZetaEditPredictionDebugInfo { · 16 unmodified lines
165	173	pub buffer: WeakEntity<Buffer>,
166	174	pub position: language::Anchor,
167	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	179	#[derive(Debug)]
172	180	pub struct ZetaSearchQueryDebugInfo {
173	181	pub project: Entity<Project>,
174	182	pub timestamp: Instant,
175		− pub queries: Vec<SearchToolQuery>,
	183	+ pub regex_by_glob: HashMap<String, String>,
176	184	}
177	185
178	186	pub type RequestDebugInfo = predict_edits_v3::DebugInfo;
@@ -224,7 +232,7 @@ impl CurrentEditPrediction { · 45 unmodified lines
224	232	{
225	233	let (old_range, old_text) = &old_edits[0];
226	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	236	} else {
229	237	true
230	238	}
@@ -539,7 +547,7 @@ impl Zeta { · 308 unmodified lines
539	547	prediction,
540	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	551	Some(BufferEditPrediction::Local { prediction })
544	552	} else if *requested_by_buffer_id == buffer.entity_id() {
545	553	Some(BufferEditPrediction::Jump { prediction })
@@ -639,7 +647,7 @@ impl Zeta { · 93 unmodified lines
639	647	pub fn request_prediction(
640	648	&mut self,
641	649	project: &Entity<Project>,
642		− buffer: &Entity<Buffer>,
	650	+ active_buffer: &Entity<Buffer>,
643	651	position: language::Anchor,
644	652	cx: &mut Context<Self>,
645	653	) -> Task<Result<Option<EditPrediction>>> {
@@ -651,8 +659,8 @@ impl Zeta { · 5 unmodified lines
651	659	.read_with(cx, \|index, _cx\| index.state().clone())
652	660	});
653	661	let options = self.options.clone();
654		− let snapshot = buffer.read(cx).snapshot();
655		− let Some(excerpt_path) = snapshot
	662	+ let active_snapshot = active_buffer.read(cx).snapshot();
	663	+ let Some(excerpt_path) = active_snapshot
656	664	.file()
657	665	.map(\|path\| -> Arc<Path> { path.full_path(cx).into() })
658	666	else {
@@ -678,12 +686,13 @@ impl Zeta { · 19 unmodified lines
678	686	})
679	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\| {
684		− let mut path = f.worktree.read(cx).absolutize(&f.path);
685		− if path.pop() { Some(path) } else { None }
686		− });
	691	+ let parent_abs_path =
	692	+ project::File::from_dyn(active_buffer.read(cx).file()).and_then(\|f\| {
	693	+ let mut path = f.worktree.read(cx).absolutize(&f.path);
	694	+ if path.pop() { Some(path) } else { None }
	695	+ });
687	696
688	697	// TODO data collection
689	698	let can_collect_data = cx.is_staff();
@@ -692,9 +701,10 @@ impl Zeta { · 2 unmodified lines
692	701	.and_then(\|project_state\| project_state.context.as_ref())
693	702	.unwrap_or(&HashMap::default())
694	703	.iter()
695		− .filter_map(\|(buffer, ranges)\| {
696		− let buffer = buffer.read(cx);
	704	+ .filter_map(\|(buffer_entity, ranges)\| {
	705	+ let buffer = buffer_entity.read(cx);
697	706	Some((
	707	+ buffer_entity.clone(),
698	708	buffer.snapshot(),
699	709	buffer.file()?.full_path(cx).into(),
700	710	ranges.clone(),
@@ -703,8 +713,7 @@ impl Zeta { · 2 unmodified lines
703	713	.collect::<Vec<_>>();
704	714
705	715	let request_task = cx.background_spawn({
706		− let snapshot = snapshot.clone();
707		− let buffer = buffer.clone();
	716	+ let active_buffer = active_buffer.clone();
708	717	async move {
709	718	let index_state = if let Some(index_state) = index_state {
710	719	Some(index_state.lock_owned().await)
@@ -712,8 +721,8 @@ impl Zeta { · 1 unmodified line
712	721	None
713	722	};
714	723
715		− let cursor_offset = position.to_offset(&snapshot);
716		− let cursor_point = cursor_offset.to_point(&snapshot);
	724	+ let cursor_offset = position.to_offset(&active_snapshot);
	725	+ let cursor_point = cursor_offset.to_point(&active_snapshot);
717	726
718	727	let before_retrieval = chrono::Utc::now();
719	728
@@ -721,29 +730,30 @@ impl Zeta { · 1 unmodified line
721	730	Self::gather_nearby_diagnostics(
722	731	cursor_offset,
723	732	&diagnostics,
724		− &snapshot,
	733	+ &active_snapshot,
725	734	options.max_diagnostic_bytes,
726	735	);
727	736
728		− let request = match options.context {
729		− ContextMode::Llm(context_options) => {
	737	+ let cloud_request = match options.context {
	738	+ ContextMode::Agentic(context_options) => {
730	739	let Some(excerpt) = EditPredictionExcerpt::select_from_buffer(
731	740	cursor_point,
732		− &snapshot,
	741	+ &active_snapshot,
733	742	&context_options.excerpt,
734	743	index_state.as_deref(),
735	744	) else {
736	745	return Ok((None, None));
737	746	};
738	747
739		− let excerpt_anchor_range = snapshot.anchor_after(excerpt.range.start)
740		− ..snapshot.anchor_before(excerpt.range.end);
	748	+ let excerpt_anchor_range = active_snapshot.anchor_after(excerpt.range.start)
	749	+ ..active_snapshot.anchor_before(excerpt.range.end);
741	750
742		− if let Some(buffer_ix) = included_files
743		− .iter()
744		− .position(\|(buffer, _, _)\| buffer.remote_id() == snapshot.remote_id())
	751	+ if let Some(buffer_ix) =
	752	+ included_files.iter().position(\|(_, snapshot, _, _)\| {
	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	757	let range_ix = ranges
748	758	.binary_search_by(\|probe\| {
749	759	probe
@@ -758,15 +768,16 @@ impl Zeta { · 8 unmodified lines
758	768	included_files.swap(buffer_ix, last_ix);
759	769	} else {
760	770	included_files.push((
761		− snapshot,
	771	+ active_buffer.clone(),
	772	+ active_snapshot,
762	773	excerpt_path.clone(),
763	774	vec![excerpt_anchor_range],
764	775	));
765	776	}
766	777
767	778	let included_files = included_files
768		− .into_iter()
769		− .map(\|(buffer, path, ranges)\| {
	779	+ .iter()
	780	+ .map(\|(_, buffer, path, ranges)\| {
770	781	let excerpts = merge_excerpts(
771	782	&buffer,
772	783	ranges.iter().map(\|range\| {
@@ -775,7 +786,7 @@ impl Zeta { · 2 unmodified lines
775	786	}),
776	787	);
777	788	predict_edits_v3::IncludedFile {
778		− path,
	789	+ path: path.clone(),
779	790	max_row: Line(buffer.max_point().row),
780	791	excerpts,
781	792	}
@@ -809,7 +820,7 @@ impl Zeta { · 27 unmodified lines
809	820	ContextMode::Syntax(context_options) => {
810	821	let Some(context) = EditPredictionContext::gather_context(
811	822	cursor_point,
812		− &snapshot,
	823	+ &active_snapshot,
813	824	parent_abs_path.as_deref(),
814	825	&context_options,
815	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	850	let retrieval_time = chrono::Utc::now() - before_retrieval;
838	851
839	852	let debug_response_tx = if let Some(debug_tx) = &debug_tx {
840	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	855	debug_tx
847		− .unbounded_send(ZetaDebugInfo::EditPredicted(ZetaEditPredictionDebugInfo {
848		− request: request.clone(),
849		− retrieval_time,
850		− buffer: buffer.downgrade(),
851		− local_prompt,
852		− position,
853		− response_rx,
854		− }))
	856	+ .unbounded_send(ZetaDebugInfo::EditPredictionRequested(
	857	+ ZetaEditPredictionDebugInfo {
	858	+ request: cloud_request.clone(),
	859	+ retrieval_time,
	860	+ buffer: active_buffer.downgrade(),
	861	+ local_prompt: match prompt_result.as_ref() {
	862	+ Ok((prompt, _)) => Ok(prompt.clone()),
	863	+ Err(err) => Err(err.to_string()),
	864	+ },
	865	+ position,
	866	+ response_rx,
	867	+ },
	868	+ ))
855	869	.ok();
856	870	Some(response_tx)
857	871	} else {
@@ -861,61 +875,144 @@ impl Zeta { · 3 unmodified lines
861	875	if cfg!(debug_assertions) && std::env::var("ZED_ZETA2_SKIP_REQUEST").is_ok() {
862	876	if let Some(debug_response_tx) = debug_response_tx {
863	877	debug_response_tx
864		− .send(Err("Request skipped".to_string()))
	878	+ .send((Err("Request skipped".to_string()), TimeDelta::zero()))
865	879	.ok();
866	880	}
867	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	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	910	if let Some(debug_response_tx) = debug_response_tx {
874	911	debug_response_tx
875		− .send(
	912	+ .send((
876	913	response
877	914	.as_ref()
878	915	.map_err(\|err\| err.to_string())
879	916	.map(\|response\| response.0.clone()),
880		− )
	917	+ request_time,
	918	+ ))
881	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	987	cx.spawn({
891		− let project = project.clone();
892	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	991	else {
895	992	return Ok(None);
896	993	};
897	994
898	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	1005	client: Arc<Client>,
906	1006	llm_token: LlmApiToken,
907	1007	app_version: SemanticVersion,
908		− request: predict_edits_v3::PredictEditsRequest,
909		− ) -> Result<(
910		− predict_edits_v3::PredictEditsResponse,
911		− Option<EditPredictionUsage>,
912		− )> {
	1008	+ request: open_ai::Request,
	1009	+ ) -> Result<(open_ai::Response, Option<EditPredictionUsage>)> {
913	1010	let url = if let Ok(predict_edits_url) = std::env::var("ZED_PREDICT_EDITS_URL") {
914	1011	http_client::Url::parse(&predict_edits_url)?
915	1012	} else {
916	1013	client
917	1014	.http_client()
918		− .build_zed_llm_url("/predict_edits/v3", &[])?
	1015	+ .build_zed_llm_url("/predict_edits/raw", &[])?
919	1016	};
920	1017
921	1018	Self::send_api_request(
@@ -1052,7 +1149,7 @@ impl Zeta { · 130 unmodified lines
1052	1149	cursor_position: language::Anchor,
1053	1150	cx: &mut Context<Self>,
1054	1151	) {
1055		− if !matches!(&self.options().context, ContextMode::Llm { .. }) {
	1152	+ if !matches!(&self.options().context, ContextMode::Agentic { .. }) {
1056	1153	return;
1057	1154	}
1058	1155
@@ -1100,36 +1197,149 @@ impl Zeta { · 41 unmodified lines
1100	1197	cursor_position: language::Anchor,
1101	1198	cx: &mut Context<Self>,
1102	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	1272	cx.spawn(async move \|this, cx\| {
1104		− let related_excerpts_result = this
1105		− .update(cx, \|this, cx\| {
1106		− let Some(zeta_project) = this.projects.get(&project.entity_id()) else {
1107		− return Task::ready(anyhow::Ok(HashMap::default()));
1108		− };
	1273	+ log::trace!("Sending search planning request");
	1274	+ let response =
	1275	+ Self::send_raw_llm_request(client, llm_token, app_version, request).await;
	1276	+ let mut response = Self::handle_api_response(&this, response, cx)?;
	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 {
1111		− return Task::ready(anyhow::Ok(HashMap::default()));
1112		− };
	1292	+ let mut regex_by_glob: HashMap<String, String> = HashMap::default();
	1293	+ for tool_call in tool_calls {
	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() {
1117		− if let Some(event) = event.to_request_event(cx) {
1118		− writeln!(&mut edit_history_unified_diff, "{event}").ok();
1119		− }
	1304	+ let input: SearchToolInput = serde_json::from_str(&function.arguments)?;
	1305	+ for query in input.queries {
	1306	+ let regex = regex_by_glob.entry(query.glob).or_default();
	1307	+ if !regex.is_empty() {
	1308	+ regex.push('\|');
1120	1309	}
	1310	+ regex.push_str(&query.regex);
	1311	+ }
	1312	+ }
1121	1313
1122		− find_related_excerpts(
1123		− buffer.clone(),
1124		− cursor_position,
1125		− &project,
1126		− edit_history_unified_diff,
1127		− options,
1128		− this.debug_tx.clone(),
1129		− cx,
1130		− )
1131		− })?
1132		− .await;
	1314	+ if let Some(debug_tx) = &debug_tx {
	1315	+ debug_tx
	1316	+ .unbounded_send(ZetaDebugInfo::SearchQueriesGenerated(
	1317	+ ZetaSearchQueryDebugInfo {
	1318	+ project: project.clone(),
	1319	+ timestamp: Instant::now(),
	1320	+ regex_by_glob: regex_by_glob.clone(),
	1321	+ },
	1322	+ ))
	1323	+ .ok();
	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	1344	this.update(cx, \|this, _cx\| {
1135	1345	let Some(zeta_project) = this.projects.get_mut(&project.entity_id()) else {
@@ -1249,7 +1459,7 @@ impl Zeta { · 113 unmodified lines
1249	1459	&snapshot,
1250	1460	parent_abs_path.as_deref(),
1251	1461	match &options.context {
1252		− ContextMode::Llm(_) => {
	1462	+ ContextMode::Agentic(_) => {
1253	1463	// TODO
1254	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	1637	#[cfg(test)]
1428	1638	mod tests {
1429		− use std::{
1430		− path::{Path, PathBuf},
1431		− sync::Arc,
1432		− };
	1639	+ use std::{path::Path, sync::Arc};
1433	1640
1434	1641	use client::UserStore;
1435	1642	use clock::FakeSystemClock;
1436		− use cloud_llm_client::predict_edits_v3::{self, Point};
1437		− use edit_prediction_context::Line;
	1643	+ use cloud_zeta2_prompt::retrieval_prompt::{SearchToolInput, SearchToolQuery};
1438	1644	use futures::{
1439	1645	AsyncReadExt, StreamExt,
1440	1646	channel::{mpsc, oneshot},
@@ -1445,7 +1651,8 @@ mod tests { · 4 unmodified lines
1445	1651	prelude::*,
1446	1652	};
1447	1653	use indoc::indoc;
1448		− use language::{LanguageServerId, OffsetRangeExt as _};
	1654	+ use language::OffsetRangeExt as _;
	1655	+ use open_ai::Usage;
1449	1656	use pretty_assertions::{assert_eq, assert_matches};
1450	1657	use project::{FakeFs, Project};
1451	1658	use serde_json::json;
@@ -1462,8 +1669,8 @@ mod tests { · 10 unmodified lines
1462	1669	fs.insert_tree(
1463	1670	"/root",
1464	1671	json!({
1465		− "1.txt": "Hello!\nHow\nBye",
1466		− "2.txt": "Hola!\nComo\nAdios"
	1672	+ "1.txt": "Hello!\nHow\nBye\n",
	1673	+ "2.txt": "Hola!\nComo\nAdios\n"
1467	1674	}),
1468	1675	)
1469	1676	.await;
@@ -1489,16 +1696,17 @@ mod tests { · 19 unmodified lines
1489	1696	zeta.refresh_prediction(&project, &buffer1, position, cx)
1490	1697	});
1491	1698	let (_request, respond_tx) = req_rx.next().await.unwrap();
	1699	+
1492	1700	respond_tx
1493		− .send(predict_edits_v3::PredictEditsResponse {
1494		− request_id: Uuid::new_v4(),
1495		− edits: vec![predict_edits_v3::Edit {
1496		− path: Path::new(path!("root/1.txt")).into(),
1497		− range: Line(0)..Line(snapshot1.max_point().row + 1),
1498		− content: "Hello!\nHow are you?\nBye".into(),
1499		− }],
1500		− debug_info: None,
1501		− })
	1701	+ .send(model_response(indoc! {r"
	1702	+ --- a/root/1.txt
	1703	+ +++ b/root/1.txt
	1704	+ @@ ... @@
	1705	+ Hello!
	1706	+ -How
	1707	+ +How are you?
	1708	+ Bye
	1709	+ "}))
1502	1710	.unwrap();
1503	1711	prediction_task.await.unwrap();
1504	1712
@@ -1509,21 +1717,67 @@ mod tests { · 4 unmodified lines
1509	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	1767	// Prediction for another file
1513	1768	let prediction_task = zeta.update(cx, \|zeta, cx\| {
1514	1769	zeta.refresh_prediction(&project, &buffer1, position, cx)
1515	1770	});
1516	1771	let (_request, respond_tx) = req_rx.next().await.unwrap();
1517	1772	respond_tx
1518		− .send(predict_edits_v3::PredictEditsResponse {
1519		− request_id: Uuid::new_v4(),
1520		− edits: vec![predict_edits_v3::Edit {
1521		− path: Path::new(path!("root/2.txt")).into(),
1522		− range: Line(0)..Line(snapshot1.max_point().row + 1),
1523		− content: "Hola!\nComo estas?\nAdios".into(),
1524		− }],
1525		− debug_info: None,
1526		− })
	1773	+ .send(model_response(indoc! {r#"
	1774	+ --- a/root/2.txt
	1775	+ +++ b/root/2.txt
	1776	+ Hola!
	1777	+ -Como
	1778	+ +Como estas?
	1779	+ Adios
	1780	+ "#}))
1527	1781	.unwrap();
1528	1782	prediction_task.await.unwrap();
1529	1783	zeta.read_with(cx, \|zeta, cx\| {
@@ -1532,7 +1786,7 @@ mod tests { · 2 unmodified lines
1532	1786	.unwrap();
1533	1787	assert_matches!(
1534	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	1813	fs.insert_tree(
1560	1814	"/root",
1561	1815	json!({
1562		− "foo.md": "Hello!\nHow\nBye"
	1816	+ "foo.md": "Hello!\nHow\nBye\n"
1563	1817	}),
1564	1818	)
1565	1819	.await;
@@ -1579,29 +1833,31 @@ mod tests { · 13 unmodified lines
1579	1833	zeta.request_prediction(&project, &buffer, position, cx)
1580	1834	});
1581	1835
1582		− let (request, respond_tx) = req_rx.next().await.unwrap();
1583		− assert_eq!(
1584		− request.excerpt_path.as_ref(),
1585		− Path::new(path!("root/foo.md"))
1586		− );
1587		− assert_eq!(
1588		− request.cursor_point,
1589		− Point {
1590		− line: Line(1),
1591		− column: 3
1592		− }
1593		− );
	1836	+ let (_, respond_tx) = req_rx.next().await.unwrap();
	1837	+
	1838	+ // TODO Put back when we have a structured request again
	1839	+ // assert_eq!(
	1840	+ // request.excerpt_path.as_ref(),
	1841	+ // Path::new(path!("root/foo.md"))
	1842	+ // );
	1843	+ // assert_eq!(
	1844	+ // request.cursor_point,
	1845	+ // Point {
	1846	+ // line: Line(1),
	1847	+ // column: 3
	1848	+ // }
	1849	+ // );
1594	1850
1595	1851	respond_tx
1596		− .send(predict_edits_v3::PredictEditsResponse {
1597		− request_id: Uuid::new_v4(),
1598		− edits: vec![predict_edits_v3::Edit {
1599		− path: Path::new(path!("root/foo.md")).into(),
1600		− range: Line(0)..Line(snapshot.max_point().row + 1),
1601		− content: "Hello!\nHow are you?\nBye".into(),
1602		− }],
1603		− debug_info: None,
1604		− })
	1852	+ .send(model_response(indoc! { r"
	1853	+ --- a/root/foo.md
	1854	+ +++ b/root/foo.md
	1855	+ @@ ... @@
	1856	+ Hello!
	1857	+ -How
	1858	+ +How are you?
	1859	+ Bye
	1860	+ "}))
1605	1861	.unwrap();
1606	1862
1607	1863	let prediction = prediction_task.await.unwrap().unwrap();
@@ -1611,7 +1867,7 @@ mod tests { · 3 unmodified lines
1611	1867	prediction.edits[0].0.to_point(&snapshot).start,
1612	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	1873	#[gpui::test]
@@ -1621,7 +1877,7 @@ mod tests { · 3 unmodified lines
1621	1877	fs.insert_tree(
1622	1878	"/root",
1623	1879	json!({
1624		− "foo.md": "Hello!\n\nBye"
	1880	+ "foo.md": "Hello!\n\nBye\n"
1625	1881	}),
1626	1882	)
1627	1883	.await;
@@ -1652,34 +1908,30 @@ mod tests { · 24 unmodified lines
1652	1908
1653	1909	let (request, respond_tx) = req_rx.next().await.unwrap();
1654	1910
1655		− assert_eq!(request.events.len(), 1);
1656		− assert_eq!(
1657		− request.events[0],
1658		− predict_edits_v3::Event::BufferChange {
1659		− path: Some(PathBuf::from(path!("root/foo.md"))),
1660		− old_path: None,
1661		− diff: indoc! {"
1662		− @@ -1,3 +1,3 @@
1663		− Hello!
1664		− -
1665		− +How
1666		− Bye
1667		− "}
1668		− .to_string(),
1669		− predicted: false
1670		− }
	1911	+ let prompt = prompt_from_request(&request);
	1912	+ assert!(
	1913	+ prompt.contains(indoc! {"
	1914	+ --- a/root/foo.md
	1915	+ +++ b/root/foo.md
	1916	+ @@ -1,3 +1,3 @@
	1917	+ Hello!
	1918	+ -
	1919	+ +How
	1920	+ Bye
	1921	+ "}),
	1922	+ "{prompt}"
1671	1923	);
1672	1924
1673	1925	respond_tx
1674		− .send(predict_edits_v3::PredictEditsResponse {
1675		− request_id: Uuid::new_v4(),
1676		− edits: vec![predict_edits_v3::Edit {
1677		− path: Path::new(path!("root/foo.md")).into(),
1678		− range: Line(0)..Line(snapshot.max_point().row + 1),
1679		− content: "Hello!\nHow are you?\nBye".into(),
1680		− }],
1681		− debug_info: None,
1682		− })
	1926	+ .send(model_response(indoc! {r#"
	1927	+ --- a/root/foo.md
	1928	+ +++ b/root/foo.md
	1929	+ @@ ... @@
	1930	+ Hello!
	1931	+ -How
	1932	+ +How are you?
	1933	+ Bye
	1934	+ "#}))
1683	1935	.unwrap();
1684	1936
1685	1937	let prediction = prediction_task.await.unwrap().unwrap();
@@ -1689,114 +1941,150 @@ mod tests { · 3 unmodified lines
1689	1941	prediction.edits[0].0.to_point(&snapshot).start,
1690	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]
1696		− async fn test_request_diagnostics(cx: &mut TestAppContext) {
1697		− let (zeta, mut req_rx) = init_test(cx);
1698		− let fs = FakeFs::new(cx.executor());
1699		− fs.insert_tree(
1700		− "/root",
1701		− json!({
1702		− "foo.md": "Hello!\nBye"
1703		− }),
1704		− )
1705		− .await;
1706		− let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
	1947	+ // Skipped until we start including diagnostics in prompt
	1948	+ // #[gpui::test]
	1949	+ // async fn test_request_diagnostics(cx: &mut TestAppContext) {
	1950	+ // let (zeta, mut req_rx) = init_test(cx);
	1951	+ // let fs = FakeFs::new(cx.executor());
	1952	+ // fs.insert_tree(
	1953	+ // "/root",
	1954	+ // json!({
	1955	+ // "foo.md": "Hello!\nBye"
	1956	+ // }),
	1957	+ // )
	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();
1709		− let diagnostic = lsp::Diagnostic {
1710		− range: lsp::Range::new(lsp::Position::new(1, 1), lsp::Position::new(1, 5)),
1711		− severity: Some(lsp::DiagnosticSeverity::ERROR),
1712		− message: "\"Hello\" deprecated. Use \"Hi\" instead".to_string(),
1713		− ..Default::default()
	2061	+ fn prompt_from_request(request: &open_ai::Request) -> &str {
	2062	+ assert_eq!(request.messages.len(), 1);
	2063	+ let open_ai::RequestMessage::User {
	2064	+ content: open_ai::MessageContent::Plain(content),
	2065	+ ..
	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		−
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		− );
	2073	+ content
1784	2074	}
1785	2075
1786	2076	fn init_test(
1787	2077	cx: &mut TestAppContext,
1788	2078	) -> (
1789	2079	Entity<Zeta>,
1790		− mpsc::UnboundedReceiver<(
1791		− predict_edits_v3::PredictEditsRequest,
1792		− oneshot::Sender<predict_edits_v3::PredictEditsResponse>,
1793		− )>,
	2080	+ mpsc::UnboundedReceiver<(open_ai::Request, oneshot::Sender<open_ai::Response>)>,
1794	2081	) {
1795	2082	cx.update(move \|cx\| {
1796	2083	let settings_store = SettingsStore::test(cx);
1797	2084	cx.set_global(settings_store);
1798	2085	language::init(cx);
1799	2086	Project::init_settings(cx);
	2087	+ zlog::init_test();
1800	2088
1801	2089	let (req_tx, req_rx) = mpsc::unbounded();
1802	2090
@@ -1811,7 +2099,7 @@ mod tests { · 8 unmodified lines
1811	2099	"token": "test"
1812	2100	}))
1813	2101	.unwrap(),
1814		− "/predict_edits/v3" => {
	2102	+ "/predict_edits/raw" => {
1815	2103	let mut buf = Vec::new();
1816	2104	body.read_to_end(&mut buf).await.ok();
1817	2105	let req = serde_json::from_slice(&buf).unwrap();

crates/zeta2_tools/Cargo.toml+2−1

@@ -27,10 +27,11 @@ log.workspace = true · 26 unmodified lines
27	27	multi_buffer.workspace = true
28	28	ordered-float.workspace = true
29	29	project.workspace = true
	30	+regex-syntax = "0.8.8"
30	31	serde.workspace = true
	32	+serde_json.workspace = true
31	33	telemetry.workspace = true
32	34	text.workspace = true
33		−regex-syntax = "0.8.8"
34	35	ui.workspace = true
35	36	ui_input.workspace = true
36	37	util.workspace = true

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

@@ -45,7 +45,6 @@ struct RetrievalRun { · 44 unmodified lines
45	45	started_at: Instant,
46	46	search_results_generated_at: Option<Instant>,
47	47	search_results_executed_at: Option<Instant>,
48		− search_results_filtered_at: Option<Instant>,
49	48	finished_at: Option<Instant>,
50	49	}
51	50
@@ -117,17 +116,12 @@ impl Zeta2ContextView { · 65 unmodified lines
117	116	self.handle_search_queries_executed(info, window, cx);
118	117	}
119	118	}
120		− ZetaDebugInfo::SearchResultsFiltered(info) => {
121		− if info.project == self.project {
122		− self.handle_search_results_filtered(info, window, cx);
123		− }
124		− }
125	119	ZetaDebugInfo::ContextRetrievalFinished(info) => {
126	120	if info.project == self.project {
127	121	self.handle_context_retrieval_finished(info, window, cx);
128	122	}
129	123	}
130		− ZetaDebugInfo::EditPredicted(_) => {}
	124	+ ZetaDebugInfo::EditPredictionRequested(_) => {}
131	125	}
132	126	}
133	127
@@ -159,7 +153,6 @@ impl Zeta2ContextView { · 25 unmodified lines
159	153	started_at: info.timestamp,
160	154	search_results_generated_at: None,
161	155	search_results_executed_at: None,
162		− search_results_filtered_at: None,
163	156	finished_at: None,
164	157	});
165	158
@@ -218,18 +211,18 @@ impl Zeta2ContextView { · 52 unmodified lines
218	211
219	212	run.search_results_generated_at = Some(info.timestamp);
220	213	run.search_queries = info
221		− .queries
	214	+ .regex_by_glob
222	215	.into_iter()
223		− .map(\|query\| {
	216	+ .map(\|(glob, regex)\| {
224	217	let mut regex_parser = regex_syntax::ast::parse::Parser::new();
225	218
226	219	GlobQueries {
227		− glob: query.glob,
228		− alternations: match regex_parser.parse(&query.regex) {
	220	+ glob,
	221	+ alternations: match regex_parser.parse(&regex) {
229	222	Ok(regex_syntax::ast::Ast::Alternation(ref alt)) => {
230	223	alt.asts.iter().map(\|ast\| ast.to_string()).collect()
231	224	}
232		− _ => vec![query.regex],
	225	+ _ => vec![regex],
233	226	},
234	227	}
235	228	})
@@ -256,20 +249,6 @@ impl Zeta2ContextView { · 20 unmodified lines
256	249	cx.notify();
257	250	}
258	251
259		− fn handle_search_results_filtered(
260		− &mut self,
261		− info: ZetaContextRetrievalDebugInfo,
262		− _window: &mut Window,
263		− cx: &mut Context<Self>,
264		− ) {
265		− let Some(run) = self.runs.back_mut() else {
266		− return;
267		− };
268		−
269		− run.search_results_filtered_at = Some(info.timestamp);
270		− cx.notify();
271		− }
272		−
273	252	fn handle_go_back(
274	253	&mut self,
275	254	_: &Zeta2ContextGoBack,
@@ -398,19 +377,10 @@ impl Zeta2ContextView { · 122 unmodified lines
398	377	};
399	378	div = div.child(format!("Ran search: {:>5} ms", (t2 - t1).as_millis()));
400	379
401		− let Some(t3) = run.search_results_filtered_at else {
402		− return pending_message(div, "Filtering results...");
403		− };
404		− div =
405		− div.child(format!("Filtered results: {:>5} ms", (t3 - t2).as_millis()));
406		−
407		− let Some(t4) = run.finished_at else {
408		− return pending_message(div, "Building excerpts");
409		− };
410		− div = div
411		− .child(format!("Build excerpts: {:>5} µs", (t4 - t3).as_micros()))
412		− .child(format!("Total: {:>5} ms", (t4 - t0).as_millis()));
413		− div
	380	+ div.child(format!(
	381	+ "Total: {:>5} ms",
	382	+ (run.finished_at.unwrap_or(t0) - t0).as_millis()
	383	+ ))
414	384	}),
415	385	)
416	386	}

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

@@ -5,7 +5,7 @@ use std::{cmp::Reverse, path::PathBuf, str::FromStr, sync::Arc, time::Duration}; · 4 unmodified lines
5	5	use chrono::TimeDelta;
6	6	use client::{Client, UserStore};
7	7	use cloud_llm_client::predict_edits_v3::{
8		− self, DeclarationScoreComponents, PredictEditsRequest, PredictEditsResponse, PromptFormat,
	8	+ DeclarationScoreComponents, PredictEditsRequest, PromptFormat,
9	9	};
10	10	use collections::HashMap;
11	11	use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
@@ -23,7 +23,7 @@ use ui_input::InputField; · 11 unmodified lines
23	23	use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
24	24	use workspace::{Item, SplitDirection, Workspace};
25	25	use zeta2::{
26		− ContextMode, DEFAULT_SYNTAX_CONTEXT_OPTIONS, LlmContextOptions, Zeta, Zeta2FeatureFlag,
	26	+ AgenticContextOptions, ContextMode, DEFAULT_SYNTAX_CONTEXT_OPTIONS, Zeta, Zeta2FeatureFlag,
27	27	ZetaDebugInfo, ZetaEditPredictionDebugInfo, ZetaOptions,
28	28	};
29	29
@@ -123,6 +123,7 @@ struct LastPrediction { · 93 unmodified lines
123	123	context_editor: Entity<Editor>,
124	124	prompt_editor: Entity<Editor>,
125	125	retrieval_time: TimeDelta,
	126	+ request_time: Option<TimeDelta>,
126	127	buffer: WeakEntity<Buffer>,
127	128	position: language::Anchor,
128	129	state: LastPredictionState,
@@ -143,7 +144,7 @@ enum LastPredictionState { · 14 unmodified lines
143	144	model_response_editor: Entity<Editor>,
144	145	feedback_editor: Entity<Editor>,
145	146	feedback: Option<Feedback>,
146		− response: predict_edits_v3::PredictEditsResponse,
	147	+ request_id: String,
147	148	},
148	149	Failed {
149	150	message: String,
@@ -217,7 +218,7 @@ impl Zeta2Inspector { · 67 unmodified lines
217	218	});
218	219
219	220	match &options.context {
220		− ContextMode::Llm(_) => {
	221	+ ContextMode::Agentic(_) => {
221	222	self.context_mode = ContextModeState::Llm;
222	223	}
223	224	ContextMode::Syntax(_) => {
@@ -307,9 +308,11 @@ impl Zeta2Inspector { · 83 unmodified lines
307	308	};
308	309
309	310	let context = match zeta_options.context {
310		− ContextMode::Llm(_context_options) => ContextMode::Llm(LlmContextOptions {
311		− excerpt: excerpt_options,
312		− }),
	311	+ ContextMode::Agentic(_context_options) => {
	312	+ ContextMode::Agentic(AgenticContextOptions {
	313	+ excerpt: excerpt_options,
	314	+ })
	315	+ }
313	316	ContextMode::Syntax(context_options) => {
314	317	let max_retrieved_declarations = match &this.context_mode {
315	318	ContextModeState::Llm => {
@@ -368,7 +371,7 @@ impl Zeta2Inspector { · 52 unmodified lines
368	371	let language_registry = self.project.read(cx).languages().clone();
369	372	async move \|this, cx\| {
370	373	let mut languages = HashMap::default();
371		− let ZetaDebugInfo::EditPredicted(prediction) = prediction else {
	374	+ let ZetaDebugInfo::EditPredictionRequested(prediction) = prediction else {
372	375	return;
373	376	};
374	377	for ext in prediction
@@ -396,6 +399,8 @@ impl Zeta2Inspector { · 21 unmodified lines
396	399	.await
397	400	.log_err();
398	401
	402	+ let json_language = language_registry.language_for_name("Json").await.log_err();
	403	+
399	404	this.update_in(cx, \|this, window, cx\| {
400	405	let context_editor = cx.new(\|cx\| {
401	406	let mut excerpt_score_components = HashMap::default();
@@ -492,25 +497,15 @@ impl Zeta2Inspector { · 90 unmodified lines
492	497
493	498	let task = cx.spawn_in(window, {
494	499	let markdown_language = markdown_language.clone();
	500	+ let json_language = json_language.clone();
495	501	async move \|this, cx\| {
496	502	let response = response_rx.await;
497	503
498	504	this.update_in(cx, \|this, window, cx\| {
499	505	if let Some(prediction) = this.last_prediction.as_mut() {
500	506	prediction.state = match response {
501		− Ok(Ok(response)) => {
502		− if let Some(debug_info) = &response.debug_info {
503		− prediction.prompt_editor.update(
504		− cx,
505		− \|prompt_editor, cx\| {
506		− prompt_editor.set_text(
507		− debug_info.prompt.as_str(),
508		− window,
509		− cx,
510		− );
511		− },
512		− );
513		− }
	507	+ Ok((Ok(response), request_time)) => {
	508	+ prediction.request_time = Some(request_time);
514	509
515	510	let feedback_editor = cx.new(\|cx\| {
516	511	let buffer = cx.new(\|cx\| {
@@ -577,16 +572,11 @@ impl Zeta2Inspector { · 60 unmodified lines
577	572	model_response_editor: cx.new(\|cx\| {
578	573	let buffer = cx.new(\|cx\| {
579	574	let mut buffer = Buffer::local(
580		− response
581		− .debug_info
582		− .as_ref()
583		− .map(\|p\| p.model_response.as_str())
584		− .unwrap_or(
585		− "(Debug info not available)",
586		− ),
	575	+ serde_json::to_string_pretty(&response)
	576	+ .unwrap_or_default(),
587	577	cx,
588	578	);
589		− buffer.set_language(markdown_language, cx);
	579	+ buffer.set_language(json_language, cx);
590	580	buffer
591	581	});
592	582	let buffer = cx.new(\|cx\| {
@@ -607,10 +597,11 @@ impl Zeta2Inspector { · 14 unmodified lines
607	597	}),
608	598	feedback_editor,
609	599	feedback: None,
610		− response,
	600	+ request_id: response.id.clone(),
611	601	}
612	602	}
613		− Ok(Err(err)) => {
	603	+ Ok((Err(err), request_time)) => {
	604	+ prediction.request_time = Some(request_time);
614	605	LastPredictionState::Failed { message: err }
615	606	}
616	607	Err(oneshot::Canceled) => LastPredictionState::Failed {
@@ -644,6 +635,7 @@ impl Zeta2Inspector { · 27 unmodified lines
644	635	editor
645	636	}),
646	637	retrieval_time,
	638	+ request_time: None,
647	639	buffer,
648	640	position,
649	641	state: LastPredictionState::Requested,
@@ -700,7 +692,7 @@ impl Zeta2Inspector { · 50 unmodified lines
700	692	feedback: feedback_state,
701	693	feedback_editor,
702	694	model_response_editor,
703		− response,
	695	+ request_id,
704	696	..
705	697	} = &mut last_prediction.state
706	698	else {
@@ -734,11 +726,10 @@ impl Zeta2Inspector { · 27 unmodified lines
734	726
735	727	telemetry::event!(
736	728	"Zeta2 Prediction Rated",
737		− id = response.request_id,
	729	+ id = request_id,
738	730	kind = kind,
739	731	text = text,
740	732	request = last_prediction.request,
741		− response = response,
742	733	project_snapshot = project_snapshot,
743	734	);
744	735	})
@@ -834,11 +825,11 @@ impl Zeta2Inspector { · 89 unmodified lines
834	825	let current_options =
835	826	this.zeta.read(cx).options().clone();
836	827	match current_options.context.clone() {
837		− ContextMode::Llm(_) => {}
	828	+ ContextMode::Agentic(_) => {}
838	829	ContextMode::Syntax(context_options) => {
839	830	let options = ZetaOptions {
840		− context: ContextMode::Llm(
841		− LlmContextOptions {
	831	+ context: ContextMode::Agentic(
	832	+ AgenticContextOptions {
842	833	excerpt: context_options.excerpt,
843	834	},
844	835	),
@@ -865,7 +856,7 @@ impl Zeta2Inspector { · 20 unmodified lines
865	856	let current_options =
866	857	this.zeta.read(cx).options().clone();
867	858	match current_options.context.clone() {
868		− ContextMode::Llm(context_options) => {
	859	+ ContextMode::Agentic(context_options) => {
869	860	let options = ZetaOptions {
870	861	context: ContextMode::Syntax(
871	862	EditPredictionContextOptions {
@@ -976,25 +967,6 @@ impl Zeta2Inspector { · 104 unmodified lines
976	967	return None;
977	968	};
978	969
979		− let (prompt_planning_time, inference_time, parsing_time) =
980		− if let LastPredictionState::Success {
981		− response:
982		− PredictEditsResponse {
983		− debug_info: Some(debug_info),
984		− ..
985		− },
986		− ..
987		− } = &prediction.state
988		− {
989		− (
990		− Some(debug_info.prompt_planning_time),
991		− Some(debug_info.inference_time),
992		− Some(debug_info.parsing_time),
993		− )
994		− } else {
995		− (None, None, None)
996		− };
997		−
998	970	Some(
999	971	v_flex()
1000	972	.p_4()
@@ -1005,12 +977,7 @@ impl Zeta2Inspector { · 4 unmodified lines
1005	977	"Context retrieval",
1006	978	Some(prediction.retrieval_time),
1007	979	))
1008		− .child(Self::render_duration(
1009		− "Prompt planning",
1010		− prompt_planning_time,
1011		− ))
1012		− .child(Self::render_duration("Inference", inference_time))
1013		− .child(Self::render_duration("Parsing", parsing_time)),
	980	+ .child(Self::render_duration("Request", prediction.request_time)),
1014	981	)
1015	982	}
1016	983

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

@@ -7,13 +7,14 @@ use std::{ · 6 unmodified lines
7	7
8	8	use anyhow::Result;
9	9	use clap::Args;
10		−use cloud_llm_client::udiff::DiffLine;
11	10	use collections::HashSet;
12	11	use gpui::AsyncApp;
	12	+use zeta2::udiff::DiffLine;
13	13
14	14	use crate::{
15	15	example::{Example, NamedExample},
16	16	headless::ZetaCliAppState,
	17	+ paths::CACHE_DIR,
17	18	predict::{PredictionDetails, zeta2_predict},
18	19	};
19	20
@@ -54,10 +55,8 @@ pub async fn run_evaluate_one( · 34 unmodified lines
54	55	app_state: Arc<ZetaCliAppState>,
55	56	cx: &mut AsyncApp,
56	57	) -> Result<EvaluationResult> {
57		− let cache_dir = Path::new(&std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default())
58		− .join("../../target/zeta-prediction-cache");
59	58	let example = NamedExample::load(&example_path).unwrap();
60		− let example_cache_path = cache_dir.join(&example_path.file_name().unwrap());
	59	+ let example_cache_path = CACHE_DIR.join(&example_path.file_name().unwrap());
61	60
62	61	let predictions = if !re_run && example_cache_path.exists() {
63	62	let file_contents = fs::read_to_string(&example_cache_path)?;
@@ -74,7 +73,7 @@ pub async fn run_evaluate_one( · 10 unmodified lines
74	73	};
75	74
76	75	if !example_cache_path.exists() {
77		− fs::create_dir_all(&cache_dir).unwrap();
	76	+ fs::create_dir_all(&*CACHE_DIR).unwrap();
78	77	fs::write(
79	78	example_cache_path,
80	79	serde_json::to_string(&predictions).unwrap(),

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

@@ -1,28 +1,31 @@
1	1	use std::{
2	2	borrow::Cow,
3	3	cell::RefCell,
4		− env,
5	4	fmt::{self, Display},
6	5	fs,
7	6	io::Write,
8	7	mem,
9		− ops::Range,
10	8	path::{Path, PathBuf},
11	9	sync::Arc,
12	10	};
13	11
14		−use anyhow::{Context as _, Result};
	12	+use anyhow::{Context as _, Result, anyhow};
15	13	use clap::ValueEnum;
16		−use collections::{HashMap, HashSet};
	14	+use cloud_zeta2_prompt::CURSOR_MARKER;
	15	+use collections::HashMap;
17	16	use futures::{
18	17	AsyncWriteExt as _,
19	18	lock::{Mutex, OwnedMutexGuard},
20	19	};
21	20	use gpui::{AsyncApp, Entity, http_client::Url};
22		−use language::Buffer;
	21	+use language::{Anchor, Buffer};
23	22	use project::{Project, ProjectPath};
24	23	use pulldown_cmark::CowStr;
25	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	30	const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
28	31	const EDIT_HISTORY_HEADING: &str = "Edit History";
@@ -215,12 +218,10 @@ impl NamedExample { · 186 unmodified lines
215	218	let (repo_owner, repo_name) = self.repo_name()?;
216	219	let file_name = self.file_name();
217	220
218		− let worktrees_dir = env::current_dir()?.join("target").join("zeta-worktrees");
219		− let repos_dir = env::current_dir()?.join("target").join("zeta-repos");
220		− fs::create_dir_all(&repos_dir)?;
221		− fs::create_dir_all(&worktrees_dir)?;
	221	+ fs::create_dir_all(&*REPOS_DIR)?;
	222	+ 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	225	let repo_lock = lock_repo(&repo_dir).await;
225	226
226	227	if !repo_dir.is_dir() {
@@ -251,7 +252,7 @@ impl NamedExample { · 24 unmodified lines
251	252	};
252	253
253	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	256	if worktree_path.is_dir() {
256	257	run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
257	258	run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
@@ -309,7 +310,6 @@ impl NamedExample { · 51 unmodified lines
309	310	.collect()
310	311	}
311	312
312		− #[allow(unused)]
313	313	fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> {
314	314	// git@github.com:owner/repo.git
315	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	397	#[must_use]
348	398	pub async fn apply_edit_history(
349	399	&self,
350	400	project: &Entity<Project>,
351	401	cx: &mut AsyncApp,
352		− ) -> Result<HashSet<Entity<Buffer>>> {
353		− apply_diff(&self.example.edit_history, project, cx).await
	402	+ ) -> Result<OpenedBuffers<'_>> {
	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	496	.lock_owned()
447	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	1	mod evaluate;
2	2	mod example;
3	3	mod headless;
	4	+mod paths;
4	5	mod predict;
5	6	mod source_location;
6	7	mod syntax_retrieval_stats;
@@ -10,28 +11,22 @@ use crate::evaluate::{EvaluateArguments, run_evaluate}; · 3 unmodified lines
10	11	use crate::example::{ExampleFormat, NamedExample};
11	12	use crate::predict::{PredictArguments, run_zeta2_predict};
12	13	use crate::syntax_retrieval_stats::retrieval_stats;
13		−use ::serde::Serialize;
14	14	use ::util::paths::PathStyle;
15		−use anyhow::{Context as _, Result, anyhow};
	15	+use anyhow::{Result, anyhow};
16	16	use clap::{Args, Parser, Subcommand};
17		−use cloud_llm_client::predict_edits_v3::{self, Excerpt};
18		−use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
	17	+use cloud_llm_client::predict_edits_v3;
19	18	use edit_prediction_context::{
20		− EditPredictionContextOptions, EditPredictionExcerpt, EditPredictionExcerptOptions,
21		− EditPredictionScoreOptions, Line,
	19	+ EditPredictionContextOptions, EditPredictionExcerptOptions, EditPredictionScoreOptions,
22	20	};
23		−use futures::StreamExt as _;
24		−use futures::channel::mpsc;
25	21	use gpui::{Application, AsyncApp, Entity, prelude::*};
26		−use language::{Bias, Buffer, BufferSnapshot, OffsetRangeExt, Point};
27		−use language_model::LanguageModelRegistry;
	22	+use language::{Bias, Buffer, BufferSnapshot, Point};
28	23	use project::{Project, Worktree};
29	24	use reqwest_client::ReqwestClient;
30	25	use serde_json::json;
31	26	use std::io::{self};
32	27	use std::time::Duration;
33	28	use std::{collections::HashSet, path::PathBuf, str::FromStr, sync::Arc};
34		−use zeta2::{ContextMode, LlmContextOptions, SearchToolQuery};
	29	+use zeta2::ContextMode;
35	30
36	31	use crate::headless::ZetaCliAppState;
37	32	use crate::source_location::SourceLocation;
@@ -79,12 +74,6 @@ enum Zeta2Command { · 41 unmodified lines
79	74	#[command(subcommand)]
80	75	command: Zeta2SyntaxCommand,
81	76	},
82		− Llm {
83		− #[clap(flatten)]
84		− args: Zeta2Args,
85		− #[command(subcommand)]
86		− command: Zeta2LlmCommand,
87		− },
88	77	Predict(PredictArguments),
89	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	99	#[derive(Debug, Args)]
119	100	#[group(requires = "worktree")]
120	101	struct ContextArgs {
@@ -388,197 +369,6 @@ async fn zeta2_syntax_context( · 267 unmodified lines
388	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	372	async fn zeta1_context(
583	373	args: ContextArgs,
584	374	app_state: &Arc<ZetaCliAppState>,
@@ -670,13 +460,6 @@ fn main() { · 85 unmodified lines
670	460	};
671	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	464	Command::ConvertExample {
682	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

@@ -1,22 +1,20 @@
1	1	use crate::example::{ActualExcerpt, NamedExample};
2		−
3	2	use crate::headless::ZetaCliAppState;
	3	+use crate::paths::LOGS_DIR;
4	4	use ::serde::Serialize;
5		−use ::util::paths::PathStyle;
6	5	use anyhow::{Context as _, Result, anyhow};
7	6	use clap::Args;
8	7	use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
9	8	use futures::StreamExt as _;
10	9	use gpui::AsyncApp;
11		−use language_model::LanguageModelRegistry;
12		−use project::{Project, ProjectPath};
	10	+use project::Project;
13	11	use serde::Deserialize;
14	12	use std::cell::Cell;
	13	+use std::fs;
15	14	use std::io::Write;
16	15	use std::path::PathBuf;
17	16	use std::sync::Arc;
18	17	use std::time::{Duration, Instant};
19		−use util::rel_path::RelPath;
20	18
21	19	#[derive(Debug, Args)]
22	20	pub struct PredictArguments {
@@ -50,21 +48,12 @@ pub async fn zeta2_predict( · 27 unmodified lines
50	48	app_state: &Arc<ZetaCliAppState>,
51	49	cx: &mut AsyncApp,
52	50	) -> Result<PredictionDetails> {
	51	+ fs::create_dir_all(&*LOGS_DIR)?;
53	52	let worktree_path = example.setup_worktree().await?;
54	53
55	54	if !AUTHENTICATED.get() {
56	55	AUTHENTICATED.set(true);
57	56
58		− cx.update(\|cx\| {
59		− LanguageModelRegistry::global(cx).update(cx, \|registry, cx\| {
60		− registry
61		− .provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
62		− .unwrap()
63		− .authenticate(cx)
64		− })
65		− })?
66		− .await?;
67		−
68	57	app_state
69	58	.client
70	59	.sign_in_with_optional_connect(true, cx)
@@ -83,6 +72,8 @@ pub async fn zeta2_predict( · 12 unmodified lines
83	72	)
84	73	})?;
85	74
	75	+ let buffer_store = project.read_with(cx, \|project, _\| project.buffer_store().clone())?;
	76	+
86	77	let worktree = project
87	78	.update(cx, \|project, cx\| {
88	79	project.create_worktree(&worktree_path, true, cx)
@@ -94,58 +85,30 @@ pub async fn zeta2_predict( · 5 unmodified lines
94	85	})?
95	86	.await;
96	87
97		− let _edited_buffers = example.apply_edit_history(&project, cx).await?;
98		−
99		− let cursor_path = RelPath::new(&example.example.cursor_path, PathStyle::Posix)?.into_arc();
100		−
101		− let cursor_buffer = project
102		− .update(cx, \|project, cx\| {
103		− project.open_buffer(
104		− ProjectPath {
105		− worktree_id: worktree.read(cx).id(),
106		− path: cursor_path,
107		− },
108		− cx,
109		− )
110		− })?
111		− .await?;
112		−
113		− let cursor_offset_within_excerpt = example
114		− .example
115		− .cursor_position
116		− .find(CURSOR_MARKER)
117		− .ok_or_else(\|\| anyhow!("missing cursor marker"))?;
118		− let mut cursor_excerpt = example.example.cursor_position.clone();
119		− cursor_excerpt.replace_range(
120		− cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
121		− "",
122		− );
123		− let excerpt_offset = cursor_buffer.read_with(cx, \|buffer, _cx\| {
124		− let text = buffer.text();
125		−
126		− let mut matches = text.match_indices(&cursor_excerpt);
127		− let Some((excerpt_offset, _)) = matches.next() else {
128		− anyhow::bail!(
129		− "Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
130		− );
131		− };
132		− assert!(matches.next().is_none());
	88	+ let zeta = cx.update(\|cx\| zeta2::Zeta::global(&app_state.client, &app_state.user_store, cx))?;
133	89
134		− Ok(excerpt_offset)
135		− })??;
	90	+ cx.subscribe(&buffer_store, {
	91	+ let project = project.clone();
	92	+ move \|_, event, cx\| match event {
	93	+ project::buffer_store::BufferStoreEvent::BufferAdded(buffer) => {
	94	+ zeta2::Zeta::try_global(cx)
	95	+ .unwrap()
	96	+ .update(cx, \|zeta, cx\| zeta.register_buffer(&buffer, &project, cx));
	97	+ }
	98	+ _ => {}
	99	+ }
	100	+ })?
	101	+ .detach();
136	102
137		− let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
138		− let cursor_anchor =
139		− cursor_buffer.read_with(cx, \|buffer, _\| buffer.anchor_after(cursor_offset))?;
	103	+ let _edited_buffers = example.apply_edit_history(&project, cx).await?;
	104	+ let (cursor_buffer, cursor_anchor) = example.cursor_position(&project, cx).await?;
140	105
141		− let zeta = cx.update(\|cx\| zeta2::Zeta::global(&app_state.client, &app_state.user_store, cx))?;
	106	+ let mut debug_rx = zeta.update(cx, \|zeta, _\| zeta.debug_info())?;
142	107
143	108	let refresh_task = zeta.update(cx, \|zeta, cx\| {
144		− zeta.register_buffer(&cursor_buffer, &project, cx);
145	109	zeta.refresh_context(project.clone(), cursor_buffer.clone(), cursor_anchor, cx)
146	110	})?;
147	111
148		− let mut debug_rx = zeta.update(cx, \|zeta, _\| zeta.debug_info())?;
149	112	let mut context_retrieval_started_at = None;
150	113	let mut context_retrieval_finished_at = None;
151	114	let mut search_queries_generated_at = None;
@@ -159,9 +122,14 @@ pub async fn zeta2_predict( · 7 unmodified lines
159	122	match event {
160	123	zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
161	124	context_retrieval_started_at = Some(info.timestamp);
	125	+ fs::write(LOGS_DIR.join("search_prompt.md"), &info.search_prompt)?;
162	126	}
163	127	zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
164	128	search_queries_generated_at = Some(info.timestamp);
	129	+ fs::write(
	130	+ LOGS_DIR.join("search_queries.json"),
	131	+ serde_json::to_string_pretty(&info.regex_by_glob).unwrap(),
	132	+ )?;
165	133	}
166	134	zeta2::ZetaDebugInfo::SearchQueriesExecuted(info) => {
167	135	search_queries_executed_at = Some(info.timestamp);
@@ -173,11 +141,21 @@ pub async fn zeta2_predict( · 5 unmodified lines
173	141	zeta.request_prediction(&project, &cursor_buffer, cursor_anchor, cx)
174	142	})?);
175	143	}
176		− zeta2::ZetaDebugInfo::EditPredicted(request) => {
	144	+ zeta2::ZetaDebugInfo::EditPredictionRequested(request) => {
177	145	prediction_started_at = Some(Instant::now());
178		− request.response_rx.await?.map_err(\|err\| anyhow!(err))?;
	146	+ fs::write(
	147	+ LOGS_DIR.join("prediction_prompt.md"),
	148	+ &request.local_prompt.unwrap_or_default(),
	149	+ )?;
	150	+
	151	+ let response = request.response_rx.await?.0.map_err(\|err\| anyhow!(err))?;
179	152	prediction_finished_at = Some(Instant::now());
180	153
	154	+ fs::write(
	155	+ LOGS_DIR.join("prediction_response.json"),
	156	+ &serde_json::to_string_pretty(&response).unwrap(),
	157	+ )?;
	158	+
181	159	for included_file in request.request.included_files {
182	160	let insertions = vec![(request.request.cursor_point, CURSOR_MARKER)];
183	161	result
@@ -201,7 +179,6 @@ pub async fn zeta2_predict( · 17 unmodified lines
201	179	}
202	180	break;
203	181	}
204		− _ => {}
205	182	}
206	183	}
207	184

@@ -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"));