1use anyhow::Result;
2use serde::{Deserialize, Serialize};
3use std::fmt::Write;
4use std::ops::Range;
5use std::path::Path;
6use std::sync::Arc;
7use strum::{EnumIter, IntoEnumIterator as _, IntoStaticStr};
8
9pub const CURSOR_MARKER: &str = "<|user_cursor|>";
10pub const MAX_PROMPT_TOKENS: usize = 4096;
11
12fn estimate_tokens(bytes: usize) -> usize {
13 bytes / 3
14}
15
16#[derive(Clone, Debug, Serialize, Deserialize)]
17pub struct ZetaPromptInput {
18 pub cursor_path: Arc<Path>,
19 pub cursor_excerpt: Arc<str>,
20 pub editable_range_in_excerpt: Range<usize>,
21 pub cursor_offset_in_excerpt: usize,
22 #[serde(default, skip_serializing_if = "Option::is_none")]
23 pub excerpt_start_row: Option<u32>,
24 pub events: Vec<Arc<Event>>,
25 pub related_files: Vec<RelatedFile>,
26}
27
28#[derive(
29 Default,
30 Clone,
31 Copy,
32 Debug,
33 PartialEq,
34 Eq,
35 Hash,
36 EnumIter,
37 IntoStaticStr,
38 Serialize,
39 Deserialize,
40)]
41#[allow(non_camel_case_types)]
42pub enum ZetaFormat {
43 V0112MiddleAtEnd,
44 V0113Ordered,
45 #[default]
46 V0114180EditableRegion,
47 V0120GitMergeMarkers,
48 V0131GitMergeMarkersPrefix,
49}
50
51impl std::fmt::Display for ZetaFormat {
52 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53 write!(f, "{}", <&'static str>::from(self))
54 }
55}
56
57impl ZetaFormat {
58 pub fn parse(format_name: &str) -> Result<Self> {
59 let mut results = ZetaFormat::iter().filter(|version| {
60 <&'static str>::from(version)
61 .to_lowercase()
62 .contains(&format_name.to_lowercase())
63 });
64 let Some(result) = results.next() else {
65 anyhow::bail!(
66 "`{format_name}` did not match any of:\n{}",
67 Self::options_as_string()
68 );
69 };
70 if results.next().is_some() {
71 anyhow::bail!(
72 "`{format_name}` matched more than one of:\n{}",
73 Self::options_as_string()
74 );
75 }
76 Ok(result)
77 }
78
79 pub fn options_as_string() -> String {
80 ZetaFormat::iter()
81 .map(|format| format!("- {}\n", <&'static str>::from(format)))
82 .collect::<Vec<_>>()
83 .concat()
84 }
85}
86
87#[derive(Clone, Debug, Serialize, Deserialize)]
88#[serde(tag = "event")]
89pub enum Event {
90 BufferChange {
91 path: Arc<Path>,
92 old_path: Arc<Path>,
93 diff: String,
94 predicted: bool,
95 in_open_source_repo: bool,
96 },
97}
98
99pub fn write_event(prompt: &mut String, event: &Event) {
100 fn write_path_as_unix_str(prompt: &mut String, path: &Path) {
101 for component in path.components() {
102 prompt.push('/');
103 write!(prompt, "{}", component.as_os_str().display()).ok();
104 }
105 }
106 match event {
107 Event::BufferChange {
108 path,
109 old_path,
110 diff,
111 predicted,
112 in_open_source_repo: _,
113 } => {
114 if *predicted {
115 prompt.push_str("// User accepted prediction:\n");
116 }
117 prompt.push_str("--- a");
118 write_path_as_unix_str(prompt, old_path.as_ref());
119 prompt.push_str("\n+++ b");
120 write_path_as_unix_str(prompt, path.as_ref());
121 prompt.push('\n');
122 prompt.push_str(diff);
123 }
124 }
125}
126
127#[derive(Clone, Debug, Serialize, Deserialize)]
128pub struct RelatedFile {
129 pub path: Arc<Path>,
130 pub max_row: u32,
131 pub excerpts: Vec<RelatedExcerpt>,
132}
133
134#[derive(Clone, Debug, Serialize, Deserialize)]
135pub struct RelatedExcerpt {
136 pub row_range: Range<u32>,
137 pub text: Arc<str>,
138}
139
140pub fn format_zeta_prompt(input: &ZetaPromptInput, format: ZetaFormat) -> String {
141 format_zeta_prompt_with_budget(input, format, MAX_PROMPT_TOKENS)
142}
143
144/// Post-processes model output for the given zeta format by stripping format-specific suffixes.
145pub fn clean_zeta2_model_output(output: &str, format: ZetaFormat) -> &str {
146 match format {
147 ZetaFormat::V0120GitMergeMarkers => output
148 .strip_suffix(v0120_git_merge_markers::END_MARKER)
149 .unwrap_or(output),
150 ZetaFormat::V0131GitMergeMarkersPrefix => output
151 .strip_suffix(v0131_git_merge_markers_prefix::END_MARKER)
152 .unwrap_or(output),
153 _ => output,
154 }
155}
156
157fn format_zeta_prompt_with_budget(
158 input: &ZetaPromptInput,
159 format: ZetaFormat,
160 max_tokens: usize,
161) -> String {
162 let mut cursor_section = String::new();
163 match format {
164 ZetaFormat::V0112MiddleAtEnd => {
165 v0112_middle_at_end::write_cursor_excerpt_section(&mut cursor_section, input);
166 }
167 ZetaFormat::V0113Ordered | ZetaFormat::V0114180EditableRegion => {
168 v0113_ordered::write_cursor_excerpt_section(&mut cursor_section, input)
169 }
170 ZetaFormat::V0120GitMergeMarkers => {
171 v0120_git_merge_markers::write_cursor_excerpt_section(&mut cursor_section, input)
172 }
173 ZetaFormat::V0131GitMergeMarkersPrefix => {
174 v0131_git_merge_markers_prefix::write_cursor_excerpt_section(&mut cursor_section, input)
175 }
176 }
177
178 let cursor_tokens = estimate_tokens(cursor_section.len());
179 let budget_after_cursor = max_tokens.saturating_sub(cursor_tokens);
180
181 let edit_history_section =
182 format_edit_history_within_budget(&input.events, budget_after_cursor);
183 let edit_history_tokens = estimate_tokens(edit_history_section.len());
184 let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens);
185
186 let related_files_section =
187 format_related_files_within_budget(&input.related_files, budget_after_edit_history);
188
189 let mut prompt = String::new();
190 prompt.push_str(&related_files_section);
191 prompt.push_str(&edit_history_section);
192 prompt.push_str(&cursor_section);
193 prompt
194}
195
196fn format_edit_history_within_budget(events: &[Arc<Event>], max_tokens: usize) -> String {
197 let header = "<|file_sep|>edit history\n";
198 let header_tokens = estimate_tokens(header.len());
199 if header_tokens >= max_tokens {
200 return String::new();
201 }
202
203 let mut event_strings: Vec<String> = Vec::new();
204 let mut total_tokens = header_tokens;
205
206 for event in events.iter().rev() {
207 let mut event_str = String::new();
208 write_event(&mut event_str, event);
209 let event_tokens = estimate_tokens(event_str.len());
210
211 if total_tokens + event_tokens > max_tokens {
212 break;
213 }
214 total_tokens += event_tokens;
215 event_strings.push(event_str);
216 }
217
218 if event_strings.is_empty() {
219 return String::new();
220 }
221
222 let mut result = String::from(header);
223 for event_str in event_strings.iter().rev() {
224 result.push_str(&event_str);
225 }
226 result
227}
228
229fn format_related_files_within_budget(related_files: &[RelatedFile], max_tokens: usize) -> String {
230 let mut result = String::new();
231 let mut total_tokens = 0;
232
233 for file in related_files {
234 let path_str = file.path.to_string_lossy();
235 let header_len = "<|file_sep|>".len() + path_str.len() + 1;
236 let header_tokens = estimate_tokens(header_len);
237
238 if total_tokens + header_tokens > max_tokens {
239 break;
240 }
241
242 let mut file_tokens = header_tokens;
243 let mut excerpts_to_include = 0;
244
245 for excerpt in &file.excerpts {
246 let needs_newline = !excerpt.text.ends_with('\n');
247 let needs_ellipsis = excerpt.row_range.end < file.max_row;
248 let excerpt_len = excerpt.text.len()
249 + if needs_newline { "\n".len() } else { "".len() }
250 + if needs_ellipsis {
251 "...\n".len()
252 } else {
253 "".len()
254 };
255
256 let excerpt_tokens = estimate_tokens(excerpt_len);
257 if total_tokens + file_tokens + excerpt_tokens > max_tokens {
258 break;
259 }
260 file_tokens += excerpt_tokens;
261 excerpts_to_include += 1;
262 }
263
264 if excerpts_to_include > 0 {
265 total_tokens += file_tokens;
266 write!(result, "<|file_sep|>{}\n", path_str).ok();
267 for excerpt in file.excerpts.iter().take(excerpts_to_include) {
268 result.push_str(&excerpt.text);
269 if !result.ends_with('\n') {
270 result.push('\n');
271 }
272 if excerpt.row_range.end < file.max_row {
273 result.push_str("...\n");
274 }
275 }
276 }
277 }
278
279 result
280}
281
282pub fn write_related_files(
283 prompt: &mut String,
284 related_files: &[RelatedFile],
285) -> Vec<Range<usize>> {
286 let mut ranges = Vec::new();
287 for file in related_files {
288 let start = prompt.len();
289 let path_str = file.path.to_string_lossy();
290 write!(prompt, "<|file_sep|>{}\n", path_str).ok();
291 for excerpt in &file.excerpts {
292 prompt.push_str(&excerpt.text);
293 if !prompt.ends_with('\n') {
294 prompt.push('\n');
295 }
296 if excerpt.row_range.end < file.max_row {
297 prompt.push_str("...\n");
298 }
299 }
300 let end = prompt.len();
301 ranges.push(start..end);
302 }
303 ranges
304}
305
306mod v0112_middle_at_end {
307 use super::*;
308
309 pub fn write_cursor_excerpt_section(prompt: &mut String, input: &ZetaPromptInput) {
310 let path_str = input.cursor_path.to_string_lossy();
311 write!(prompt, "<|file_sep|>{}\n", path_str).ok();
312
313 prompt.push_str("<|fim_prefix|>\n");
314 prompt.push_str(&input.cursor_excerpt[..input.editable_range_in_excerpt.start]);
315
316 prompt.push_str("<|fim_suffix|>\n");
317 prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
318 if !prompt.ends_with('\n') {
319 prompt.push('\n');
320 }
321
322 prompt.push_str("<|fim_middle|>current\n");
323 prompt.push_str(
324 &input.cursor_excerpt
325 [input.editable_range_in_excerpt.start..input.cursor_offset_in_excerpt],
326 );
327 prompt.push_str(CURSOR_MARKER);
328 prompt.push_str(
329 &input.cursor_excerpt
330 [input.cursor_offset_in_excerpt..input.editable_range_in_excerpt.end],
331 );
332 if !prompt.ends_with('\n') {
333 prompt.push('\n');
334 }
335
336 prompt.push_str("<|fim_middle|>updated\n");
337 }
338}
339
340mod v0113_ordered {
341 use super::*;
342
343 pub fn write_cursor_excerpt_section(prompt: &mut String, input: &ZetaPromptInput) {
344 let path_str = input.cursor_path.to_string_lossy();
345 write!(prompt, "<|file_sep|>{}\n", path_str).ok();
346
347 prompt.push_str("<|fim_prefix|>\n");
348 prompt.push_str(&input.cursor_excerpt[..input.editable_range_in_excerpt.start]);
349 if !prompt.ends_with('\n') {
350 prompt.push('\n');
351 }
352
353 prompt.push_str("<|fim_middle|>current\n");
354 prompt.push_str(
355 &input.cursor_excerpt
356 [input.editable_range_in_excerpt.start..input.cursor_offset_in_excerpt],
357 );
358 prompt.push_str(CURSOR_MARKER);
359 prompt.push_str(
360 &input.cursor_excerpt
361 [input.cursor_offset_in_excerpt..input.editable_range_in_excerpt.end],
362 );
363 if !prompt.ends_with('\n') {
364 prompt.push('\n');
365 }
366
367 prompt.push_str("<|fim_suffix|>\n");
368 prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
369 if !prompt.ends_with('\n') {
370 prompt.push('\n');
371 }
372
373 prompt.push_str("<|fim_middle|>updated\n");
374 }
375}
376
377pub mod v0120_git_merge_markers {
378 //! A prompt that uses git-style merge conflict markers to represent the editable region.
379 //!
380 //! Example prompt:
381 //!
382 //! <|file_sep|>path/to/target_file.py
383 //! <|fim_prefix|>
384 //! code before editable region
385 //! <|fim_suffix|>
386 //! code after editable region
387 //! <|fim_middle|>
388 //! <<<<<<< CURRENT
389 //! code that
390 //! needs to<|user_cursor|>
391 //! be rewritten
392 //! =======
393 //!
394 //! Expected output (should be generated by the model):
395 //!
396 //! updated
397 //! code with
398 //! changes applied
399 //! >>>>>>> UPDATED
400
401 use super::*;
402
403 pub const START_MARKER: &str = "<<<<<<< CURRENT\n";
404 pub const SEPARATOR: &str = "=======\n";
405 pub const END_MARKER: &str = ">>>>>>> UPDATED\n";
406
407 pub fn write_cursor_excerpt_section(prompt: &mut String, input: &ZetaPromptInput) {
408 let path_str = input.cursor_path.to_string_lossy();
409 write!(prompt, "<|file_sep|>{}\n", path_str).ok();
410
411 prompt.push_str("<|fim_prefix|>");
412 prompt.push_str(&input.cursor_excerpt[..input.editable_range_in_excerpt.start]);
413
414 prompt.push_str("<|fim_suffix|>");
415 prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
416 if !prompt.ends_with('\n') {
417 prompt.push('\n');
418 }
419
420 prompt.push_str("<|fim_middle|>");
421 prompt.push_str(START_MARKER);
422 prompt.push_str(
423 &input.cursor_excerpt
424 [input.editable_range_in_excerpt.start..input.cursor_offset_in_excerpt],
425 );
426 prompt.push_str(CURSOR_MARKER);
427 prompt.push_str(
428 &input.cursor_excerpt
429 [input.cursor_offset_in_excerpt..input.editable_range_in_excerpt.end],
430 );
431 if !prompt.ends_with('\n') {
432 prompt.push('\n');
433 }
434 prompt.push_str(SEPARATOR);
435 }
436}
437
438pub mod v0131_git_merge_markers_prefix {
439 //! A prompt that uses git-style merge conflict markers to represent the editable region.
440 //!
441 //! Example prompt:
442 //!
443 //! <|file_sep|>path/to/target_file.py
444 //! <|fim_prefix|>
445 //! code before editable region
446 //! <<<<<<< CURRENT
447 //! code that
448 //! needs to<|user_cursor|>
449 //! be rewritten
450 //! =======
451 //! <|fim_suffix|>
452 //! code after editable region
453 //! <|fim_middle|>
454 //!
455 //! Expected output (should be generated by the model):
456 //!
457 //! updated
458 //! code with
459 //! changes applied
460 //! >>>>>>> UPDATED
461
462 use super::*;
463
464 pub const START_MARKER: &str = "<<<<<<< CURRENT\n";
465 pub const SEPARATOR: &str = "=======\n";
466 pub const END_MARKER: &str = ">>>>>>> UPDATED\n";
467
468 pub fn write_cursor_excerpt_section(prompt: &mut String, input: &ZetaPromptInput) {
469 let path_str = input.cursor_path.to_string_lossy();
470 write!(prompt, "<|file_sep|>{}\n", path_str).ok();
471
472 prompt.push_str("<|fim_prefix|>");
473 prompt.push_str(&input.cursor_excerpt[..input.editable_range_in_excerpt.start]);
474 prompt.push_str(START_MARKER);
475 prompt.push_str(
476 &input.cursor_excerpt
477 [input.editable_range_in_excerpt.start..input.cursor_offset_in_excerpt],
478 );
479 prompt.push_str(CURSOR_MARKER);
480 prompt.push_str(
481 &input.cursor_excerpt
482 [input.cursor_offset_in_excerpt..input.editable_range_in_excerpt.end],
483 );
484 if !prompt.ends_with('\n') {
485 prompt.push('\n');
486 }
487 prompt.push_str(SEPARATOR);
488
489 prompt.push_str("<|fim_suffix|>");
490 prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
491 if !prompt.ends_with('\n') {
492 prompt.push('\n');
493 }
494
495 prompt.push_str("<|fim_middle|>");
496 }
497}
498
499/// The zeta1 prompt format
500pub mod zeta1 {
501 pub const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
502 pub const START_OF_FILE_MARKER: &str = "<|start_of_file|>";
503 pub const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>";
504 pub const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>";
505
506 const INSTRUCTION_HEADER: &str = concat!(
507 "### Instruction:\n",
508 "You are a code completion assistant and your task is to analyze user edits and then rewrite an ",
509 "excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking ",
510 "into account the cursor location.\n\n",
511 "### User Edits:\n\n"
512 );
513 const EXCERPT_HEADER: &str = "\n\n### User Excerpt:\n\n";
514 const RESPONSE_HEADER: &str = "\n\n### Response:\n";
515
516 /// Formats a complete zeta1 prompt from the input events and excerpt.
517 pub fn format_zeta1_prompt(input_events: &str, input_excerpt: &str) -> String {
518 let mut prompt = String::with_capacity(
519 INSTRUCTION_HEADER.len()
520 + input_events.len()
521 + EXCERPT_HEADER.len()
522 + input_excerpt.len()
523 + RESPONSE_HEADER.len(),
524 );
525 prompt.push_str(INSTRUCTION_HEADER);
526 prompt.push_str(input_events);
527 prompt.push_str(EXCERPT_HEADER);
528 prompt.push_str(input_excerpt);
529 prompt.push_str(RESPONSE_HEADER);
530 prompt
531 }
532}
533
534#[cfg(test)]
535mod tests {
536 use super::*;
537 use indoc::indoc;
538
539 fn make_input(
540 cursor_excerpt: &str,
541 editable_range: Range<usize>,
542 cursor_offset: usize,
543 events: Vec<Event>,
544 related_files: Vec<RelatedFile>,
545 ) -> ZetaPromptInput {
546 ZetaPromptInput {
547 cursor_path: Path::new("test.rs").into(),
548 cursor_excerpt: cursor_excerpt.into(),
549 editable_range_in_excerpt: editable_range,
550 cursor_offset_in_excerpt: cursor_offset,
551 excerpt_start_row: None,
552 events: events.into_iter().map(Arc::new).collect(),
553 related_files,
554 }
555 }
556
557 fn make_event(path: &str, diff: &str) -> Event {
558 Event::BufferChange {
559 path: Path::new(path).into(),
560 old_path: Path::new(path).into(),
561 diff: diff.to_string(),
562 predicted: false,
563 in_open_source_repo: false,
564 }
565 }
566
567 fn make_related_file(path: &str, content: &str) -> RelatedFile {
568 RelatedFile {
569 path: Path::new(path).into(),
570 max_row: content.lines().count() as u32,
571 excerpts: vec![RelatedExcerpt {
572 row_range: 0..content.lines().count() as u32,
573 text: content.into(),
574 }],
575 }
576 }
577
578 fn format_with_budget(input: &ZetaPromptInput, max_tokens: usize) -> String {
579 format_zeta_prompt_with_budget(input, ZetaFormat::V0114180EditableRegion, max_tokens)
580 }
581
582 #[test]
583 fn test_no_truncation_when_within_budget() {
584 let input = make_input(
585 "prefix\neditable\nsuffix",
586 7..15,
587 10,
588 vec![make_event("a.rs", "-old\n+new\n")],
589 vec![make_related_file("related.rs", "fn helper() {}\n")],
590 );
591
592 assert_eq!(
593 format_with_budget(&input, 10000),
594 indoc! {r#"
595 <|file_sep|>related.rs
596 fn helper() {}
597 <|file_sep|>edit history
598 --- a/a.rs
599 +++ b/a.rs
600 -old
601 +new
602 <|file_sep|>test.rs
603 <|fim_prefix|>
604 prefix
605 <|fim_middle|>current
606 edi<|user_cursor|>table
607 <|fim_suffix|>
608
609 suffix
610 <|fim_middle|>updated
611 "#}
612 );
613 }
614
615 #[test]
616 fn test_truncation_drops_edit_history_when_budget_tight() {
617 let input = make_input(
618 "code",
619 0..4,
620 2,
621 vec![make_event("a.rs", "-x\n+y\n")],
622 vec![
623 make_related_file("r1.rs", "a\n"),
624 make_related_file("r2.rs", "b\n"),
625 ],
626 );
627
628 assert_eq!(
629 format_with_budget(&input, 10000),
630 indoc! {r#"
631 <|file_sep|>r1.rs
632 a
633 <|file_sep|>r2.rs
634 b
635 <|file_sep|>edit history
636 --- a/a.rs
637 +++ b/a.rs
638 -x
639 +y
640 <|file_sep|>test.rs
641 <|fim_prefix|>
642 <|fim_middle|>current
643 co<|user_cursor|>de
644 <|fim_suffix|>
645 <|fim_middle|>updated
646 "#}
647 );
648
649 assert_eq!(
650 format_with_budget(&input, 50),
651 indoc! {r#"
652 <|file_sep|>r1.rs
653 a
654 <|file_sep|>r2.rs
655 b
656 <|file_sep|>test.rs
657 <|fim_prefix|>
658 <|fim_middle|>current
659 co<|user_cursor|>de
660 <|fim_suffix|>
661 <|fim_middle|>updated
662 "#}
663 );
664 }
665
666 #[test]
667 fn test_truncation_includes_partial_excerpts() {
668 let input = make_input(
669 "x",
670 0..1,
671 0,
672 vec![],
673 vec![RelatedFile {
674 path: Path::new("big.rs").into(),
675 max_row: 30,
676 excerpts: vec![
677 RelatedExcerpt {
678 row_range: 0..10,
679 text: "first excerpt\n".into(),
680 },
681 RelatedExcerpt {
682 row_range: 10..20,
683 text: "second excerpt\n".into(),
684 },
685 RelatedExcerpt {
686 row_range: 20..30,
687 text: "third excerpt\n".into(),
688 },
689 ],
690 }],
691 );
692
693 assert_eq!(
694 format_with_budget(&input, 10000),
695 indoc! {r#"
696 <|file_sep|>big.rs
697 first excerpt
698 ...
699 second excerpt
700 ...
701 third excerpt
702 <|file_sep|>test.rs
703 <|fim_prefix|>
704 <|fim_middle|>current
705 <|user_cursor|>x
706 <|fim_suffix|>
707 <|fim_middle|>updated
708 "#}
709 );
710
711 assert_eq!(
712 format_with_budget(&input, 50),
713 indoc! {r#"
714 <|file_sep|>big.rs
715 first excerpt
716 ...
717 <|file_sep|>test.rs
718 <|fim_prefix|>
719 <|fim_middle|>current
720 <|user_cursor|>x
721 <|fim_suffix|>
722 <|fim_middle|>updated
723 "#}
724 );
725 }
726
727 #[test]
728 fn test_truncation_drops_older_events_first() {
729 let input = make_input(
730 "x",
731 0..1,
732 0,
733 vec![make_event("old.rs", "-1\n"), make_event("new.rs", "-2\n")],
734 vec![],
735 );
736
737 assert_eq!(
738 format_with_budget(&input, 10000),
739 indoc! {r#"
740 <|file_sep|>edit history
741 --- a/old.rs
742 +++ b/old.rs
743 -1
744 --- a/new.rs
745 +++ b/new.rs
746 -2
747 <|file_sep|>test.rs
748 <|fim_prefix|>
749 <|fim_middle|>current
750 <|user_cursor|>x
751 <|fim_suffix|>
752 <|fim_middle|>updated
753 "#}
754 );
755
756 assert_eq!(
757 format_with_budget(&input, 55),
758 indoc! {r#"
759 <|file_sep|>edit history
760 --- a/new.rs
761 +++ b/new.rs
762 -2
763 <|file_sep|>test.rs
764 <|fim_prefix|>
765 <|fim_middle|>current
766 <|user_cursor|>x
767 <|fim_suffix|>
768 <|fim_middle|>updated
769 "#}
770 );
771 }
772
773 #[test]
774 fn test_cursor_excerpt_always_included_with_minimal_budget() {
775 let input = make_input(
776 "fn main() {}",
777 0..12,
778 3,
779 vec![make_event("a.rs", "-old\n+new\n")],
780 vec![make_related_file("related.rs", "helper\n")],
781 );
782
783 assert_eq!(
784 format_with_budget(&input, 30),
785 indoc! {r#"
786 <|file_sep|>test.rs
787 <|fim_prefix|>
788 <|fim_middle|>current
789 fn <|user_cursor|>main() {}
790 <|fim_suffix|>
791 <|fim_middle|>updated
792 "#}
793 );
794 }
795}