Skip to repository content1679 lines · 66.0 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T05:01:13.651Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
format_prompt.rs
1use crate::{
2 FormatPromptArgs, PredictionProvider,
3 example::{ActualCursor, Example, ExamplePrompt},
4 headless::EpAppState,
5 progress::{ExampleProgress, Step},
6 retrieve_context::{ContextRetrievalType, run_context_retrieval},
7};
8use anyhow::{Context as _, Result, anyhow};
9use gpui::AsyncApp;
10use std::ops::Range;
11use std::sync::Arc;
12use zeta_prompt::{
13 Zeta2PromptInput, ZetaFormat, format_edit_history_within_budget, format_expected_output,
14 format_zeta_prompt,
15 hashed_regions::{self, SnippetMarkers},
16 max_edit_event_count_for_format, resolve_cursor_region,
17};
18
19fn resolved_excerpt_ranges_for_format(
20 input: &zeta_prompt::Zeta2PromptInput,
21 format: ZetaFormat,
22) -> (Range<usize>, Range<usize>) {
23 let (_, editable_range_in_context, context_range, _) = resolve_cursor_region(input, format);
24 let editable_range = (context_range.start + editable_range_in_context.start)
25 ..(context_range.start + editable_range_in_context.end);
26 (editable_range, context_range)
27}
28
29pub async fn run_format_prompt(
30 example: &mut Example,
31 args: &FormatPromptArgs,
32 app_state: Arc<EpAppState>,
33 example_progress: &ExampleProgress,
34 cx: AsyncApp,
35) -> Result<()> {
36 run_context_retrieval(
37 example,
38 app_state.clone(),
39 example_progress,
40 vec![ContextRetrievalType::Lsp],
41 false,
42 cx.clone(),
43 )
44 .await?;
45
46 // Teacher-jumps addresses every edit through related-file excerpts and
47 // hard-errors unless the cursor file is covered by one. Settled-data
48 // samples carry a `cursor_excerpt` but were not run through current-file
49 // context retrieval, so normalize the input to synthesize a current-file
50 // excerpt from it when the cursor isn't already covered (a no-op for proper
51 // `ep context --type=current-file` runs).
52 if matches!(
53 args.provider,
54 PredictionProvider::TeacherJumps(_) | PredictionProvider::TeacherJumpsNonBatching(_)
55 ) {
56 if let Some(prompt_inputs) = example.prompt_inputs.as_mut() {
57 hashed_regions::ensure_cursor_file_excerpt(prompt_inputs);
58 }
59 }
60
61 let step_progress = example_progress.start(Step::FormatPrompt);
62
63 let prompt_inputs = example
64 .prompt_inputs
65 .as_ref()
66 .context("prompt_inputs must be set after context retrieval")?;
67
68 match args.provider {
69 PredictionProvider::Teacher(_, zeta_format)
70 | PredictionProvider::TeacherNonBatching(_, zeta_format) => {
71 step_progress.set_substatus("formatting teacher prompt");
72
73 let (editable_range, context_range) =
74 resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);
75
76 let include_diagnostics = matches!(zeta_format, ZetaFormat::V0420Diagnostics);
77
78 let prompt = TeacherPrompt::format_prompt(
79 example,
80 editable_range,
81 context_range,
82 include_diagnostics,
83 );
84 example.prompt = Some(ExamplePrompt {
85 input: prompt,
86 expected_output: None,
87 rejected_output: None,
88 prefill: None,
89 provider: args.provider,
90 });
91 }
92 PredictionProvider::TeacherJumps(_) | PredictionProvider::TeacherJumpsNonBatching(_) => {
93 step_progress.set_substatus("formatting teacher jumps prompt");
94
95 let prompt = TeacherJumpsPrompt::format_prompt(example, args.related_files_budget)?;
96 example.prompt = Some(ExamplePrompt {
97 input: prompt,
98 expected_output: None,
99 rejected_output: None,
100 prefill: None,
101 provider: args.provider,
102 });
103 }
104 PredictionProvider::Zeta2(zeta_format) => {
105 step_progress.set_substatus("formatting zeta2 prompt");
106
107 let prompt = format_zeta_prompt(prompt_inputs, zeta_format);
108 let prefill = zeta_prompt::get_prefill(prompt_inputs, zeta_format);
109 let expected_output = example
110 .spec
111 .expected_patches_with_cursor_positions()
112 .into_iter()
113 .next()
114 .and_then(|(expected_patch, expected_cursor_offset)| {
115 format_expected_output(
116 prompt_inputs,
117 zeta_format,
118 &expected_patch,
119 expected_cursor_offset,
120 )
121 .ok()
122 });
123
124 let rejected_output = example.spec.rejected_patch.as_ref().and_then(|patch| {
125 format_expected_output(prompt_inputs, zeta_format, patch, None).ok()
126 });
127
128 example.prompt = prompt.map(|prompt| ExamplePrompt {
129 input: prompt,
130 expected_output,
131 rejected_output,
132 provider: args.provider,
133 prefill: Some(prefill),
134 });
135 }
136 _ => {
137 panic!("Cannot format prompt for {:?}", args.provider);
138 }
139 };
140 Ok(())
141}
142
143pub struct TeacherPrompt;
144
145impl TeacherPrompt {
146 pub(crate) const EDITABLE_REGION_START: &str = "<|editable_region_start|>\n";
147 pub(crate) const EDITABLE_REGION_END: &str = "\n<|editable_region_end|>";
148 pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";
149 pub(crate) const NO_EDITS: &str = "NO_EDITS";
150
151 /// Truncate edit history to this number of last lines
152 const MAX_HISTORY_LINES: usize = 128;
153
154 pub fn format_prompt(
155 example: &Example,
156 editable_range: Range<usize>,
157 context_range: Range<usize>,
158 include_diagnostics: bool,
159 ) -> String {
160 let edit_history = Self::format_edit_history(&example.spec.edit_history);
161 let context = Self::format_context(example);
162 let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range);
163 let diagnostics = include_diagnostics
164 .then(|| Self::format_diagnostics(example))
165 .map(|diagnostics| format!("# 4. Diagnostics\n\n{diagnostics}"));
166
167 let prompt_template = crate::prompt_assets::get_prompt("teacher.md");
168 let prompt = prompt_template
169 .replace("{{context}}", &context)
170 .replace("{{edit_history}}", &edit_history)
171 .replace("{{diagnostics}}", diagnostics.as_deref().unwrap_or(""))
172 .replace("{{cursor_excerpt}}", &cursor_excerpt);
173
174 prompt
175 }
176
177 pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {
178 // Check if the model indicated no edits are needed
179 let no_edits = (String::new(), None);
180 if let Some(last_codeblock) = extract_last_codeblock(&response) {
181 if last_codeblock.trim() == Self::NO_EDITS {
182 return Ok(no_edits);
183 }
184 }
185
186 if response
187 .trim_end_matches(&[' ', '\n', '`'])
188 .ends_with(Self::NO_EDITS)
189 {
190 return Ok(no_edits);
191 }
192
193 // Extract updated (new) editable region from the model response.
194 let new_editable_region = Self::extract_editable_region(&response)?;
195 let cursor_offset = new_editable_region.find(Self::USER_CURSOR_MARKER);
196 let mut new_editable_region = new_editable_region.replace(Self::USER_CURSOR_MARKER, "");
197 let old_editable_region = Self::extract_editable_region(
198 &example
199 .prompt
200 .as_ref()
201 .context("example prompt missing")?
202 .input,
203 )?
204 .replace(Self::USER_CURSOR_MARKER, "");
205
206 let prompt_inputs = example
207 .prompt_inputs
208 .as_ref()
209 .context("example is missing prompt inputs")?;
210
211 // Normalize leading newlines: if old starts with newline but new doesn't,
212 // prepend newline to new to preserve whitespace structure.
213 // This handles the case where the model drops the leading blank line.
214 if old_editable_region.starts_with('\n') && !new_editable_region.starts_with('\n') {
215 new_editable_region.insert(0, '\n');
216 }
217
218 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
219 let (editable_region_offset, _) = excerpt
220 .match_indices(&old_editable_region)
221 .min_by_key(|(index, _)| index.abs_diff(prompt_inputs.cursor_offset_in_excerpt))
222 .context("editable region not found in prompt content")?;
223 let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
224
225 let editable_region_lines = old_editable_region.lines().count() as u32;
226 let diff = language::unified_diff_with_context(
227 &old_editable_region,
228 &new_editable_region,
229 editable_region_start_line as u32,
230 editable_region_start_line as u32,
231 editable_region_lines,
232 );
233
234 let diff = indoc::formatdoc! {"
235 --- a/{path}
236 +++ b/{path}
237 {diff}",
238 path = example.spec.cursor_path.to_string_lossy(),
239 diff = diff,
240 };
241
242 let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
243 ActualCursor::from_editable_region(
244 &example.spec.cursor_path,
245 editable_region_cursor_offset,
246 &new_editable_region,
247 excerpt,
248 editable_region_offset,
249 editable_region_start_line,
250 )
251 });
252
253 Ok((diff, actual_cursor))
254 }
255
256 fn format_edit_history(edit_history: &str) -> String {
257 let lines: Vec<&str> = edit_history.lines().collect();
258
259 if lines.is_empty() {
260 return "(No edit history)".to_string();
261 }
262
263 if lines.len() > Self::MAX_HISTORY_LINES {
264 let truncated = lines[lines.len() - Self::MAX_HISTORY_LINES..].join("\n");
265 format!("{truncated}\n[...truncated...]")
266 } else {
267 lines.join("\n")
268 }
269 }
270
271 pub fn format_context(example: &Example) -> String {
272 let related_files = example
273 .prompt_inputs
274 .as_ref()
275 .and_then(|pi| pi.related_files.as_deref());
276
277 let Some(related_files) = related_files else {
278 return "(No context)".to_string();
279 };
280
281 if related_files.is_empty() {
282 return "(No context)".to_string();
283 }
284
285 let prefix = "`````";
286 let suffix = "`````\n\n";
287 let max_tokens = 1024;
288 zeta_prompt::format_related_files_within_budget(related_files, &prefix, &suffix, max_tokens)
289 }
290
291 fn format_cursor_excerpt(
292 example: &Example,
293 editable_range: Range<usize>,
294 context_range: Range<usize>,
295 ) -> String {
296 let mut result = String::new();
297
298 let prompt_inputs = example.prompt_inputs.as_ref().unwrap();
299 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
300 let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
301
302 let path_str = example.spec.cursor_path.to_string_lossy();
303 result.push_str(&format!("`````{path_str}\n"));
304 result.push_str(&excerpt[context_range.start..editable_range.start]);
305 result.push_str(Self::EDITABLE_REGION_START);
306 result.push_str(&excerpt[editable_range.start..cursor_offset]);
307 result.push_str(Self::USER_CURSOR_MARKER);
308 result.push_str(&excerpt[cursor_offset..editable_range.end]);
309 result.push_str(Self::EDITABLE_REGION_END);
310 result.push_str(&excerpt[editable_range.end..context_range.end]);
311 result.push_str("\n`````");
312
313 result
314 }
315
316 pub fn extract_editable_region(text: &str) -> Result<String> {
317 let start = text
318 .rfind(Self::EDITABLE_REGION_START)
319 .map_or(0, |pos| pos + Self::EDITABLE_REGION_START.len());
320 let end = text.rfind(Self::EDITABLE_REGION_END).unwrap_or(text.len());
321
322 if start >= end {
323 return Err(anyhow!("Invalid editable region markers"));
324 }
325
326 let region = &text[start..end];
327 Ok(region.strip_suffix('\n').unwrap_or(region).to_string())
328 }
329
330 fn format_diagnostics(example: &Example) -> String {
331 let Some(prompt_inputs) = example.prompt_inputs.as_ref() else {
332 return "No Diagnostics".to_string();
333 };
334
335 let cursor_buffer_row = prompt_inputs.excerpt_start_row.map(|excerpt_start_row| {
336 excerpt_start_row
337 + prompt_inputs.cursor_excerpt[..prompt_inputs.cursor_offset_in_excerpt]
338 .bytes()
339 .filter(|byte| *byte == b'\n')
340 .count() as u32
341 });
342 let diagnostics = zeta_prompt::format_active_buffer_diagnostics_with_budget(
343 &prompt_inputs.active_buffer_diagnostics,
344 cursor_buffer_row,
345 2_000,
346 );
347
348 let diagnostics = diagnostics
349 .strip_prefix("<filename>diagnostics\n")
350 .unwrap_or(&diagnostics);
351
352 if diagnostics.is_empty() {
353 "No Diagnostics".to_string()
354 } else {
355 diagnostics.to_string()
356 }
357 }
358}
359
360/// Teacher prompt for long-range edit prediction ("jumps"). All prompt
361/// context — the cursor file and every related-file excerpt — is annotated
362/// with hashed region markers (V0609HashedRegions), and the teacher may
363/// output a sequence of marker-bounded edits targeting any of it.
364pub struct TeacherJumpsPrompt;
365
366impl TeacherJumpsPrompt {
367 pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";
368 pub(crate) const NO_EDITS: &str = "NO_EDITS";
369
370 const MAX_HISTORY_TOKENS: usize = 4000;
371
372 pub const DEFAULT_RELATED_FILES_BUDGET: usize = 8192;
373
374 pub fn format_prompt(example: &Example, related_files_budget: usize) -> Result<String> {
375 let prompt_inputs = example
376 .prompt_inputs
377 .as_ref()
378 .context("example is missing prompt inputs")?;
379 let marker_table = hashed_regions::build_marker_table(prompt_inputs);
380 let cursor = hashed_regions::locate_cursor_in_related_files(prompt_inputs).context(
381 "cursor position is not covered by any related-file excerpt of the cursor file; \
382 teacher-jumps requires current-file context retrieval (e.g. `ep context --type=current-file,...`)",
383 )?;
384
385 let edit_history = Self::format_edit_history(&prompt_inputs);
386 let context = Self::format_context(
387 prompt_inputs,
388 &marker_table,
389 related_files_budget,
390 cursor.file_ix,
391 );
392 let cursor_excerpt =
393 Self::format_cursor_excerpt(example, prompt_inputs, &marker_table, &cursor)?;
394
395 let prompt_template = crate::prompt_assets::get_prompt("teacher_jumps.md");
396 let prompt = prompt_template
397 .replace("{{context}}", &context)
398 .replace("{{edit_history}}", &edit_history)
399 .replace("{{cursor_excerpt}}", &cursor_excerpt);
400
401 Ok(prompt)
402 }
403
404 pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {
405 let no_edits = (String::new(), None);
406 if let Some(last_codeblock) = extract_last_codeblock(&response) {
407 if last_codeblock.trim() == Self::NO_EDITS {
408 return Ok(no_edits);
409 }
410 }
411
412 if response.trim().ends_with(Self::NO_EDITS) {
413 return Ok(no_edits);
414 }
415
416 let prompt_inputs = example
417 .prompt_inputs
418 .as_ref()
419 .context("example is missing prompt inputs")?;
420
421 // The teacher emits reasoning plus a sequence of markdown code fences,
422 // one per edit, each a marker-bounded span. Extract the spans from the
423 // fences, then hand off to the shared hash-region patch assembler that
424 // the student parser also uses.
425 let codeblocks: Vec<String> = extract_all_codeblocks(response)
426 .into_iter()
427 .filter(|block| block.contains(hashed_regions::MARKER_TAG_PREFIX))
428 .collect();
429 if codeblocks.is_empty() {
430 return Err(anyhow!(
431 "no marker-bounded edit codeblocks found in model response"
432 ));
433 }
434
435 let mut spans = Vec::with_capacity(codeblocks.len());
436 for codeblock in &codeblocks {
437 spans.push(hashed_regions::extract_marker_span(codeblock)?);
438 }
439
440 let (patch, cursor) = hashed_regions::build_patch_from_spans(
441 prompt_inputs,
442 &spans,
443 Self::USER_CURSOR_MARKER,
444 )?;
445
446 let actual_cursor = cursor.map(|cursor| {
447 ActualCursor::from_editable_region(
448 &cursor.path,
449 cursor.cursor_offset_in_new_text,
450 &cursor.new_text,
451 &cursor.old_text,
452 0,
453 cursor.start_row as usize,
454 )
455 });
456
457 Ok((patch, actual_cursor))
458 }
459
460 fn format_edit_history(prompt_inputs: &Zeta2PromptInput) -> String {
461 format_edit_history_within_budget(
462 &prompt_inputs.events,
463 "",
464 "",
465 Self::MAX_HISTORY_TOKENS,
466 max_edit_event_count_for_format(&ZetaFormat::V0327SingleFile),
467 )
468 }
469
470 /// Render related files with hashed region markers, within a token
471 /// budget. Mirrors `zeta_prompt::format_related_files_within_budget`,
472 /// but inserts marker tags into every included excerpt. The cursor file
473 /// is skipped: it renders in its own prompt section via
474 /// `format_cursor_excerpt`, and including it here would duplicate it.
475 fn format_context(
476 prompt_inputs: &Zeta2PromptInput,
477 marker_table: &[SnippetMarkers],
478 max_tokens: usize,
479 cursor_file_ix: usize,
480 ) -> String {
481 let Some(related_files) = prompt_inputs.related_files.as_deref() else {
482 return "(No context)".to_string();
483 };
484 if related_files.is_empty() {
485 return "(No context)".to_string();
486 }
487
488 let estimate_tokens = |bytes: usize| bytes / 3;
489
490 struct RenderedExcerpt {
491 file_ix: usize,
492 excerpt_ix: usize,
493 order: usize,
494 rendered: String,
495 }
496
497 let mut candidates = Vec::new();
498 for (file_ix, file) in related_files.iter().enumerate() {
499 if file_ix == cursor_file_ix {
500 continue;
501 }
502 for (excerpt_ix, excerpt) in file.excerpts.iter().enumerate() {
503 let markers = marker_table.iter().find_map(|snippet| {
504 (snippet.file_ix == file_ix && snippet.excerpt_ix == excerpt_ix)
505 .then_some(&snippet.markers)
506 });
507 let mut rendered = String::new();
508 match markers {
509 Some(markers) => hashed_regions::write_snippet_with_markers(
510 &mut rendered,
511 &excerpt.text,
512 markers,
513 None,
514 ),
515 None => rendered.push_str(&excerpt.text),
516 }
517 if !rendered.ends_with('\n') {
518 rendered.push('\n');
519 }
520 candidates.push(RenderedExcerpt {
521 file_ix,
522 excerpt_ix,
523 order: excerpt.order,
524 rendered,
525 });
526 }
527 }
528
529 let file_headers: Vec<String> = related_files
530 .iter()
531 .map(|file| format!("`````{}\n", file.path.to_string_lossy()))
532 .collect();
533 let file_suffix = "`````\n\n";
534
535 let mut selection_order: Vec<usize> = (0..candidates.len()).collect();
536 selection_order.sort_by_key(|&candidate_ix| {
537 let candidate = &candidates[candidate_ix];
538 (candidate.order, candidate.file_ix, candidate.excerpt_ix)
539 });
540
541 let mut total_tokens = 0;
542 let mut included = vec![false; candidates.len()];
543 let mut file_included = vec![false; related_files.len()];
544 for &candidate_ix in &selection_order {
545 let candidate = &candidates[candidate_ix];
546 let header_cost = if file_included[candidate.file_ix] {
547 0
548 } else {
549 estimate_tokens(file_headers[candidate.file_ix].len() + file_suffix.len())
550 };
551 let excerpt_cost = estimate_tokens(candidate.rendered.len());
552 if total_tokens + header_cost + excerpt_cost > max_tokens {
553 break;
554 }
555 total_tokens += header_cost + excerpt_cost;
556 file_included[candidate.file_ix] = true;
557 included[candidate_ix] = true;
558 }
559
560 let mut result = String::new();
561 let mut last_file_ix = None;
562 for (candidate_ix, candidate) in candidates.iter().enumerate() {
563 if !included[candidate_ix] {
564 continue;
565 }
566 if last_file_ix != Some(candidate.file_ix) {
567 if last_file_ix.is_some() {
568 result.push_str(file_suffix);
569 }
570 result.push_str(&file_headers[candidate.file_ix]);
571 last_file_ix = Some(candidate.file_ix);
572 }
573 result.push_str(&candidate.rendered);
574
575 let file = &related_files[candidate.file_ix];
576 let excerpt = &file.excerpts[candidate.excerpt_ix];
577 let next_excerpt_start = candidates
578 .iter()
579 .enumerate()
580 .skip(candidate_ix + 1)
581 .find(|(next_ix, next)| included[*next_ix] && next.file_ix == candidate.file_ix)
582 .map(|(_, next)| file.excerpts[next.excerpt_ix].row_range.start);
583 if zeta_prompt::rows_omitted_after_excerpt(excerpt, next_excerpt_start, file.max_row) {
584 result.push_str("...\n");
585 }
586 }
587 if last_file_ix.is_some() {
588 result.push_str(file_suffix);
589 }
590
591 if result.is_empty() {
592 "(No context)".to_string()
593 } else {
594 result
595 }
596 }
597
598 /// Render the current file from its related-file entry, with marker tags
599 /// and the user cursor injected. The current file gets its own prompt
600 /// section but shares the related-file snippets and markers, so its
601 /// content appears in the prompt exactly once.
602 fn format_cursor_excerpt(
603 example: &Example,
604 prompt_inputs: &Zeta2PromptInput,
605 marker_table: &[SnippetMarkers],
606 cursor: &hashed_regions::RelatedFileCursor,
607 ) -> Result<String> {
608 let related_files = prompt_inputs
609 .related_files
610 .as_deref()
611 .context("prompt inputs are missing related files")?;
612 let file = related_files
613 .get(cursor.file_ix)
614 .context("cursor file index out of range")?;
615
616 let path_str = example.spec.cursor_path.to_string_lossy();
617 let mut result = format!("`````{path_str}\n");
618 for (excerpt_ix, excerpt) in file.excerpts.iter().enumerate() {
619 let markers = marker_table
620 .iter()
621 .find_map(|snippet| {
622 (snippet.file_ix == cursor.file_ix && snippet.excerpt_ix == excerpt_ix)
623 .then_some(&snippet.markers)
624 })
625 .context("marker table is missing a cursor file snippet")?;
626 let cursor_in_excerpt = (excerpt_ix == cursor.excerpt_ix)
627 .then_some((cursor.offset_in_excerpt, Self::USER_CURSOR_MARKER));
628 hashed_regions::write_snippet_with_markers(
629 &mut result,
630 &excerpt.text,
631 markers,
632 cursor_in_excerpt,
633 );
634 if !result.ends_with('\n') {
635 result.push('\n');
636 }
637 if excerpt.row_range.end < file.max_row {
638 result.push_str("...\n");
639 }
640 }
641 result.push_str("`````");
642
643 Ok(result)
644 }
645}
646
647pub(crate) fn line_start_offset(text: &str, row: usize) -> Option<usize> {
648 let mut offset = 0;
649 for _ in 0..row {
650 offset += text[offset..].find('\n')? + 1;
651 }
652 Some(offset)
653}
654
655/// Extract the cursor excerpt from an example.
656/// First tries to extract from an existing prompt, then falls back to constructing from prompt_inputs.
657pub fn extract_cursor_excerpt_from_example(example: &Example) -> Option<String> {
658 // If we have the original prompt, extract the cursor excerpt from it
659 if let Some(prompt) = &example.prompt {
660 // Find "# 3. Current File" section and extract the content
661 if let Some(start) = prompt.input.find("# 3. Current File") {
662 let content_start = prompt.input[start..].find('`').map(|i| start + i)?;
663 let backtick_count = prompt.input[content_start..]
664 .chars()
665 .take_while(|&c| c == '`')
666 .count();
667 let content_start = content_start + backtick_count;
668
669 // Find the path line and skip it
670 let newline_pos = prompt.input[content_start..].find('\n')?;
671 let text_start = content_start + newline_pos + 1;
672
673 // Find the closing backticks
674 let closing_pattern = "`".repeat(backtick_count);
675 let text_end = prompt.input[text_start..].find(&closing_pattern)?;
676 let cursor_excerpt = &prompt.input[text_start..text_start + text_end];
677
678 let path_str = example.spec.cursor_path.to_string_lossy();
679 return Some(format!("`````{path_str}\n{cursor_excerpt}`````"));
680 }
681 }
682
683 // Fallback: construct from prompt_inputs if available
684 let prompt_inputs = example.prompt_inputs.as_ref()?;
685 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
686 let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
687
688 // Simple fallback: just show content around cursor with markers
689 let path_str = example.spec.cursor_path.to_string_lossy();
690 let mut result = format!("`````{path_str}\n");
691 result.push_str(TeacherPrompt::EDITABLE_REGION_START);
692 result.push_str(&excerpt[..cursor_offset]);
693 result.push_str(TeacherPrompt::USER_CURSOR_MARKER);
694 result.push_str(&excerpt[cursor_offset..]);
695 result.push_str(TeacherPrompt::EDITABLE_REGION_END);
696 result.push_str("\n`````");
697
698 Some(result)
699}
700
701/// Extract all top-level fenced codeblocks from `text`, in order.
702///
703/// A fence opens with 3+ backticks (optionally followed by an info string)
704/// and closes with a line of at least as many backticks, so codeblocks that
705/// themselves contain shorter fences are handled.
706pub(crate) fn extract_all_codeblocks(text: &str) -> Vec<String> {
707 let mut codeblocks = Vec::new();
708 let mut current_block: Option<(usize, Vec<&str>)> = None;
709
710 for line in text.lines() {
711 match &mut current_block {
712 None => {
713 let backtick_count = line.chars().take_while(|&c| c == '`').count();
714 if backtick_count >= 3 {
715 current_block = Some((backtick_count, Vec::new()));
716 }
717 }
718 Some((opening_count, lines)) => {
719 let trimmed = line.trim();
720 if trimmed.len() >= *opening_count && trimmed.chars().all(|c| c == '`') {
721 let mut content = lines.join("\n");
722 if !content.is_empty() {
723 content.push('\n');
724 }
725 codeblocks.push(content);
726 current_block = None;
727 } else {
728 lines.push(line);
729 }
730 }
731 }
732 }
733
734 codeblocks
735}
736
737pub(crate) fn extract_last_codeblock(text: &str) -> Option<String> {
738 let lines: Vec<&str> = text.lines().collect();
739
740 // Search from the end for a closing fence (line containing only backticks, 3+)
741 let mut closing_line_idx = None;
742 let mut backtick_count = 0;
743
744 for i in (0..lines.len()).rev() {
745 let line = lines[i].trim();
746 if line.len() >= 3 && line.chars().all(|c| c == '`') {
747 closing_line_idx = Some(i);
748 backtick_count = line.len();
749 break;
750 }
751 }
752
753 let closing_idx = closing_line_idx?;
754
755 // Search backwards for matching opening fence
756 // Opening fence starts with same backtick count, possibly followed by language/metadata
757 let opening_pattern = "`".repeat(backtick_count);
758
759 for i in (0..closing_idx).rev() {
760 let line = lines[i];
761 if line.starts_with(&opening_pattern) {
762 // Ensure it's exactly the right number of backticks (not more)
763 let rest = &line[backtick_count..];
764 if rest.is_empty() || !rest.starts_with('`') {
765 // Found matching opening fence
766 // Extract content between opening and closing (exclusive)
767 if closing_idx > i + 1 {
768 let content = lines[i + 1..closing_idx].join("\n");
769 // Preserve trailing newline to match previous behavior
770 return Some(format!("{}\n", content));
771 } else {
772 // Empty block
773 return Some(String::new());
774 }
775 }
776 }
777 }
778
779 None
780}
781
782#[cfg(test)]
783mod tests {
784 use super::*;
785 use zeta_prompt::multi_region;
786
787 fn make_example(
788 cursor_excerpt: &str,
789 cursor_offset: usize,
790 related: &[(&str, &[(&str, u32)])],
791 ) -> Example {
792 // The cursor file is included as the first related file, mirroring
793 // `ContextSource::CurrentFile` context retrieval.
794 let cursor_file = zeta_prompt::RelatedFile {
795 path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
796 max_row: 1000,
797 excerpts: vec![zeta_prompt::RelatedExcerpt {
798 row_range: 0..cursor_excerpt.matches('\n').count() as u32,
799 text: std::sync::Arc::from(cursor_excerpt),
800 order: 0,
801 context_source: zeta_prompt::ContextSource::CurrentFile,
802 }],
803 in_open_source_repo: false,
804 };
805 let related_files = std::iter::once(cursor_file)
806 .chain(related.iter().map(|(path, excerpts)| {
807 zeta_prompt::RelatedFile {
808 path: std::sync::Arc::from(std::path::Path::new(path)),
809 max_row: 1000,
810 excerpts: excerpts
811 .iter()
812 .map(|(text, start_row)| zeta_prompt::RelatedExcerpt {
813 row_range: *start_row..*start_row + text.matches('\n').count() as u32,
814 text: std::sync::Arc::from(*text),
815 order: 0,
816 context_source: zeta_prompt::ContextSource::CurrentFile,
817 })
818 .collect(),
819 in_open_source_repo: false,
820 }
821 }))
822 .collect();
823
824 Example {
825 spec: edit_prediction::example_spec::ExampleSpec {
826 name: "test".to_string(),
827 repository_url: "https://github.com/zed-industries/zed.git".to_string(),
828 revision: "HEAD".to_string(),
829 tags: Vec::new(),
830 reasoning: None,
831 uncommitted_diff: String::new(),
832 recently_opened_files: Vec::new(),
833 recently_viewed_files: Vec::new(),
834 uncommitted_diff_contains_edit_history: false,
835 cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
836 cursor_position: "0:0".to_string(),
837 edit_history: String::new(),
838 expected_patches: Vec::new(),
839 rejected_patch: None,
840 telemetry: None,
841 human_feedback: Vec::new(),
842 rating: None,
843 },
844 prompt_inputs: Some(zeta_prompt::Zeta2PromptInput {
845 cursor_path: std::path::Path::new("src/main.rs").into(),
846 cursor_excerpt: cursor_excerpt.into(),
847 cursor_offset_in_excerpt: cursor_offset,
848 excerpt_start_row: Some(0),
849 events: Vec::new(),
850 related_files: Some(related_files),
851 active_buffer_diagnostics: Vec::new(),
852 excerpt_ranges: zeta_prompt::ExcerptRanges::default(),
853 syntax_ranges: None,
854 in_open_source_repo: false,
855 can_collect_data: false,
856 repo_url: None,
857 }),
858 prompt: None,
859 predictions: Vec::new(),
860 score: Vec::new(),
861 qa: Vec::new(),
862 zed_version: None,
863 state: None,
864 }
865 }
866
867 #[test]
868 fn test_teacher_jumps_format_prompt_markers_everywhere() {
869 let example = make_example(
870 "fn main() {\n let x = 1;\n}\n",
871 16,
872 &[("src/lib.rs", &[("pub fn helper() {}\n", 5)])],
873 );
874 let prompt = TeacherJumpsPrompt::format_prompt(&example, 8192).unwrap();
875
876 assert!(prompt.contains(TeacherJumpsPrompt::USER_CURSOR_MARKER));
877 assert!(prompt.contains("`````src/main.rs\n"));
878 assert!(prompt.contains("`````src/lib.rs\n"));
879 // Markers in both the current file and the related excerpt.
880 let marker_table =
881 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
882 for snippet in &marker_table {
883 for (id, _) in &snippet.markers {
884 assert!(
885 prompt.contains(&hashed_regions::marker_tag(id)),
886 "prompt is missing marker {id}"
887 );
888 }
889 }
890 // The current file appears exactly once, in its own section, with the
891 // user cursor injected.
892 assert_eq!(prompt.matches("let x = 1;").count(), 1);
893 assert!(prompt.contains("<|user_cursor|>let x = 1;"));
894 }
895
896 #[test]
897 fn test_teacher_jumps_cursor_file_with_coinciding_worktree_root_name() {
898 // Worktree root `jaq` contains a `jaq/` subdirectory: the cursor path
899 // is `jaq/src/main.rs` while the related-file entry is prefixed with
900 // the root name (`jaq/jaq/src/main.rs`).
901 let mut example = make_example("fn main() {\n let x = 1;\n}\n", 16, &[]);
902 example.spec.cursor_path = std::sync::Arc::from(std::path::Path::new("jaq/src/main.rs"));
903 {
904 let prompt_inputs = example.prompt_inputs.as_mut().unwrap();
905 prompt_inputs.cursor_path = std::path::Path::new("jaq/src/main.rs").into();
906 prompt_inputs.related_files.as_mut().unwrap()[0].path =
907 std::sync::Arc::from(std::path::Path::new("jaq/jaq/src/main.rs"));
908 }
909
910 let prompt = TeacherJumpsPrompt::format_prompt(&example, 8192).unwrap();
911 assert!(prompt.contains("<|user_cursor|>let x = 1;"));
912
913 let marker_table =
914 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
915 let cursor_markers = &marker_table[0].markers;
916 let start_tag = hashed_regions::marker_tag(&cursor_markers[0].0);
917 let end_tag = hashed_regions::marker_tag(&cursor_markers[cursor_markers.len() - 1].0);
918 let response =
919 format!("`````\n{start_tag}\nfn main() {{\n let x = 2;\n}}\n{end_tag}\n`````\n");
920 let (patch, _) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
921 assert!(patch.contains("--- a/jaq/src/main.rs"), "patch: {patch}");
922 }
923
924 #[test]
925 fn test_teacher_jumps_format_prompt_requires_current_file_context() {
926 let mut example = make_example("fn main() {}\n", 0, &[]);
927 example.prompt_inputs.as_mut().unwrap().related_files = Some(Vec::new());
928 assert!(TeacherJumpsPrompt::format_prompt(&example, 8192).is_err());
929 }
930
931 #[test]
932 fn test_teacher_jumps_synthesizes_missing_cursor_file_excerpt() {
933 // Simulate a settled-data sample: `cursor_excerpt` is present, but the
934 // related-file excerpts of the cursor file don't cover the cursor (only
935 // an unrelated fragment elsewhere in the file).
936 let mut example = make_example("fn main() {\n let x = 1;\n}\n", 16, &[]);
937 {
938 let prompt_inputs = example.prompt_inputs.as_mut().unwrap();
939 prompt_inputs.related_files.as_mut().unwrap()[0].excerpts =
940 vec![zeta_prompt::RelatedExcerpt {
941 row_range: 40..42,
942 text: std::sync::Arc::from("// unrelated\n// fragment\n"),
943 order: 0,
944 context_source: zeta_prompt::ContextSource::Bm25,
945 }];
946 }
947
948 // Without a covering cursor-file excerpt, formatting hard-errors.
949 assert!(TeacherJumpsPrompt::format_prompt(&example, 8192).is_err());
950
951 // `ensure_cursor_file_excerpt` (in zeta_prompt) synthesizes one from
952 // `cursor_excerpt`, so the prompt formats with the cursor in the
953 // current-file window and the unrelated fragment is replaced (no
954 // duplicated content with overlapping markers).
955 hashed_regions::ensure_cursor_file_excerpt(example.prompt_inputs.as_mut().unwrap());
956 let prompt = TeacherJumpsPrompt::format_prompt(&example, 8192).unwrap();
957 assert!(prompt.contains("<|user_cursor|>let x = 1;"));
958 assert!(!prompt.contains("// unrelated"));
959 assert_eq!(prompt.matches("let x = 1;").count(), 1);
960 }
961
962 #[test]
963 fn test_teacher_jumps_cursor_file_hunks_are_file_absolute() {
964 let mut example = make_example("fn main() {\n let x = 1;\n}\n", 16, &[]);
965 {
966 let prompt_inputs = example.prompt_inputs.as_mut().unwrap();
967 prompt_inputs.excerpt_start_row = Some(10);
968 prompt_inputs.related_files.as_mut().unwrap()[0].excerpts[0].row_range = 10..13;
969 }
970 let marker_table =
971 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
972 let cursor_markers = &marker_table[0].markers;
973 let start_tag = hashed_regions::marker_tag(&cursor_markers[0].0);
974 let end_tag = hashed_regions::marker_tag(&cursor_markers[cursor_markers.len() - 1].0);
975
976 let response = format!(
977 "The user is changing x.\n\n`````\n{start_tag}\nfn main() {{\n let x = 2;<|user_cursor|>\n}}\n{end_tag}\n`````\n"
978 );
979 let (patch, cursor) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
980
981 // Hunk rows are file-absolute (1-based in the hunk header, excerpt
982 // starts at 0-based row 10).
983 assert!(patch.contains("@@ -11,"), "patch: {patch}");
984 let cursor = cursor.unwrap();
985 assert_eq!(cursor.row, 11);
986 }
987
988 #[test]
989 fn test_teacher_jumps_parse_single_edit_in_cursor_file() {
990 let example = make_example("fn main() {\n let x = 1;\n}\n", 16, &[]);
991 let marker_table =
992 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
993 let cursor_markers = &marker_table[0].markers;
994 let start_tag = hashed_regions::marker_tag(&cursor_markers[0].0);
995 let end_tag = hashed_regions::marker_tag(&cursor_markers[cursor_markers.len() - 1].0);
996
997 let response = format!(
998 "The user is changing x.\n\n`````\n{start_tag}\nfn main() {{\n let x = 2;<|user_cursor|>\n}}\n{end_tag}\n`````\n"
999 );
1000 let (patch, cursor) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1001
1002 assert!(patch.contains("--- a/src/main.rs"), "patch: {patch}");
1003 assert!(patch.contains("- let x = 1;"), "patch: {patch}");
1004 assert!(patch.contains("+ let x = 2;"), "patch: {patch}");
1005 let cursor = cursor.unwrap();
1006 assert_eq!(cursor.path, "src/main.rs");
1007 assert_eq!(cursor.row, 1);
1008 }
1009
1010 #[test]
1011 fn test_teacher_jumps_parse_sequence_across_files() {
1012 let example = make_example(
1013 "fn fetch_user_cached() {}\n",
1014 0,
1015 &[(
1016 "src/server.rs",
1017 &[("fn handle() {\n fetch_user();\n}\n", 10)],
1018 )],
1019 );
1020 let marker_table =
1021 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1022 assert_eq!(marker_table.len(), 2);
1023 let related_markers = &marker_table[1].markers;
1024 let start_tag = hashed_regions::marker_tag(&related_markers[0].0);
1025 let end_tag = hashed_regions::marker_tag(&related_markers[related_markers.len() - 1].0);
1026
1027 let response = format!(
1028 "Updating the call site to use the new name.\n\n\
1029 `````\n{start_tag}\nfn handle() {{\n fetch_user_cached();\n}}\n{end_tag}\n`````\n"
1030 );
1031 let (patch, cursor) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1032
1033 assert!(patch.contains("--- a/src/server.rs"), "patch: {patch}");
1034 assert!(patch.contains("- fetch_user();"), "patch: {patch}");
1035 assert!(
1036 patch.contains("+ fetch_user_cached();"),
1037 "patch: {patch}"
1038 );
1039 // Hunk rows are file-absolute for related files (1-based in the
1040 // hunk header, excerpt starts at 0-based row 10).
1041 assert!(patch.contains("@@ -11,"), "patch: {patch}");
1042 assert!(cursor.is_none());
1043 }
1044
1045 #[test]
1046 fn test_teacher_jumps_parse_multiple_edits_same_file() {
1047 let cursor_excerpt = "\
1048 fn alpha() {\n one();\n}\n\nfn beta() {\n two();\n}\n\n\
1049 fn gamma() {\n three();\n}\n\nfn delta() {\n four();\n}\n";
1050 let example = make_example(cursor_excerpt, 0, &[]);
1051 let marker_table =
1052 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1053 let markers = &marker_table[0].markers;
1054 assert!(
1055 markers.len() >= 3,
1056 "expected internal markers, got {markers:?}"
1057 );
1058
1059 // First edit: between the first two markers; second edit: between the
1060 // second and last markers.
1061 let tag = |ix: usize| hashed_regions::marker_tag(&markers[ix].0);
1062 let old_first_span = &cursor_excerpt[markers[0].1..markers[1].1];
1063 let old_second_span = &cursor_excerpt[markers[1].1..markers[markers.len() - 1].1];
1064 let new_first_span = old_first_span.replace("one()", "uno()");
1065 let new_second_span = old_second_span.replace("four()", "cuatro()");
1066
1067 let response = format!(
1068 "Renaming calls.\n\n`````\n{}\n{}{}\n`````\n\n`````\n{}\n{}{}\n`````\n",
1069 tag(0),
1070 new_first_span,
1071 tag(1),
1072 tag(1),
1073 new_second_span,
1074 tag(markers.len() - 1),
1075 );
1076 let (patch, _) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1077
1078 assert!(patch.contains("+ uno();"), "patch: {patch}");
1079 assert!(patch.contains("+ cuatro();"), "patch: {patch}");
1080 assert_eq!(patch.matches("--- a/src/main.rs").count(), 1);
1081 }
1082
1083 #[test]
1084 fn test_teacher_jumps_parse_no_edits() {
1085 let example = make_example("fn main() {}\n", 0, &[]);
1086 let (patch, cursor) =
1087 TeacherJumpsPrompt::parse(&example, "All good.\n\n`````\nNO_EDITS\n`````\n").unwrap();
1088 assert!(patch.is_empty());
1089 assert!(cursor.is_none());
1090 }
1091
1092 #[test]
1093 fn test_teacher_jumps_parse_rejects_truncated_span() {
1094 let cursor_excerpt = "\
1095 fn alpha() {\n one();\n}\n\nfn beta() {\n two();\n}\n\n\
1096 fn gamma() {\n three();\n}\n\nfn delta() {\n four();\n}\n";
1097 let example = make_example(cursor_excerpt, 0, &[]);
1098 let marker_table =
1099 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1100 let markers = &marker_table[0].markers;
1101 assert!(markers.len() >= 3);
1102 let start_tag = hashed_regions::marker_tag(&markers[0].0);
1103 let end_tag = hashed_regions::marker_tag(&markers[markers.len() - 1].0);
1104
1105 // The model reproduces only the head of the span and stops before the
1106 // end marker; accepting this would silently delete the rest.
1107 let head = &cursor_excerpt[markers[0].1..markers[1].1];
1108 let response = format!("Minor cleanup.\n\n`````\n{start_tag}\n{head}{end_tag}\n`````\n");
1109 let error = TeacherJumpsPrompt::parse(&example, &response).unwrap_err();
1110 assert!(
1111 error.to_string().contains("looks truncated"),
1112 "unexpected error: {error}"
1113 );
1114 }
1115
1116 #[test]
1117 fn test_teacher_jumps_parse_rejects_tail_deletion_after_head_edit() {
1118 let cursor_excerpt = "\
1119 fn alpha() {\n one();\n}\n\nfn beta() {\n two();\n}\n\n\
1120 fn gamma() {\n three();\n}\n\nfn delta() {\n four();\n}\n";
1121 let example = make_example(cursor_excerpt, 0, &[]);
1122 let marker_table =
1123 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1124 let markers = &marker_table[0].markers;
1125 assert!(markers.len() >= 3);
1126 let start_tag = hashed_regions::marker_tag(&markers[0].0);
1127 let end_tag = hashed_regions::marker_tag(&markers[markers.len() - 1].0);
1128
1129 // The dropped tail is large enough to trip the trailing-deletion check.
1130 let tail = &cursor_excerpt[markers[1].1..];
1131 assert!(tail.lines().filter(|line| !line.trim().is_empty()).count() > 3);
1132
1133 // The model makes a real edit at the head of the span, reproduces
1134 // some context, and then stops before the end marker. The replacement
1135 // is not a verbatim prefix of the span, but the tail is still
1136 // silently deleted.
1137 let head = &cursor_excerpt[markers[0].1..markers[1].1];
1138 assert!(head.contains("fn alpha()"));
1139 let edited_head = head.replacen("fn alpha()", "fn alpha_renamed()", 1);
1140 let response =
1141 format!("Renaming alpha.\n\n`````\n{start_tag}\n{edited_head}{end_tag}\n`````\n");
1142 let error = TeacherJumpsPrompt::parse(&example, &response).unwrap_err();
1143 assert!(
1144 error.to_string().contains("looks truncated"),
1145 "unexpected error: {error}"
1146 );
1147 }
1148
1149 #[test]
1150 fn test_teacher_jumps_parse_allows_mid_span_deletion() {
1151 let cursor_excerpt = "\
1152 fn alpha() {\n one();\n}\n\nfn beta() {\n two();\n}\n\n\
1153 fn gamma() {\n three();\n}\n\nfn delta() {\n four();\n}\n";
1154 let example = make_example(cursor_excerpt, 0, &[]);
1155 let marker_table =
1156 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1157 let markers = &marker_table[0].markers;
1158 assert!(markers.len() >= 3);
1159 let start_tag = hashed_regions::marker_tag(&markers[0].0);
1160 let end_tag = hashed_regions::marker_tag(&markers[markers.len() - 1].0);
1161
1162 // Deleting code in the middle while reproducing the span's tail shows
1163 // the model kept writing to the end marker, so it must be accepted.
1164 let head = &cursor_excerpt[markers[0].1..markers[1].1];
1165 let reproduced_tail = &cursor_excerpt[markers[markers.len() - 2].1..];
1166 assert!(!reproduced_tail.trim().is_empty());
1167 let response = format!(
1168 "Removing the middle.\n\n`````\n{start_tag}\n{head}{reproduced_tail}{end_tag}\n`````\n"
1169 );
1170 let (patch, _) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1171 assert!(patch.contains("-fn beta() {"), "patch: {patch}");
1172 }
1173
1174 #[test]
1175 fn test_teacher_jumps_parse_allows_small_tail_deletion() {
1176 let cursor_excerpt = "\
1177 fn alpha() {\n one();\n}\n\nfn beta() {\n two();\n}\n\n\
1178 fn gamma() {\n three();\n}\n\nfn delta() {\n four();\n}\n";
1179 let example = make_example(cursor_excerpt, 0, &[]);
1180 let marker_table =
1181 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1182 let markers = &marker_table[0].markers;
1183 let start_tag = hashed_regions::marker_tag(&markers[0].0);
1184 let end_tag = hashed_regions::marker_tag(&markers[markers.len() - 1].0);
1185
1186 // Dropping only the last line of the span may be a genuine
1187 // end-of-snippet deletion, so it stays below the threshold.
1188 let new_span = cursor_excerpt.strip_suffix("}\n").unwrap();
1189 let response =
1190 format!("Dropping the brace.\n\n`````\n{start_tag}\n{new_span}{end_tag}\n`````\n");
1191 let (patch, _) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1192 assert!(patch.contains("-}"), "patch: {patch}");
1193 }
1194
1195 #[test]
1196 fn test_teacher_jumps_parse_allows_empty_span_deletion() {
1197 let cursor_excerpt = "\
1198 fn alpha() {\n one();\n}\n\nfn beta() {\n two();\n}\n\n\
1199 fn gamma() {\n three();\n}\n\nfn delta() {\n four();\n}\n";
1200 let example = make_example(cursor_excerpt, 0, &[]);
1201 let marker_table =
1202 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1203 let markers = &marker_table[0].markers;
1204 assert!(markers.len() >= 3);
1205 let start_tag = hashed_regions::marker_tag(&markers[0].0);
1206 let end_tag = hashed_regions::marker_tag(&markers[1].0);
1207
1208 // Deleting an entire span by replacing it with nothing is fine.
1209 let response = format!("Removing alpha.\n\n`````\n{start_tag}\n{end_tag}\n`````\n");
1210 let (patch, _) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1211 assert!(patch.contains("-fn alpha() {"), "patch: {patch}");
1212 }
1213
1214 #[test]
1215 fn test_teacher_jumps_parse_span_across_contiguous_excerpts() {
1216 // Two excerpts of src/lib.rs with touching row ranges (5..8 and
1217 // 8..11) render seamlessly in the prompt, so the model may span a
1218 // single edit across the excerpt boundary.
1219 let example = make_example(
1220 "fn main() {}\n",
1221 0,
1222 &[(
1223 "src/lib.rs",
1224 &[
1225 ("fn a() {\n one();\n}\n", 5),
1226 ("fn b() {\n two();\n}\n", 8),
1227 ],
1228 )],
1229 );
1230 let marker_table =
1231 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1232 assert_eq!(marker_table.len(), 3);
1233 let start_tag = hashed_regions::marker_tag(&marker_table[1].markers[0].0);
1234 let last_markers = &marker_table[2].markers;
1235 let end_tag = hashed_regions::marker_tag(&last_markers[last_markers.len() - 1].0);
1236
1237 let response = format!(
1238 "Renaming both.\n\n`````\n{start_tag}\nfn a() {{\n uno();\n}}\nfn b() {{\n dos();\n}}\n{end_tag}\n`````\n"
1239 );
1240 let (patch, _) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1241
1242 assert!(patch.contains("+ uno();"), "patch: {patch}");
1243 assert!(patch.contains("+ dos();"), "patch: {patch}");
1244 assert_eq!(patch.matches("--- a/src/lib.rs").count(), 1);
1245 // Hunk rows are file-absolute (merged region starts at 0-based row 5).
1246 assert!(patch.contains("@@ -6,"), "patch: {patch}");
1247 }
1248
1249 #[test]
1250 fn test_teacher_jumps_parse_insertion_at_contiguous_excerpt_seam() {
1251 // The two markers at the seam between contiguous excerpts map to the
1252 // same merged offset; bracketing them expresses a pure insertion.
1253 let example = make_example(
1254 "fn main() {}\n",
1255 0,
1256 &[(
1257 "src/lib.rs",
1258 &[
1259 ("fn a() {\n one();\n}\n", 5),
1260 ("fn b() {\n two();\n}\n", 8),
1261 ],
1262 )],
1263 );
1264 let marker_table =
1265 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1266 assert_eq!(marker_table.len(), 3);
1267 let first_markers = &marker_table[1].markers;
1268 let start_tag = hashed_regions::marker_tag(&first_markers[first_markers.len() - 1].0);
1269 let end_tag = hashed_regions::marker_tag(&marker_table[2].markers[0].0);
1270
1271 let response = format!(
1272 "Adding a function between a and b.\n\n`````\n{start_tag}\nfn between() {{}}\n{end_tag}\n`````\n"
1273 );
1274 let (patch, _) = TeacherJumpsPrompt::parse(&example, &response).unwrap();
1275
1276 assert!(patch.contains("+fn between() {}"), "patch: {patch}");
1277 assert!(
1278 !patch.contains("-\n"),
1279 "patch should be pure insertion: {patch}"
1280 );
1281 // Insertion lands between the excerpts (after 0-based row 7).
1282 assert!(patch.contains("@@ -6,"), "patch: {patch}");
1283 }
1284
1285 #[test]
1286 fn test_teacher_jumps_parse_rejects_span_across_gapped_excerpts() {
1287 // Same file, but the excerpts don't touch (5..8 and 20..23): rows in
1288 // between were never shown to the model, so a span across them is
1289 // invalid.
1290 let example = make_example(
1291 "fn main() {}\n",
1292 0,
1293 &[(
1294 "src/lib.rs",
1295 &[
1296 ("fn a() {\n one();\n}\n", 5),
1297 ("fn b() {\n two();\n}\n", 20),
1298 ],
1299 )],
1300 );
1301 let marker_table =
1302 hashed_regions::build_marker_table(example.prompt_inputs.as_ref().unwrap());
1303 assert_eq!(marker_table.len(), 3);
1304 let start_tag = hashed_regions::marker_tag(&marker_table[1].markers[0].0);
1305 let last_markers = &marker_table[2].markers;
1306 let end_tag = hashed_regions::marker_tag(&last_markers[last_markers.len() - 1].0);
1307
1308 let response = format!(
1309 "Renaming both.\n\n`````\n{start_tag}\nfn a() {{\n uno();\n}}\nfn b() {{\n dos();\n}}\n{end_tag}\n`````\n"
1310 );
1311 let error = TeacherJumpsPrompt::parse(&example, &response).unwrap_err();
1312 assert!(
1313 error.to_string().contains("different context snippets"),
1314 "unexpected error: {error}"
1315 );
1316 }
1317
1318 #[test]
1319 fn test_teacher_jumps_parse_rejects_unknown_marker() {
1320 let example = make_example("fn main() {}\n", 0, &[]);
1321 let response = "`````\n<|marker_zzzz|>\nnew\n<|marker_yyyy|>\n`````\n";
1322 assert!(TeacherJumpsPrompt::parse(&example, response).is_err());
1323 }
1324
1325 #[test]
1326 fn test_extract_all_codeblocks_multiple() {
1327 let text = indoc::indoc! {"
1328 First edit:
1329
1330 `````
1331 block one
1332 `````
1333
1334 Second edit:
1335
1336 `````
1337 block two
1338 with ``` nested
1339 `````
1340 "};
1341 let blocks = extract_all_codeblocks(text);
1342 assert_eq!(
1343 blocks,
1344 vec![
1345 "block one\n".to_string(),
1346 "block two\nwith ``` nested\n".to_string()
1347 ]
1348 );
1349 }
1350
1351 #[test]
1352 fn test_extract_last_code_block() {
1353 let text = indoc::indoc! {"
1354 Some thinking
1355
1356 ```
1357 first block
1358 ```
1359
1360 `````path='something' lines=1:2
1361 last block
1362 `````
1363 "};
1364 let last_block = extract_last_codeblock(text).unwrap();
1365 assert_eq!(last_block, "last block\n");
1366 }
1367
1368 #[test]
1369 fn test_extract_codeblock_with_nested_fences() {
1370 let text = indoc::indoc! {"
1371 `````
1372 content with ``` inline
1373 and ```python nested
1374 more content
1375 `````
1376 "};
1377 let last_block = extract_last_codeblock(text).unwrap();
1378 assert_eq!(
1379 last_block,
1380 "content with ``` inline\nand ```python nested\nmore content\n"
1381 );
1382 }
1383
1384 #[test]
1385 fn test_extract_codeblock_ignores_inline_backticks() {
1386 let text = indoc::indoc! {"
1387 `````
1388 here is some `code` with inline backticks
1389 and here```more```stuff
1390 `````
1391 "};
1392 let last_block = extract_last_codeblock(text).unwrap();
1393 assert_eq!(
1394 last_block,
1395 "here is some `code` with inline backticks\nand here```more```stuff\n"
1396 );
1397 }
1398
1399 #[test]
1400 fn test_extract_editable_region_old_format() {
1401 let text = indoc::indoc! {"
1402 some lines
1403 are
1404 here
1405 <|editable_region_start|>
1406 one
1407 two three
1408
1409 <|editable_region_end|>
1410 more
1411 lines here
1412 "};
1413 let parsed = TeacherPrompt::extract_editable_region(text).unwrap();
1414 assert_eq!(
1415 parsed,
1416 indoc::indoc! {"
1417 one
1418 two three"}
1419 );
1420 }
1421
1422 #[test]
1423 fn test_extract_editable_region_marker_format() {
1424 let text = indoc::indoc! {"
1425 some context
1426 <|marker_1|>
1427 one
1428 two three
1429 <|marker_2|>
1430 more context
1431 "};
1432 let parsed = multi_region::extract_editable_region_from_markers(text).unwrap();
1433 assert_eq!(parsed, "one\ntwo three");
1434 }
1435
1436 #[test]
1437 fn test_extract_editable_region_multi_markers() {
1438 let text = indoc::indoc! {"
1439 prefix
1440 <|marker_1|>
1441 aaa
1442 bbb
1443 <|marker_2|>
1444 ccc
1445 ddd
1446 <|marker_3|>
1447 suffix
1448 "};
1449 let parsed = multi_region::extract_editable_region_from_markers(text).unwrap();
1450 // Intermediate marker and its trailing \n are stripped
1451 assert_eq!(parsed, "aaa\nbbb\nccc\nddd");
1452 }
1453
1454 #[test]
1455 fn test_extract_last_codeblock_nested_bibtex() {
1456 let text = indoc::indoc! {r#"
1457 Looking at the edit history, I can see that a Citation section was just added.
1458
1459 `````
1460 ## Collaborations
1461 Our mission is to create a 4D generative model.
1462
1463 ## Citation
1464
1465 If you found Unique3D helpful, please cite our report:
1466 ```bibtex
1467 @misc{wu2024unique3d,
1468 title={Unique3D},
1469 }
1470 ```
1471 `````
1472 "#};
1473 let last_block = extract_last_codeblock(text).unwrap();
1474 assert_eq!(
1475 last_block,
1476 indoc::indoc! {r#"
1477 ## Collaborations
1478 Our mission is to create a 4D generative model.
1479
1480 ## Citation
1481
1482 If you found Unique3D helpful, please cite our report:
1483 ```bibtex
1484 @misc{wu2024unique3d,
1485 title={Unique3D},
1486 }
1487 ```
1488 "#}
1489 );
1490 }
1491
1492 #[test]
1493 fn test_extract_editable_region_no_markers() {
1494 let text = indoc::indoc! {"
1495 one
1496 two three"};
1497 let parsed = TeacherPrompt::extract_editable_region(text).unwrap();
1498 assert_eq!(
1499 parsed,
1500 indoc::indoc! {"
1501 one
1502 two three"}
1503 );
1504 }
1505
1506 #[test]
1507 fn test_parse_no_edits_response() {
1508 let response = indoc::indoc! {"
1509 The code is already complete. There is no clear next edit to make.
1510
1511 `````
1512 NO_EDITS
1513 `````
1514 "};
1515 let codeblock = extract_last_codeblock(response).unwrap();
1516 assert_eq!(codeblock.trim(), TeacherPrompt::NO_EDITS);
1517 }
1518
1519 #[test]
1520 fn test_extract_codeblock_no_valid_block() {
1521 // Text with no code blocks should return None
1522 let text = "Just some plain text without any code blocks";
1523 assert!(extract_last_codeblock(text).is_none());
1524
1525 // Unclosed code block should return None
1526 let text = indoc::indoc! {"
1527 ```
1528 unclosed block
1529 "};
1530 assert!(extract_last_codeblock(text).is_none());
1531
1532 // Analysis text with nested markdown but no proper outer block
1533 let text = indoc::indoc! {"
1534 # Analysis
1535 Looking at this:
1536 ```
1537 some code
1538 ```
1539 But then more analysis without wrapping block
1540 "};
1541 // This should find the inner block
1542 let result = extract_last_codeblock(text).unwrap();
1543 assert_eq!(result, "some code\n");
1544 }
1545
1546 #[test]
1547 fn test_extract_codeblock_no_trailing_newline() {
1548 // Text ending without trailing newline after closing fence
1549 let text = "`````\ncontent here\n`````";
1550 let result = extract_last_codeblock(text).unwrap();
1551 assert_eq!(result, "content here\n");
1552 }
1553
1554 #[test]
1555 fn test_parse_no_edits_response_with_trailing_backticks() {
1556 let response = "NO_EDITS```";
1557
1558 let parsed = TeacherPrompt::parse(
1559 &Example {
1560 spec: edit_prediction::example_spec::ExampleSpec {
1561 name: "test".to_string(),
1562 repository_url: "https://github.com/zed-industries/zed.git".to_string(),
1563 revision: "HEAD".to_string(),
1564 tags: Vec::new(),
1565 reasoning: None,
1566 uncommitted_diff: String::new(),
1567 recently_opened_files: Vec::new(),
1568 recently_viewed_files: Vec::new(),
1569 uncommitted_diff_contains_edit_history: false,
1570 cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
1571 cursor_position: "0:0".to_string(),
1572 edit_history: String::new(),
1573 expected_patches: Vec::new(),
1574 rejected_patch: None,
1575 telemetry: None,
1576 human_feedback: Vec::new(),
1577 rating: None,
1578 },
1579 prompt_inputs: None,
1580 prompt: None,
1581 predictions: Vec::new(),
1582 score: Vec::new(),
1583 qa: Vec::new(),
1584 zed_version: None,
1585 state: None,
1586 },
1587 response,
1588 )
1589 .unwrap();
1590
1591 assert!(parsed.0.is_empty());
1592 assert!(parsed.1.is_none());
1593 }
1594
1595 #[test]
1596 fn test_v0327_teacher_prompt_uses_resolved_ranges() {
1597 let excerpt = (0..80)
1598 .map(|index| format!("line{index:02}\n"))
1599 .collect::<String>();
1600 let cursor_offset = excerpt.find("line40").expect("cursor line exists");
1601 let prompt_inputs = zeta_prompt::Zeta2PromptInput {
1602 cursor_path: std::path::Path::new("src/main.rs").into(),
1603 cursor_excerpt: excerpt.clone().into(),
1604 cursor_offset_in_excerpt: cursor_offset,
1605 excerpt_start_row: None,
1606 events: Vec::new(),
1607 related_files: Some(Vec::new()),
1608 active_buffer_diagnostics: Vec::new(),
1609 excerpt_ranges: zeta_prompt::ExcerptRanges {
1610 editable_150: 0..32,
1611 editable_180: 0..32,
1612 editable_350: 0..32,
1613 editable_512: None,
1614 editable_150_context_350: 0..48,
1615 editable_180_context_350: 0..48,
1616 editable_350_context_150: 20..50,
1617 editable_350_context_512: None,
1618 editable_350_context_1024: None,
1619 context_4096: None,
1620 context_8192: Some(30..excerpt.len()),
1621 },
1622 syntax_ranges: None,
1623 in_open_source_repo: false,
1624 can_collect_data: false,
1625 repo_url: None,
1626 };
1627
1628 let (stored_editable_range, stored_context_range) = zeta_prompt::excerpt_range_for_format(
1629 ZetaFormat::V0327SingleFile,
1630 &prompt_inputs.excerpt_ranges,
1631 );
1632 assert!(stored_context_range.start > stored_editable_range.start);
1633
1634 let (editable_range, context_range) =
1635 resolved_excerpt_ranges_for_format(&prompt_inputs, ZetaFormat::V0327SingleFile);
1636 assert_eq!(context_range, 0..excerpt.len());
1637 assert!(editable_range.start < cursor_offset);
1638 assert!(editable_range.end > cursor_offset);
1639
1640 let prompt = TeacherPrompt::format_prompt(
1641 &Example {
1642 spec: edit_prediction::example_spec::ExampleSpec {
1643 name: "test".to_string(),
1644 repository_url: "https://github.com/zed-industries/zed.git".to_string(),
1645 revision: "HEAD".to_string(),
1646 tags: Vec::new(),
1647 reasoning: None,
1648 uncommitted_diff: String::new(),
1649 recently_opened_files: Vec::new(),
1650 recently_viewed_files: Vec::new(),
1651 uncommitted_diff_contains_edit_history: false,
1652 cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
1653 cursor_position: "0:0".to_string(),
1654 edit_history: String::new(),
1655 expected_patches: Vec::new(),
1656 rejected_patch: None,
1657 telemetry: None,
1658 human_feedback: Vec::new(),
1659 rating: None,
1660 },
1661 prompt_inputs: Some(prompt_inputs),
1662 prompt: None,
1663 predictions: Vec::new(),
1664 score: Vec::new(),
1665 qa: Vec::new(),
1666 zed_version: None,
1667 state: None,
1668 },
1669 editable_range,
1670 context_range,
1671 false,
1672 );
1673
1674 assert!(prompt.contains(TeacherPrompt::EDITABLE_REGION_START));
1675 assert!(prompt.contains(TeacherPrompt::USER_CURSOR_MARKER));
1676 assert!(prompt.contains("line40"));
1677 }
1678}
1679