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