tenant.openagents/omega

1pub mod excerpt_ranges;
2pub mod hashed_regions;
3pub mod multi_region;
4pub mod udiff;
5
6use anyhow::{Result, anyhow};
7use serde::{Deserialize, Serialize};
8use std::fmt::Write;
9use std::ops::Range;
10use std::path::Path;
11use std::sync::Arc;
12use strum::{EnumIter, IntoEnumIterator as _, IntoStaticStr};
13
14pub use crate::excerpt_ranges::{
15    ExcerptRanges, compute_editable_and_context_ranges, compute_legacy_excerpt_ranges,
16};
17
18pub const CURSOR_MARKER: &str = "<|user_cursor|>";
19
20/// Use up to this amount of the editable region for prefill.
21/// Larger values may result in more robust generation, but
22/// this region becomes non-editable.
23pub const PREFILL_RATIO: f64 = 0.1; // 10%
24
25fn estimate_tokens(bytes: usize) -> usize {
26    bytes / 3
27}
28
29/// Leave some slack to avoid overflow.
30fn apply_prompt_budget_margin(max_tokens: usize) -> usize {
31    (max_tokens as f64 * 0.9).floor() as usize
32}
33
34/// Ensure text fits into the tokens budget; trim by line boundaries if needed.
35pub fn clamp_text_to_token_count(text: &str, max_tokens: usize) -> &str {
36    if estimate_tokens(text.len()) <= max_tokens {
37        return text;
38    }
39
40    let mut end_byte_offset = 0;
41
42    for line in text.split_inclusive('\n') {
43        if estimate_tokens(line.len() + end_byte_offset) > max_tokens {
44            break;
45        }
46
47        end_byte_offset += line.len();
48    }
49
50    &text[..end_byte_offset]
51}
52
53#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
54pub struct Zeta2PromptInput {
55    pub cursor_path: Arc<Path>,
56    pub cursor_excerpt: Arc<str>,
57    pub cursor_offset_in_excerpt: usize,
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub excerpt_start_row: Option<u32>,
60    pub events: Vec<Arc<Event>>,
61    #[serde(default)]
62    pub related_files: Option<Vec<RelatedFile>>,
63    #[serde(default, skip_serializing_if = "Vec::is_empty")]
64    pub active_buffer_diagnostics: Vec<ActiveBufferDiagnostic>,
65    /// These ranges let the server select model-appropriate subsets.
66    pub excerpt_ranges: ExcerptRanges,
67    /// Byte offset ranges within `cursor_excerpt` for all syntax nodes that
68    /// contain `cursor_offset_in_excerpt`, ordered from innermost to outermost.
69    /// When present, the server uses these to compute editable/context ranges
70    /// instead of `excerpt_ranges`.
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub syntax_ranges: Option<Vec<Range<usize>>>,
73    #[serde(default)]
74    pub in_open_source_repo: bool,
75    #[serde(default)]
76    pub can_collect_data: bool,
77    #[serde(default, skip_serializing_if = "Option::is_none")]
78    pub repo_url: Option<String>,
79}
80
81#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
82pub struct FilePosition {
83    pub row: u32,
84    pub column: u32,
85}
86
87#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
88pub struct Zeta3PromptInput {
89    pub cursor_path: Arc<Path>,
90    pub cursor_position: FilePosition,
91    pub events: Vec<Arc<Event>>,
92    pub editable_context: Vec<RelatedFile>,
93    #[serde(default, skip_serializing_if = "Vec::is_empty")]
94    pub syntax_ranges: Vec<Range<usize>>,
95    #[serde(default, skip_serializing_if = "Vec::is_empty")]
96    pub active_buffer_diagnostics: Vec<ActiveBufferDiagnostic>,
97    #[serde(default)]
98    pub in_open_source_repo: bool,
99    #[serde(default)]
100    pub can_collect_data: bool,
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub repo_url: Option<String>,
103}
104
105#[derive(
106    Default,
107    Clone,
108    Copy,
109    Debug,
110    PartialEq,
111    Eq,
112    Hash,
113    EnumIter,
114    IntoStaticStr,
115    Serialize,
116    Deserialize,
117)]
118#[allow(non_camel_case_types)]
119pub enum ZetaFormat {
120    V0112MiddleAtEnd,
121    V0113Ordered,
122    V0114180EditableRegion,
123    V0120GitMergeMarkers,
124    #[default]
125    V0131GitMergeMarkersPrefix,
126    V0211Prefill,
127    #[serde(alias = "Zeta2")]
128    V0211SeedCoder,
129    V0331SeedCoderModelPy,
130    v0226Hashline,
131    V0304VariableEdit,
132    V0304SeedNoEdits,
133    /// Multi-block marker spans with NO_EDITS sentinel.
134    V0306SeedMultiRegions,
135    /// Byte-exact marker spans; all intermediate markers emitted; repeated marker means no-edit.
136    V0316SeedMultiRegions,
137    /// V0316, but marker numbers are relative to the cursor block (e.g. -1, -0, +1).
138    V0317SeedMultiRegions,
139    /// V0316 with larger block sizes.
140    #[serde(alias = "Zeta2.1")]
141    V0318SeedMultiRegions,
142    /// V0318-style markers over the full available current file excerpt with no related files.
143    V0327SingleFile,
144    /// V0318-style prompt with buffer diagnostics
145    V0420Diagnostics,
146    /// V0318-style multi-region format using Qwen FIM tokens and PSM ordering.
147    V0608QwenMultiRegions,
148
149    /// V0318-style marker-span output, but with content-hashed marker tags over rendered
150    /// related-file context so the model can target jump edits. There is no cursor-centered
151    /// editable region for this format.
152    V0615HashRegions,
153}
154
155impl std::fmt::Display for ZetaFormat {
156    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157        write!(f, "{}", <&'static str>::from(self))
158    }
159}
160
161impl ZetaFormat {
162    pub fn parse(format_name: &str) -> Result<Self> {
163        let lower = format_name.to_lowercase();
164
165        // Exact case-insensitive match takes priority, bypassing ambiguity checks.
166        for variant in ZetaFormat::iter() {
167            if <&'static str>::from(&variant).to_lowercase() == lower {
168                return Ok(variant);
169            }
170        }
171
172        let mut results = ZetaFormat::iter().filter(|version| {
173            <&'static str>::from(version)
174                .to_lowercase()
175                .contains(&lower)
176        });
177        let Some(result) = results.next() else {
178            anyhow::bail!(
179                "`{format_name}` did not match any of:\n{}",
180                Self::options_as_string()
181            );
182        };
183        if results.next().is_some() {
184            anyhow::bail!(
185                "`{format_name}` matched more than one of:\n{}",
186                Self::options_as_string()
187            );
188        }
189        Ok(result)
190    }
191
192    pub fn options_as_string() -> String {
193        ZetaFormat::iter()
194            .map(|format| format!("- {}\n", <&'static str>::from(format)))
195            .collect::<Vec<_>>()
196            .concat()
197    }
198}
199
200fn empty_range() -> Range<usize> {
201    0..0
202}
203
204#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
205#[serde(tag = "event")]
206pub enum Event {
207    BufferChange {
208        path: Arc<Path>,
209        old_path: Arc<Path>,
210        diff: String,
211        #[serde(default = "empty_range")]
212        old_range: Range<usize>,
213        #[serde(default = "empty_range")]
214        new_range: Range<usize>,
215        predicted: bool,
216        in_open_source_repo: bool,
217    },
218}
219
220impl Event {
221    pub fn in_open_source_repo(&self) -> bool {
222        match self {
223            Event::BufferChange {
224                in_open_source_repo,
225                ..
226            } => *in_open_source_repo,
227        }
228    }
229}
230
231pub fn write_event(prompt: &mut String, event: &Event) {
232    fn write_path_as_unix_str(prompt: &mut String, path: &Path) {
233        for component in path.components() {
234            prompt.push('/');
235            write!(prompt, "{}", component.as_os_str().display()).ok();
236        }
237    }
238    match event {
239        Event::BufferChange {
240            path,
241            old_path,
242            diff,
243            predicted,
244            ..
245        } => {
246            if *predicted {
247                prompt.push_str("// User accepted prediction:\n");
248            }
249            prompt.push_str("--- a");
250            write_path_as_unix_str(prompt, old_path.as_ref());
251            prompt.push_str("\n+++ b");
252            write_path_as_unix_str(prompt, path.as_ref());
253            prompt.push('\n');
254            prompt.push_str(diff);
255        }
256    }
257}
258
259#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
260pub struct ActiveBufferDiagnostic {
261    pub severity: Option<i32>,
262    pub message: String,
263    pub snippet: String,
264    pub snippet_buffer_row_range: Range<u32>,
265    pub diagnostic_range_in_snippet: Range<usize>,
266}
267
268#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
269pub struct RelatedFile {
270    pub path: Arc<Path>,
271    pub max_row: u32,
272    pub excerpts: Vec<RelatedExcerpt>,
273    #[serde(default)]
274    pub in_open_source_repo: bool,
275}
276
277#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
278pub struct RelatedExcerpt {
279    pub row_range: Range<u32>,
280    pub text: Arc<str>,
281    #[serde(default)]
282    pub order: usize,
283    #[serde(default)]
284    pub context_source: ContextSource,
285}
286
287#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
288#[serde(rename_all = "snake_case")]
289pub enum ContextSource {
290    #[default]
291    Lsp,
292    CursorExcerpt,
293    CurrentFile,
294    EditHistory,
295    EditHistoryFile,
296    GitLog,
297    Bm25,
298    OracleFile,
299    OracleSnippet,
300}
301
302pub fn prompt_input_contains_special_tokens(input: &Zeta2PromptInput, format: ZetaFormat) -> bool {
303    special_tokens_for_format(format).iter().any(|token| {
304        if let Some(line_token) = token.strip_suffix('\n') {
305            input.cursor_excerpt.lines().any(|line| line == line_token)
306        } else {
307            input.cursor_excerpt.contains(token)
308        }
309    })
310}
311
312pub fn format_zeta_prompt(input: &Zeta2PromptInput, format: ZetaFormat) -> Option<String> {
313    format_prompt_with_budget_for_format(input, format, max_prompt_tokens_for_format(format))
314}
315
316pub fn format_zeta3_prompt(input: &Zeta3PromptInput, format: ZetaFormat) -> Option<String> {
317    match format {
318        ZetaFormat::V0318SeedMultiRegions => {}
319        _ => return None,
320    }
321
322    let (current_excerpt, cursor_offset_in_excerpt) = zeta3_current_file_excerpt(input)?;
323    let (context, editable_range, context_range, cursor_offset) = resolve_zeta3_cursor_region(
324        current_excerpt.text.as_ref(),
325        cursor_offset_in_excerpt,
326        &input.syntax_ranges,
327        format,
328    );
329    let relative_row_range =
330        offset_range_to_row_range(current_excerpt.text.as_ref(), context_range);
331    let cursor_row_range = current_excerpt.row_range.start + relative_row_range.start
332        ..current_excerpt.row_range.start + relative_row_range.end;
333    let related_files = filter_redundant_excerpts(
334        zeta3_related_files(input, current_excerpt),
335        input.cursor_path.as_ref(),
336        cursor_row_range,
337    );
338
339    format_resolved_prompt_with_budget(
340        format,
341        input.cursor_path.as_ref(),
342        context,
343        &editable_range,
344        cursor_offset,
345        &input.events,
346        &related_files,
347        &input.active_buffer_diagnostics,
348        Some(input.cursor_position.row),
349        max_prompt_tokens_for_format(format),
350    )
351}
352
353fn max_prompt_tokens_for_format(format: ZetaFormat) -> usize {
354    match format {
355        ZetaFormat::V0112MiddleAtEnd
356        | ZetaFormat::V0113Ordered
357        | ZetaFormat::V0114180EditableRegion
358        | ZetaFormat::V0120GitMergeMarkers
359        | ZetaFormat::V0131GitMergeMarkersPrefix
360        | ZetaFormat::V0211Prefill
361        | ZetaFormat::V0211SeedCoder
362        | ZetaFormat::v0226Hashline
363        | ZetaFormat::V0304VariableEdit
364        | ZetaFormat::V0304SeedNoEdits
365        | ZetaFormat::V0306SeedMultiRegions
366        | ZetaFormat::V0316SeedMultiRegions
367        | ZetaFormat::V0317SeedMultiRegions
368        | ZetaFormat::V0331SeedCoderModelPy
369        | ZetaFormat::V0318SeedMultiRegions
370        | ZetaFormat::V0608QwenMultiRegions => 4096,
371        ZetaFormat::V0615HashRegions => 8000,
372        ZetaFormat::V0420Diagnostics => 8192,
373        ZetaFormat::V0327SingleFile => 16384,
374    }
375}
376
377fn zeta3_current_file_excerpt(input: &Zeta3PromptInput) -> Option<(&RelatedExcerpt, usize)> {
378    input
379        .editable_context
380        .iter()
381        .filter(|file| file.path == input.cursor_path)
382        .flat_map(|file| file.excerpts.iter())
383        .find_map(|excerpt| {
384            if excerpt.context_source != ContextSource::CurrentFile {
385                return None;
386            }
387            Some((
388                excerpt,
389                offset_for_position_in_excerpt(excerpt, input.cursor_position)?,
390            ))
391        })
392}
393
394fn offset_for_position_in_excerpt(
395    excerpt: &RelatedExcerpt,
396    position: FilePosition,
397) -> Option<usize> {
398    if position.row < excerpt.row_range.start {
399        return None;
400    }
401
402    let relative_row = (position.row - excerpt.row_range.start) as usize;
403    let text = excerpt.text.as_ref();
404    let mut row_start = 0;
405
406    for row in 0..=relative_row {
407        if row == relative_row {
408            let row_end = text[row_start..]
409                .find('\n')
410                .map_or(text.len(), |offset| row_start + offset);
411            let row_text = &text[row_start..row_end];
412            let column =
413                row_text.floor_char_boundary((position.column as usize).min(row_text.len()));
414            return Some(row_start + column);
415        }
416
417        row_start += text[row_start..].find('\n')? + 1;
418    }
419
420    None
421}
422
423fn zeta3_related_files(
424    input: &Zeta3PromptInput,
425    current_excerpt: &RelatedExcerpt,
426) -> Vec<RelatedFile> {
427    input
428        .editable_context
429        .iter()
430        .filter_map(|file| {
431            let mut file = file.clone();
432            if file.path == input.cursor_path {
433                file.excerpts.retain(|excerpt| excerpt != current_excerpt);
434            }
435            (!file.excerpts.is_empty()).then_some(file)
436        })
437        .collect()
438}
439
440fn resolve_zeta3_cursor_region<'a>(
441    cursor_excerpt: &'a str,
442    cursor_offset: usize,
443    syntax_ranges: &[Range<usize>],
444    format: ZetaFormat,
445) -> (&'a str, Range<usize>, Range<usize>, usize) {
446    let (editable_tokens, context_tokens) = token_limits_for_format(format);
447    let (editable_range, context_range) = compute_editable_and_context_ranges(
448        cursor_excerpt,
449        cursor_offset,
450        syntax_ranges,
451        editable_tokens,
452        context_tokens,
453    );
454
455    adjust_cursor_region(cursor_excerpt, cursor_offset, editable_range, context_range)
456}
457
458pub fn special_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] {
459    match format {
460        ZetaFormat::V0112MiddleAtEnd => v0112_middle_at_end::special_tokens(),
461        ZetaFormat::V0113Ordered => v0113_ordered::special_tokens(),
462        ZetaFormat::V0114180EditableRegion => v0114180_editable_region::special_tokens(),
463        ZetaFormat::V0120GitMergeMarkers => v0120_git_merge_markers::special_tokens(),
464        ZetaFormat::V0131GitMergeMarkersPrefix => v0131_git_merge_markers_prefix::special_tokens(),
465        ZetaFormat::V0211Prefill => v0211_prefill::special_tokens(),
466        ZetaFormat::V0211SeedCoder | ZetaFormat::V0331SeedCoderModelPy => {
467            seed_coder::special_tokens()
468        }
469        ZetaFormat::v0226Hashline => hashline::special_tokens(),
470        ZetaFormat::V0304VariableEdit => v0304_variable_edit::special_tokens(),
471        ZetaFormat::V0304SeedNoEdits => seed_coder::special_tokens(),
472        ZetaFormat::V0316SeedMultiRegions => {
473            static TOKENS: &[&str] = &[
474                seed_coder::FIM_SUFFIX,
475                seed_coder::FIM_PREFIX,
476                seed_coder::FIM_MIDDLE,
477                seed_coder::FILE_MARKER,
478                multi_region::V0316_END_MARKER,
479                CURSOR_MARKER,
480                multi_region::MARKER_TAG_PREFIX,
481            ];
482            TOKENS
483        }
484        ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => {
485            static TOKENS: &[&str] = &[
486                seed_coder::FIM_SUFFIX,
487                seed_coder::FIM_PREFIX,
488                seed_coder::FIM_MIDDLE,
489                seed_coder::FILE_MARKER,
490                multi_region::V0318_END_MARKER,
491                CURSOR_MARKER,
492                multi_region::MARKER_TAG_PREFIX,
493            ];
494            TOKENS
495        }
496        ZetaFormat::V0608QwenMultiRegions => {
497            static TOKENS: &[&str] = &[
498                qwen::FIM_PREFIX,
499                qwen::FIM_SUFFIX,
500                qwen::FIM_MIDDLE,
501                qwen::FILE_MARKER,
502                qwen::END_MARKER,
503                CURSOR_MARKER,
504                multi_region::MARKER_TAG_PREFIX,
505            ];
506            TOKENS
507        }
508        ZetaFormat::V0317SeedMultiRegions => {
509            static TOKENS: &[&str] = &[
510                seed_coder::FIM_SUFFIX,
511                seed_coder::FIM_PREFIX,
512                seed_coder::FIM_MIDDLE,
513                seed_coder::FILE_MARKER,
514                multi_region::V0317_END_MARKER,
515                CURSOR_MARKER,
516                multi_region::RELATIVE_MARKER_TAG_PREFIX,
517            ];
518            TOKENS
519        }
520        ZetaFormat::V0615HashRegions => {
521            static TOKENS: &[&str] = &[
522                seed_coder::FIM_SUFFIX,
523                seed_coder::FIM_PREFIX,
524                seed_coder::FIM_MIDDLE,
525                seed_coder::FILE_MARKER,
526                hashed_regions::V0615_END_MARKER,
527                CURSOR_MARKER,
528                hashed_regions::MARKER_TAG_PREFIX,
529            ];
530            TOKENS
531        }
532        ZetaFormat::V0327SingleFile => {
533            static TOKENS: &[&str] = &[
534                seed_coder::FIM_SUFFIX,
535                seed_coder::FIM_PREFIX,
536                seed_coder::FIM_MIDDLE,
537                seed_coder::FILE_MARKER,
538                multi_region::V0327_END_MARKER,
539                CURSOR_MARKER,
540                multi_region::MARKER_TAG_PREFIX,
541            ];
542            TOKENS
543        }
544        ZetaFormat::V0306SeedMultiRegions => {
545            static TOKENS: &[&str] = &[
546                seed_coder::FIM_SUFFIX,
547                seed_coder::FIM_PREFIX,
548                seed_coder::FIM_MIDDLE,
549                seed_coder::FILE_MARKER,
550                seed_coder::START_MARKER,
551                seed_coder::SEPARATOR,
552                seed_coder::END_MARKER,
553                CURSOR_MARKER,
554                multi_region::MARKER_TAG_PREFIX,
555            ];
556            TOKENS
557        }
558    }
559}
560
561/// Returns the (editable_token_limit, context_token_limit) for a given format.
562pub fn token_limits_for_format(format: ZetaFormat) -> (usize, usize) {
563    match format {
564        ZetaFormat::V0112MiddleAtEnd | ZetaFormat::V0113Ordered => (150, 350),
565        ZetaFormat::V0114180EditableRegion => (180, 350),
566        ZetaFormat::V0120GitMergeMarkers
567        | ZetaFormat::V0131GitMergeMarkersPrefix
568        | ZetaFormat::V0211Prefill
569        | ZetaFormat::V0211SeedCoder
570        | ZetaFormat::V0331SeedCoderModelPy
571        | ZetaFormat::v0226Hashline
572        | ZetaFormat::V0306SeedMultiRegions
573        | ZetaFormat::V0316SeedMultiRegions
574        | ZetaFormat::V0318SeedMultiRegions
575        | ZetaFormat::V0420Diagnostics
576        | ZetaFormat::V0608QwenMultiRegions
577        | ZetaFormat::V0317SeedMultiRegions
578        | ZetaFormat::V0327SingleFile
579        | ZetaFormat::V0304SeedNoEdits => (350, 150),
580        ZetaFormat::V0615HashRegions => (8000, 0),
581
582        ZetaFormat::V0304VariableEdit => (1024, 0),
583    }
584}
585
586pub fn stop_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] {
587    match format {
588        ZetaFormat::v0226Hashline => &[hashline::NO_EDITS_COMMAND_MARKER],
589        ZetaFormat::V0112MiddleAtEnd
590        | ZetaFormat::V0113Ordered
591        | ZetaFormat::V0114180EditableRegion
592        | ZetaFormat::V0120GitMergeMarkers
593        | ZetaFormat::V0131GitMergeMarkersPrefix
594        | ZetaFormat::V0211Prefill
595        | ZetaFormat::V0211SeedCoder
596        | ZetaFormat::V0331SeedCoderModelPy
597        | ZetaFormat::V0304VariableEdit
598        | ZetaFormat::V0306SeedMultiRegions
599        | ZetaFormat::V0304SeedNoEdits => &[],
600        ZetaFormat::V0316SeedMultiRegions => &[multi_region::V0316_END_MARKER],
601        ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => {
602            &[multi_region::V0318_END_MARKER]
603        }
604        ZetaFormat::V0608QwenMultiRegions => &[qwen::END_MARKER],
605        ZetaFormat::V0317SeedMultiRegions => &[multi_region::V0317_END_MARKER],
606        ZetaFormat::V0327SingleFile => &[multi_region::V0327_END_MARKER],
607        ZetaFormat::V0615HashRegions => &[hashed_regions::V0615_END_MARKER],
608    }
609}
610
611/// Delimiters used by response-only SFT (e.g. Unsloth `train_on_responses_only`)
612/// to mask the prompt and train only on the model's completion.
613///
614/// Both strings must appear verbatim in the prompt produced by
615/// [`format_zeta_prompt`] for the same format: `instruction_part` marks the
616/// start of an example, and `response_part` is the final marker before the
617/// completion begins.
618pub struct TrainingDelimiters {
619    pub instruction_part: &'static str,
620    pub response_part: &'static str,
621}
622
623/// Return the response-only training delimiters for a format.
624///
625/// This match is intentionally exhaustive with no wildcard arm so that adding a
626/// new [`ZetaFormat`] fails to compile until its delimiters are specified.
627pub fn training_delimiters_for_format(format: ZetaFormat) -> TrainingDelimiters {
628    match format {
629        ZetaFormat::V0211SeedCoder
630        | ZetaFormat::V0331SeedCoderModelPy
631        | ZetaFormat::V0304SeedNoEdits
632        | ZetaFormat::V0306SeedMultiRegions
633        | ZetaFormat::V0316SeedMultiRegions
634        | ZetaFormat::V0317SeedMultiRegions
635        | ZetaFormat::V0318SeedMultiRegions
636        | ZetaFormat::V0327SingleFile
637        | ZetaFormat::V0420Diagnostics
638        | ZetaFormat::V0615HashRegions => TrainingDelimiters {
639            instruction_part: seed_coder::FIM_SUFFIX,
640            response_part: seed_coder::FIM_MIDDLE,
641        },
642        ZetaFormat::V0608QwenMultiRegions => TrainingDelimiters {
643            instruction_part: qwen::FIM_PREFIX,
644            response_part: qwen::FIM_MIDDLE,
645        },
646        ZetaFormat::V0112MiddleAtEnd
647        | ZetaFormat::V0113Ordered
648        | ZetaFormat::V0114180EditableRegion => TrainingDelimiters {
649            instruction_part: "<|file_sep|>",
650            response_part: "<|fim_middle|>updated\n",
651        },
652        ZetaFormat::V0120GitMergeMarkers => TrainingDelimiters {
653            instruction_part: "<|file_sep|>",
654            response_part: v0120_git_merge_markers::SEPARATOR,
655        },
656        ZetaFormat::V0131GitMergeMarkersPrefix | ZetaFormat::V0211Prefill => TrainingDelimiters {
657            instruction_part: "<|file_sep|>",
658            response_part: "<|fim_middle|>",
659        },
660        ZetaFormat::v0226Hashline => TrainingDelimiters {
661            instruction_part: "<|file_sep|>",
662            response_part: hashline::END_MARKER,
663        },
664        ZetaFormat::V0304VariableEdit => TrainingDelimiters {
665            instruction_part: "<|file_sep|>",
666            response_part: "<|fim_prefix|>",
667        },
668    }
669}
670
671/// Return (editable_range, context_range) for the prompt format
672pub fn excerpt_ranges_for_format(
673    format: ZetaFormat,
674    ranges: &ExcerptRanges,
675) -> (Range<usize>, Range<usize>) {
676    match format {
677        ZetaFormat::V0112MiddleAtEnd | ZetaFormat::V0113Ordered => (
678            ranges.editable_150.clone(),
679            ranges.editable_150_context_350.clone(),
680        ),
681        ZetaFormat::V0114180EditableRegion => (
682            ranges.editable_180.clone(),
683            ranges.editable_180_context_350.clone(),
684        ),
685        ZetaFormat::V0120GitMergeMarkers
686        | ZetaFormat::V0131GitMergeMarkersPrefix
687        | ZetaFormat::V0211Prefill
688        | ZetaFormat::V0211SeedCoder
689        | ZetaFormat::V0331SeedCoderModelPy
690        | ZetaFormat::v0226Hashline
691        | ZetaFormat::V0304SeedNoEdits
692        | ZetaFormat::V0306SeedMultiRegions
693        | ZetaFormat::V0316SeedMultiRegions
694        | ZetaFormat::V0318SeedMultiRegions
695        | ZetaFormat::V0317SeedMultiRegions
696        | ZetaFormat::V0420Diagnostics
697        | ZetaFormat::V0608QwenMultiRegions => (
698            ranges.editable_350.clone(),
699            ranges.editable_350_context_150.clone(),
700        ),
701        ZetaFormat::V0327SingleFile => (
702            ranges.editable_350_context_150.clone(),
703            ranges.context_8192.clone().unwrap_or(
704                // shouldn't be used, only for compat with old data/clients
705                ranges.editable_350_context_150.clone(),
706            ),
707        ),
708        ZetaFormat::V0615HashRegions => (
709            ranges
710                .context_8192
711                .clone()
712                .unwrap_or_else(|| ranges.editable_350_context_150.clone()),
713            ranges
714                .context_8192
715                .clone()
716                .unwrap_or_else(|| ranges.editable_350_context_150.clone()),
717        ),
718
719        ZetaFormat::V0304VariableEdit => {
720            let context = ranges
721                .editable_350_context_1024
722                .clone()
723                .or(ranges.editable_350_context_512.clone())
724                .unwrap_or_else(|| ranges.editable_350_context_150.clone());
725            (context.clone(), context)
726        }
727    }
728}
729
730pub fn write_cursor_excerpt_section_for_format(
731    format: ZetaFormat,
732    prompt: &mut String,
733    path: &Path,
734    context: &str,
735    editable_range: &Range<usize>,
736    cursor_offset: usize,
737) {
738    match format {
739        ZetaFormat::V0112MiddleAtEnd => v0112_middle_at_end::write_cursor_excerpt_section(
740            prompt,
741            path,
742            context,
743            editable_range,
744            cursor_offset,
745        ),
746        ZetaFormat::V0113Ordered | ZetaFormat::V0114180EditableRegion => {
747            v0113_ordered::write_cursor_excerpt_section(
748                prompt,
749                path,
750                context,
751                editable_range,
752                cursor_offset,
753            )
754        }
755        ZetaFormat::V0120GitMergeMarkers => v0120_git_merge_markers::write_cursor_excerpt_section(
756            prompt,
757            path,
758            context,
759            editable_range,
760            cursor_offset,
761        ),
762        ZetaFormat::V0131GitMergeMarkersPrefix | ZetaFormat::V0211Prefill => {
763            v0131_git_merge_markers_prefix::write_cursor_excerpt_section(
764                prompt,
765                path,
766                context,
767                editable_range,
768                cursor_offset,
769            )
770        }
771        ZetaFormat::V0211SeedCoder
772        | ZetaFormat::V0331SeedCoderModelPy
773        | ZetaFormat::V0304SeedNoEdits => seed_coder::write_cursor_excerpt_section(
774            prompt,
775            path,
776            context,
777            editable_range,
778            cursor_offset,
779        ),
780        ZetaFormat::v0226Hashline => hashline::write_cursor_excerpt_section(
781            prompt,
782            path,
783            context,
784            editable_range,
785            cursor_offset,
786        ),
787        ZetaFormat::V0304VariableEdit => {
788            v0304_variable_edit::write_cursor_excerpt_section(prompt, path, context, cursor_offset)
789        }
790        ZetaFormat::V0306SeedMultiRegions => {
791            prompt.push_str(&build_v0306_cursor_prefix(
792                path,
793                context,
794                editable_range,
795                cursor_offset,
796            ));
797        }
798        ZetaFormat::V0316SeedMultiRegions => {
799            prompt.push_str(&build_v0316_cursor_prefix(
800                path,
801                context,
802                editable_range,
803                cursor_offset,
804            ));
805        }
806        ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => {
807            prompt.push_str(&build_v0318_cursor_prefix(
808                path,
809                context,
810                editable_range,
811                cursor_offset,
812                seed_coder::FILE_MARKER,
813            ));
814        }
815        ZetaFormat::V0608QwenMultiRegions => {
816            prompt.push_str(&build_v0318_cursor_prefix(
817                path,
818                context,
819                editable_range,
820                cursor_offset,
821                qwen::FILE_MARKER,
822            ));
823        }
824        ZetaFormat::V0317SeedMultiRegions => {
825            prompt.push_str(&build_v0317_cursor_prefix(
826                path,
827                context,
828                editable_range,
829                cursor_offset,
830            ));
831        }
832        ZetaFormat::V0327SingleFile => {
833            prompt.push_str(&build_v0318_cursor_prefix(
834                path,
835                context,
836                editable_range,
837                cursor_offset,
838                seed_coder::FILE_MARKER,
839            ));
840        }
841        ZetaFormat::V0615HashRegions => {
842            prompt.push_str(&build_v0615_cursor_prefix(
843                path,
844                context,
845                editable_range,
846                cursor_offset,
847            ));
848        }
849    }
850}
851
852fn build_v0306_cursor_prefix(
853    path: &Path,
854    context: &str,
855    editable_range: &Range<usize>,
856    cursor_offset: usize,
857) -> String {
858    let mut section = String::new();
859    let path_str = path.to_string_lossy();
860    write!(section, "{}{}\n", seed_coder::FILE_MARKER, path_str).ok();
861
862    section.push_str(&context[..editable_range.start]);
863    section.push_str(seed_coder::START_MARKER);
864
865    let editable_text = &context[editable_range.clone()];
866    let cursor_in_editable = cursor_offset - editable_range.start;
867    multi_region::write_editable_with_markers(
868        &mut section,
869        editable_text,
870        cursor_in_editable,
871        CURSOR_MARKER,
872    );
873
874    if !section.ends_with('\n') {
875        section.push('\n');
876    }
877    section.push_str(seed_coder::SEPARATOR);
878    section
879}
880
881fn build_v0316_cursor_prefix(
882    path: &Path,
883    context: &str,
884    editable_range: &Range<usize>,
885    cursor_offset: usize,
886) -> String {
887    let mut section = String::new();
888    let path_str = path.to_string_lossy();
889    write!(section, "{}{}\n", seed_coder::FILE_MARKER, path_str).ok();
890
891    section.push_str(&context[..editable_range.start]);
892
893    let editable_text = &context[editable_range.clone()];
894    let cursor_in_editable = cursor_offset - editable_range.start;
895    multi_region::write_editable_with_markers_v0316(
896        &mut section,
897        editable_text,
898        cursor_in_editable,
899        CURSOR_MARKER,
900    );
901
902    if !section.ends_with('\n') {
903        section.push('\n');
904    }
905    section
906}
907
908fn build_v0318_cursor_prefix(
909    path: &Path,
910    context: &str,
911    editable_range: &Range<usize>,
912    cursor_offset: usize,
913    file_marker: &str,
914) -> String {
915    let mut section = String::new();
916    let path_str = path.to_string_lossy();
917    write!(section, "{}{}\n", file_marker, path_str).ok();
918
919    section.push_str(&context[..editable_range.start]);
920
921    let editable_text = &context[editable_range.clone()];
922    let cursor_in_editable = cursor_offset - editable_range.start;
923    multi_region::write_editable_with_markers_v0318(
924        &mut section,
925        editable_text,
926        cursor_in_editable,
927        CURSOR_MARKER,
928    );
929
930    if !section.ends_with('\n') {
931        section.push('\n');
932    }
933    section
934}
935
936fn build_v0615_cursor_prefix(
937    path: &Path,
938    context: &str,
939    editable_range: &Range<usize>,
940    cursor_offset: usize,
941) -> String {
942    let mut section = String::new();
943    let path_str = path.to_string_lossy();
944    write!(section, "{}{}\n", seed_coder::FILE_MARKER, path_str).ok();
945
946    section.push_str(&context[..editable_range.start]);
947
948    let editable_text = &context[editable_range.clone()];
949    let cursor_in_editable = cursor_offset - editable_range.start;
950    let markers = hashed_regions::markers_for_text(editable_text);
951    hashed_regions::write_snippet_with_markers(
952        &mut section,
953        editable_text,
954        &markers,
955        Some((cursor_in_editable, CURSOR_MARKER)),
956    );
957
958    if !section.ends_with('\n') {
959        section.push('\n');
960    }
961    section
962}
963
964fn build_v0317_cursor_prefix(
965    path: &Path,
966    context: &str,
967    editable_range: &Range<usize>,
968    cursor_offset: usize,
969) -> String {
970    let mut section = String::new();
971    let path_str = path.to_string_lossy();
972    write!(section, "{}{}\n", seed_coder::FILE_MARKER, path_str).ok();
973
974    section.push_str(&context[..editable_range.start]);
975
976    let editable_text = &context[editable_range.clone()];
977    let cursor_in_editable = cursor_offset - editable_range.start;
978    multi_region::write_editable_with_markers_v0317(
979        &mut section,
980        editable_text,
981        cursor_in_editable,
982        CURSOR_MARKER,
983    );
984
985    if !section.ends_with('\n') {
986        section.push('\n');
987    }
988    section
989}
990
991fn offset_range_to_row_range(text: &str, range: Range<usize>) -> Range<u32> {
992    let start_row = text[0..range.start].matches('\n').count() as u32;
993    let mut end_row = start_row + text[range.clone()].matches('\n').count() as u32;
994    if !text[..range.end].ends_with('\n') {
995        end_row += 1;
996    }
997    return start_row..end_row;
998}
999
1000fn assemble_single_file_fim_prompt(
1001    context: &str,
1002    editable_range: &Range<usize>,
1003    cursor_prefix_section: &str,
1004    events: &[Arc<Event>],
1005    max_tokens: usize,
1006) -> String {
1007    let suffix_section = seed_coder::build_suffix_section(context, editable_range);
1008
1009    let suffix_tokens = estimate_tokens(suffix_section.len() + seed_coder::FIM_PREFIX.len());
1010    let cursor_prefix_tokens =
1011        estimate_tokens(cursor_prefix_section.len() + seed_coder::FIM_MIDDLE.len());
1012    let budget_after_cursor = max_tokens.saturating_sub(suffix_tokens + cursor_prefix_tokens);
1013
1014    let edit_history_section = format_edit_history_within_budget(
1015        events,
1016        seed_coder::FILE_MARKER,
1017        "edit_history",
1018        budget_after_cursor,
1019        max_edit_event_count_for_format(&ZetaFormat::V0327SingleFile),
1020    );
1021
1022    let mut prompt = String::new();
1023    prompt.push_str(&suffix_section);
1024    prompt.push_str(seed_coder::FIM_PREFIX);
1025    prompt.push_str(&edit_history_section);
1026    if !edit_history_section.is_empty() {
1027        prompt.push('\n');
1028    }
1029    prompt.push_str(cursor_prefix_section);
1030    prompt.push_str(seed_coder::FIM_MIDDLE);
1031    prompt
1032}
1033
1034fn format_hash_region_related_files_within_budget(
1035    input: &Zeta2PromptInput,
1036    marker_table: &[hashed_regions::SnippetMarkers],
1037    cursor: &hashed_regions::RelatedFileCursor,
1038    max_tokens: usize,
1039) -> Option<String> {
1040    let related_files = input.related_files.as_deref()?;
1041
1042    struct RenderedExcerpt {
1043        file_ix: usize,
1044        excerpt_ix: usize,
1045        order: usize,
1046        rendered: String,
1047    }
1048
1049    let mut candidates = Vec::new();
1050    let mut required_candidate_ix = None;
1051    for (file_ix, file) in related_files.iter().enumerate() {
1052        for (excerpt_ix, excerpt) in file.excerpts.iter().enumerate() {
1053            let markers =
1054                hashed_regions::marker_table_for_excerpt(marker_table, file_ix, excerpt_ix);
1055            let mut rendered = String::new();
1056            if let Some(markers) = markers {
1057                let cursor_in_excerpt = (file_ix == cursor.file_ix
1058                    && excerpt_ix == cursor.excerpt_ix)
1059                    .then_some((cursor.offset_in_excerpt, CURSOR_MARKER));
1060                hashed_regions::write_snippet_with_markers(
1061                    &mut rendered,
1062                    &excerpt.text,
1063                    markers,
1064                    cursor_in_excerpt,
1065                );
1066            } else {
1067                rendered.push_str(&excerpt.text);
1068            }
1069            if !rendered.ends_with('\n') {
1070                rendered.push('\n');
1071            }
1072
1073            if file_ix == cursor.file_ix && excerpt_ix == cursor.excerpt_ix {
1074                required_candidate_ix = Some(candidates.len());
1075            }
1076
1077            candidates.push(RenderedExcerpt {
1078                file_ix,
1079                excerpt_ix,
1080                order: excerpt.order,
1081                rendered,
1082            });
1083        }
1084    }
1085
1086    let required_candidate_ix = required_candidate_ix?;
1087    let file_headers: Vec<String> = related_files
1088        .iter()
1089        .map(|file| {
1090            let path = hashed_regions::related_file_patch_path(&input.cursor_path, &file.path)
1091                .iter()
1092                .map(|component| component.to_string_lossy())
1093                .collect::<Vec<_>>()
1094                .join("/");
1095            format!("{}{path}\n", seed_coder::FILE_MARKER)
1096        })
1097        .collect();
1098
1099    let mut total_tokens = 0;
1100    let mut included = vec![false; candidates.len()];
1101    let mut file_included = vec![false; related_files.len()];
1102
1103    let required = &candidates[required_candidate_ix];
1104    let required_cost =
1105        estimate_tokens(file_headers[required.file_ix].len() + required.rendered.len());
1106    if required_cost > max_tokens {
1107        return None;
1108    }
1109    total_tokens += required_cost;
1110    included[required_candidate_ix] = true;
1111    file_included[required.file_ix] = true;
1112
1113    let mut selection_order: Vec<usize> = (0..candidates.len()).collect();
1114    selection_order.sort_by_key(|&candidate_ix| {
1115        let candidate = &candidates[candidate_ix];
1116        (candidate.order, candidate.file_ix, candidate.excerpt_ix)
1117    });
1118
1119    for candidate_ix in selection_order {
1120        if included[candidate_ix] {
1121            continue;
1122        }
1123        let candidate = &candidates[candidate_ix];
1124        let header_cost = if file_included[candidate.file_ix] {
1125            0
1126        } else {
1127            estimate_tokens(file_headers[candidate.file_ix].len())
1128        };
1129        let excerpt_cost = estimate_tokens(candidate.rendered.len());
1130        if total_tokens + header_cost + excerpt_cost > max_tokens {
1131            continue;
1132        }
1133        total_tokens += header_cost + excerpt_cost;
1134        included[candidate_ix] = true;
1135        file_included[candidate.file_ix] = true;
1136    }
1137
1138    let mut result = String::new();
1139    let mut last_file_ix = None;
1140    for (candidate_ix, candidate) in candidates.iter().enumerate() {
1141        if !included[candidate_ix] {
1142            continue;
1143        }
1144        if last_file_ix != Some(candidate.file_ix) {
1145            result.push_str(&file_headers[candidate.file_ix]);
1146            last_file_ix = Some(candidate.file_ix);
1147        }
1148        result.push_str(&candidate.rendered);
1149
1150        let file = &related_files[candidate.file_ix];
1151        let excerpt = &file.excerpts[candidate.excerpt_ix];
1152        let next_excerpt_start = candidates
1153            .iter()
1154            .enumerate()
1155            .skip(candidate_ix + 1)
1156            .find(|(next_ix, next)| included[*next_ix] && next.file_ix == candidate.file_ix)
1157            .map(|(_, next)| file.excerpts[next.excerpt_ix].row_range.start);
1158        if rows_omitted_after_excerpt(excerpt, next_excerpt_start, file.max_row) {
1159            result.push_str("...\n");
1160        }
1161    }
1162
1163    Some(result)
1164}
1165
1166fn format_hash_regions_prompt_with_budget(
1167    input: &Zeta2PromptInput,
1168    max_tokens: usize,
1169) -> Option<String> {
1170    let marker_table = hashed_regions::build_marker_table(input);
1171    let cursor = hashed_regions::locate_cursor_in_related_files(input)?;
1172    hashed_regions::marker_table_for_excerpt(&marker_table, cursor.file_ix, cursor.excerpt_ix)?;
1173
1174    let fixed_tokens = estimate_tokens(
1175        seed_coder::FIM_SUFFIX.len()
1176            + "\n".len()
1177            + seed_coder::FIM_PREFIX.len()
1178            + seed_coder::FIM_MIDDLE.len(),
1179    );
1180    let related_files_budget = max_tokens.saturating_sub(fixed_tokens);
1181    let related_files_section = format_hash_region_related_files_within_budget(
1182        input,
1183        &marker_table,
1184        &cursor,
1185        related_files_budget,
1186    )?;
1187
1188    let mut prompt = String::new();
1189    prompt.push_str(seed_coder::FIM_SUFFIX);
1190    prompt.push('\n');
1191    prompt.push_str(seed_coder::FIM_PREFIX);
1192    prompt.push_str(&related_files_section);
1193    if !prompt.ends_with('\n') {
1194        prompt.push('\n');
1195    }
1196    prompt.push_str(seed_coder::FIM_MIDDLE);
1197    Some(prompt)
1198}
1199
1200pub fn format_prompt_with_budget_for_format(
1201    input: &Zeta2PromptInput,
1202    format: ZetaFormat,
1203    max_tokens: usize,
1204) -> Option<String> {
1205    if format == ZetaFormat::V0615HashRegions {
1206        return format_hash_regions_prompt_with_budget(
1207            input,
1208            apply_prompt_budget_margin(max_tokens),
1209        );
1210    }
1211
1212    let (context, editable_range, context_range, cursor_offset) =
1213        resolve_cursor_region(input, format);
1214    let empty_files = Vec::new();
1215    let input_related_files = input.related_files.as_deref().unwrap_or(&empty_files);
1216    let filtered_related_files = if format == ZetaFormat::V0615HashRegions {
1217        input_related_files.to_vec()
1218    } else if let Some(cursor_excerpt_start_row) = input.excerpt_start_row {
1219        let relative_row_range =
1220            offset_range_to_row_range(&input.cursor_excerpt, context_range.clone());
1221        let row_range = relative_row_range.start + cursor_excerpt_start_row
1222            ..relative_row_range.end + cursor_excerpt_start_row;
1223        filter_redundant_excerpts(
1224            input_related_files.to_vec(),
1225            input.cursor_path.as_ref(),
1226            row_range,
1227        )
1228    } else {
1229        input_related_files.to_vec()
1230    };
1231    let cursor_buffer_row = input.excerpt_start_row.map(|excerpt_start_row| {
1232        excerpt_start_row
1233            + input.cursor_excerpt[..context_range.start + cursor_offset]
1234                .bytes()
1235                .filter(|byte| *byte == b'\n')
1236                .count() as u32
1237    });
1238
1239    format_resolved_prompt_with_budget(
1240        format,
1241        input.cursor_path.as_ref(),
1242        context,
1243        &editable_range,
1244        cursor_offset,
1245        &input.events,
1246        &filtered_related_files,
1247        &input.active_buffer_diagnostics,
1248        cursor_buffer_row,
1249        max_tokens,
1250    )
1251}
1252
1253fn format_resolved_prompt_with_budget(
1254    format: ZetaFormat,
1255    path: &Path,
1256    context: &str,
1257    editable_range: &Range<usize>,
1258    cursor_offset: usize,
1259    events: &[Arc<Event>],
1260    related_files: &[RelatedFile],
1261    active_buffer_diagnostics: &[ActiveBufferDiagnostic],
1262    cursor_buffer_row: Option<u32>,
1263    max_tokens: usize,
1264) -> Option<String> {
1265    let prompt = match format {
1266        ZetaFormat::V0211SeedCoder
1267        | ZetaFormat::V0331SeedCoderModelPy
1268        | ZetaFormat::V0304SeedNoEdits
1269        | ZetaFormat::V0306SeedMultiRegions
1270        | ZetaFormat::V0316SeedMultiRegions
1271        | ZetaFormat::V0318SeedMultiRegions
1272        | ZetaFormat::V0317SeedMultiRegions
1273        | ZetaFormat::V0420Diagnostics => {
1274            let mut cursor_section = String::new();
1275
1276            write_cursor_excerpt_section_for_format(
1277                format,
1278                &mut cursor_section,
1279                path,
1280                context,
1281                editable_range,
1282                cursor_offset,
1283            );
1284
1285            let budget_with_margin = apply_prompt_budget_margin(max_tokens);
1286            seed_coder::assemble_fim_prompt(
1287                context,
1288                editable_range,
1289                &cursor_section,
1290                events,
1291                related_files,
1292                if format == ZetaFormat::V0420Diagnostics {
1293                    active_buffer_diagnostics
1294                } else {
1295                    &[]
1296                },
1297                cursor_buffer_row,
1298                budget_with_margin,
1299            )
1300        }
1301        ZetaFormat::V0608QwenMultiRegions => {
1302            let mut cursor_section = String::new();
1303
1304            write_cursor_excerpt_section_for_format(
1305                format,
1306                &mut cursor_section,
1307                path,
1308                context,
1309                editable_range,
1310                cursor_offset,
1311            );
1312
1313            qwen::assemble_fim_prompt(
1314                context,
1315                editable_range,
1316                &cursor_section,
1317                events,
1318                related_files,
1319                apply_prompt_budget_margin(max_tokens),
1320            )
1321        }
1322        ZetaFormat::V0327SingleFile => {
1323            let mut cursor_section = String::new();
1324            write_cursor_excerpt_section_for_format(
1325                format,
1326                &mut cursor_section,
1327                path,
1328                context,
1329                editable_range,
1330                cursor_offset,
1331            );
1332
1333            assemble_single_file_fim_prompt(
1334                context,
1335                editable_range,
1336                &cursor_section,
1337                events,
1338                apply_prompt_budget_margin(max_tokens),
1339            )
1340        }
1341        _ => {
1342            let mut cursor_section = String::new();
1343            write_cursor_excerpt_section_for_format(
1344                format,
1345                &mut cursor_section,
1346                path,
1347                context,
1348                editable_range,
1349                cursor_offset,
1350            );
1351
1352            let mut remaining_budget = apply_prompt_budget_margin(max_tokens);
1353            let cursor_tokens = estimate_tokens(cursor_section.len());
1354            remaining_budget = remaining_budget.saturating_sub(cursor_tokens);
1355
1356            let edit_history_section = format_edit_history_within_budget(
1357                events,
1358                "<|file_sep|>",
1359                "edit history",
1360                remaining_budget,
1361                max_edit_event_count_for_format(&format),
1362            );
1363            let edit_history_tokens = estimate_tokens(edit_history_section.len());
1364            remaining_budget = remaining_budget.saturating_sub(edit_history_tokens);
1365
1366            let related_files_section = format_related_files_within_budget(
1367                related_files,
1368                "<|file_sep|>",
1369                "",
1370                remaining_budget,
1371            );
1372
1373            let mut prompt = String::new();
1374            prompt.push_str(&related_files_section);
1375            prompt.push_str(&edit_history_section);
1376            prompt.push_str(&cursor_section);
1377            prompt
1378        }
1379    };
1380    let prompt_tokens = estimate_tokens(prompt.len());
1381    if prompt_tokens > max_tokens {
1382        return None;
1383    }
1384    return Some(prompt);
1385}
1386
1387pub fn format_active_buffer_diagnostics_with_budget(
1388    diagnostics: &[ActiveBufferDiagnostic],
1389    cursor_buffer_row: Option<u32>,
1390    budget: usize,
1391) -> String {
1392    if diagnostics.is_empty() || budget == 0 {
1393        return String::new();
1394    }
1395
1396    const MAX_DIAGNOSTICS: usize = 10;
1397
1398    let mut diagnostic_indices = (0..diagnostics.len()).collect::<Vec<_>>();
1399    if let Some(cursor_buffer_row) = cursor_buffer_row {
1400        let distance = |index: &usize| {
1401            let range = &diagnostics[*index].snippet_buffer_row_range;
1402            u32::abs_diff(cursor_buffer_row, range.start)
1403                + u32::abs_diff(cursor_buffer_row, range.end)
1404        };
1405        // Only the closest `MAX_DIAGNOSTICS` are rendered below, so select that
1406        // prefix instead of fully sorting every diagnostic.
1407        if diagnostic_indices.len() > MAX_DIAGNOSTICS {
1408            diagnostic_indices.select_nth_unstable_by_key(MAX_DIAGNOSTICS, &distance);
1409            diagnostic_indices.truncate(MAX_DIAGNOSTICS);
1410        }
1411        diagnostic_indices.sort_unstable_by_key(&distance);
1412    }
1413
1414    let mut output = format!("{}diagnostics\n", seed_coder::FILE_MARKER);
1415    let header_tokens = estimate_tokens(output.len());
1416    if header_tokens > budget {
1417        return String::new();
1418    }
1419
1420    let mut used_tokens = header_tokens;
1421    let mut included_diagnostics = 0;
1422    for diagnostic_index in diagnostic_indices.into_iter().take(MAX_DIAGNOSTICS) {
1423        let diagnostic = &diagnostics[diagnostic_index];
1424        let snippet = clamp_text_to_token_count(&diagnostic.snippet, 256);
1425
1426        let diagnostic_section = if snippet.is_empty() {
1427            format!("*{}*\n", diagnostic.message)
1428        } else {
1429            format!(
1430                "*{}*:\n```\n{}{}\n```\n",
1431                diagnostic.message,
1432                snippet,
1433                if snippet.len() < diagnostic.snippet.len() {
1434                    "..."
1435                } else {
1436                    ""
1437                }
1438            )
1439        };
1440        let diagnostic_tokens = estimate_tokens(diagnostic_section.len());
1441        if used_tokens + diagnostic_tokens > budget {
1442            break;
1443        }
1444        output.push_str(&diagnostic_section);
1445        used_tokens += diagnostic_tokens;
1446        included_diagnostics += 1;
1447    }
1448
1449    if included_diagnostics == 0 {
1450        String::new()
1451    } else {
1452        output
1453    }
1454}
1455
1456pub fn filter_redundant_excerpts(
1457    mut related_files: Vec<RelatedFile>,
1458    cursor_path: &Path,
1459    cursor_row_range: Range<u32>,
1460) -> Vec<RelatedFile> {
1461    for file in &mut related_files {
1462        if file.path.as_ref() == cursor_path {
1463            file.excerpts.retain(|excerpt| {
1464                excerpt.row_range.start < cursor_row_range.start
1465                    || excerpt.row_range.end > cursor_row_range.end
1466            });
1467        }
1468    }
1469    related_files.retain(|file| !file.excerpts.is_empty());
1470    related_files
1471}
1472
1473pub fn max_edit_event_count_for_format(format: &ZetaFormat) -> usize {
1474    match format {
1475        ZetaFormat::V0112MiddleAtEnd
1476        | ZetaFormat::V0113Ordered
1477        | ZetaFormat::V0114180EditableRegion
1478        | ZetaFormat::V0120GitMergeMarkers
1479        | ZetaFormat::V0131GitMergeMarkersPrefix
1480        | ZetaFormat::V0211Prefill
1481        | ZetaFormat::V0211SeedCoder
1482        | ZetaFormat::V0331SeedCoderModelPy
1483        | ZetaFormat::v0226Hashline
1484        | ZetaFormat::V0304SeedNoEdits
1485        | ZetaFormat::V0304VariableEdit
1486        | ZetaFormat::V0306SeedMultiRegions
1487        | ZetaFormat::V0316SeedMultiRegions
1488        | ZetaFormat::V0318SeedMultiRegions
1489        | ZetaFormat::V0317SeedMultiRegions
1490        | ZetaFormat::V0420Diagnostics
1491        | ZetaFormat::V0608QwenMultiRegions
1492        | ZetaFormat::V0327SingleFile
1493        | ZetaFormat::V0615HashRegions => 6,
1494    }
1495}
1496
1497pub fn get_prefill_for_format(
1498    format: ZetaFormat,
1499    context: &str,
1500    editable_range: &Range<usize>,
1501) -> String {
1502    match format {
1503        ZetaFormat::V0211Prefill => v0211_prefill::get_prefill(context, editable_range),
1504        ZetaFormat::V0112MiddleAtEnd
1505        | ZetaFormat::V0113Ordered
1506        | ZetaFormat::V0114180EditableRegion
1507        | ZetaFormat::V0120GitMergeMarkers
1508        | ZetaFormat::V0131GitMergeMarkersPrefix
1509        | ZetaFormat::V0211SeedCoder
1510        | ZetaFormat::V0331SeedCoderModelPy
1511        | ZetaFormat::v0226Hashline
1512        | ZetaFormat::V0304VariableEdit => String::new(),
1513        ZetaFormat::V0304SeedNoEdits
1514        | ZetaFormat::V0306SeedMultiRegions
1515        | ZetaFormat::V0316SeedMultiRegions
1516        | ZetaFormat::V0318SeedMultiRegions
1517        | ZetaFormat::V0317SeedMultiRegions
1518        | ZetaFormat::V0420Diagnostics
1519        | ZetaFormat::V0608QwenMultiRegions
1520        | ZetaFormat::V0327SingleFile
1521        | ZetaFormat::V0615HashRegions => String::new(),
1522    }
1523}
1524
1525pub fn output_end_marker_for_format(format: ZetaFormat) -> Option<&'static str> {
1526    match format {
1527        ZetaFormat::V0120GitMergeMarkers => Some(v0120_git_merge_markers::END_MARKER),
1528        ZetaFormat::V0131GitMergeMarkersPrefix => Some(v0131_git_merge_markers_prefix::END_MARKER),
1529        ZetaFormat::V0211Prefill => Some(v0131_git_merge_markers_prefix::END_MARKER),
1530        ZetaFormat::V0211SeedCoder
1531        | ZetaFormat::V0331SeedCoderModelPy
1532        | ZetaFormat::V0304SeedNoEdits
1533        | ZetaFormat::V0306SeedMultiRegions => Some(seed_coder::END_MARKER),
1534        ZetaFormat::V0316SeedMultiRegions => Some(multi_region::V0316_END_MARKER),
1535        ZetaFormat::V0318SeedMultiRegions => Some(multi_region::V0318_END_MARKER),
1536        ZetaFormat::V0420Diagnostics => Some(multi_region::V0318_END_MARKER),
1537        ZetaFormat::V0608QwenMultiRegions => Some(qwen::END_MARKER),
1538        ZetaFormat::V0317SeedMultiRegions => Some(multi_region::V0317_END_MARKER),
1539        ZetaFormat::V0327SingleFile => Some(multi_region::V0327_END_MARKER),
1540        ZetaFormat::V0615HashRegions => Some(hashed_regions::V0615_END_MARKER),
1541
1542        ZetaFormat::V0112MiddleAtEnd
1543        | ZetaFormat::V0113Ordered
1544        | ZetaFormat::V0114180EditableRegion
1545        | ZetaFormat::v0226Hashline
1546        | ZetaFormat::V0304VariableEdit => None,
1547    }
1548}
1549
1550pub fn encode_patch_as_output_for_format(
1551    format: ZetaFormat,
1552    old_editable_region: &str,
1553    patch: &str,
1554    cursor_offset: Option<usize>,
1555) -> Result<Option<String>> {
1556    match format {
1557        ZetaFormat::v0226Hashline => {
1558            hashline::patch_to_edit_commands(old_editable_region, patch, cursor_offset).map(Some)
1559        }
1560        ZetaFormat::V0304VariableEdit => v0304_variable_edit::patch_to_variable_edit_output(
1561            old_editable_region,
1562            patch,
1563            cursor_offset,
1564        )
1565        .map(Some),
1566        ZetaFormat::V0304SeedNoEdits | ZetaFormat::V0306SeedMultiRegions => {
1567            Ok(seed_coder::no_edits(patch))
1568        }
1569        ZetaFormat::V0316SeedMultiRegions => {
1570            let empty_patch = patch.lines().count() <= 3;
1571            if empty_patch {
1572                let marker_offsets = multi_region::compute_marker_offsets(old_editable_region);
1573                let marker_num =
1574                    multi_region::nearest_marker_number(cursor_offset, &marker_offsets);
1575                let tag = multi_region::marker_tag(marker_num);
1576                Ok(Some(format!(
1577                    "{tag}{tag}{}",
1578                    multi_region::V0316_END_MARKER
1579                )))
1580            } else {
1581                Ok(None)
1582            }
1583        }
1584        ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => {
1585            let empty_patch = patch.lines().count() <= 3;
1586            if empty_patch {
1587                let marker_offsets =
1588                    multi_region::compute_marker_offsets_v0318(old_editable_region);
1589                let marker_num =
1590                    multi_region::nearest_marker_number(cursor_offset, &marker_offsets);
1591                let tag = multi_region::marker_tag(marker_num);
1592                Ok(Some(format!(
1593                    "{tag}{tag}{}",
1594                    multi_region::V0318_END_MARKER
1595                )))
1596            } else {
1597                Ok(None)
1598            }
1599        }
1600        ZetaFormat::V0608QwenMultiRegions => {
1601            let empty_patch = patch.lines().count() <= 3;
1602            if empty_patch {
1603                let marker_offsets =
1604                    multi_region::compute_marker_offsets_v0318(old_editable_region);
1605                let marker_num =
1606                    multi_region::nearest_marker_number(cursor_offset, &marker_offsets);
1607                let tag = multi_region::marker_tag(marker_num);
1608                Ok(Some(format!("{tag}{tag}{}", qwen::END_MARKER)))
1609            } else {
1610                Ok(None)
1611            }
1612        }
1613        ZetaFormat::V0317SeedMultiRegions => {
1614            let empty_patch = patch.lines().count() <= 3;
1615            if empty_patch {
1616                let tag = multi_region::marker_tag_relative(0);
1617                Ok(Some(format!(
1618                    "{tag}{tag}{}",
1619                    multi_region::V0317_END_MARKER
1620                )))
1621            } else {
1622                Ok(None)
1623            }
1624        }
1625        ZetaFormat::V0327SingleFile => {
1626            let empty_patch = patch.lines().count() <= 3;
1627            if empty_patch {
1628                let marker_offsets =
1629                    multi_region::compute_marker_offsets_v0318(old_editable_region);
1630                let marker_num =
1631                    multi_region::nearest_marker_number(cursor_offset, &marker_offsets);
1632                let tag = multi_region::marker_tag(marker_num);
1633                Ok(Some(format!(
1634                    "{tag}{tag}{}",
1635                    multi_region::V0327_END_MARKER
1636                )))
1637            } else {
1638                Ok(None)
1639            }
1640        }
1641        ZetaFormat::V0615HashRegions => Ok(None),
1642        _ => Ok(None),
1643    }
1644}
1645
1646/// Given a `Zeta2PromptInput`, a format, and a patch (with cursor already
1647/// extracted), produce the expected model output string for training.
1648pub fn format_expected_output(
1649    input: &Zeta2PromptInput,
1650    format: ZetaFormat,
1651    patch: &str,
1652    cursor_offset: Option<usize>,
1653) -> Result<String> {
1654    if format == ZetaFormat::V0615HashRegions {
1655        return hashed_regions::encode_patch_as_output(input, patch, cursor_offset, CURSOR_MARKER);
1656    }
1657
1658    let (context, editable_range, _, _) = resolve_cursor_region(input, format);
1659    let mut old_editable = context[editable_range].to_string();
1660    if !old_editable.is_empty() && !old_editable.ends_with('\n') {
1661        old_editable.push('\n');
1662    }
1663
1664    // Formats with their own output encoding (hashline, variable-edit,
1665    // multi-region empty patches) are handled here.
1666    if let Some(output) =
1667        encode_patch_as_output_for_format(format, &old_editable, patch, cursor_offset)?
1668    {
1669        return Ok(output);
1670    }
1671
1672    let empty_patch = patch.lines().count() <= 3;
1673
1674    match format {
1675        // Multi-region formats: non-empty patches need diff application
1676        // then marker-span encoding.
1677        ZetaFormat::V0316SeedMultiRegions => {
1678            let (new_editable, first_hunk_offset) =
1679                udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?;
1680            let cursor_in_new = cursor_in_new_text(cursor_offset, first_hunk_offset, &new_editable);
1681            multi_region::encode_from_old_and_new_v0316(
1682                &old_editable,
1683                &new_editable,
1684                cursor_in_new,
1685                CURSOR_MARKER,
1686                multi_region::V0316_END_MARKER,
1687            )
1688        }
1689        ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => {
1690            let (new_editable, first_hunk_offset) =
1691                udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?;
1692            let cursor_in_new = cursor_in_new_text(cursor_offset, first_hunk_offset, &new_editable);
1693            multi_region::encode_from_old_and_new_v0318(
1694                &old_editable,
1695                &new_editable,
1696                cursor_in_new,
1697                CURSOR_MARKER,
1698                multi_region::V0318_END_MARKER,
1699            )
1700        }
1701        ZetaFormat::V0608QwenMultiRegions => {
1702            let (new_editable, first_hunk_offset) =
1703                udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?;
1704            let cursor_in_new = cursor_in_new_text(cursor_offset, first_hunk_offset, &new_editable);
1705            multi_region::encode_from_old_and_new_v0318(
1706                &old_editable,
1707                &new_editable,
1708                cursor_in_new,
1709                CURSOR_MARKER,
1710                qwen::END_MARKER,
1711            )
1712        }
1713        ZetaFormat::V0327SingleFile => {
1714            let (new_editable, first_hunk_offset) =
1715                udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?;
1716            let cursor_in_new = cursor_in_new_text(cursor_offset, first_hunk_offset, &new_editable);
1717            multi_region::encode_from_old_and_new_v0318(
1718                &old_editable,
1719                &new_editable,
1720                cursor_in_new,
1721                CURSOR_MARKER,
1722                multi_region::V0327_END_MARKER,
1723            )
1724        }
1725        ZetaFormat::V0317SeedMultiRegions => {
1726            let (new_editable, first_hunk_offset) =
1727                udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?;
1728            let cursor_in_new = cursor_in_new_text(cursor_offset, first_hunk_offset, &new_editable);
1729            multi_region::encode_from_old_and_new_v0317(
1730                &old_editable,
1731                &new_editable,
1732                cursor_in_new,
1733                CURSOR_MARKER,
1734                multi_region::V0317_END_MARKER,
1735            )
1736        }
1737        // V0131-style formats and fallback: produce new editable text with
1738        // cursor marker inserted, followed by the end marker.
1739        ZetaFormat::V0112MiddleAtEnd
1740        | ZetaFormat::V0113Ordered
1741        | ZetaFormat::V0114180EditableRegion
1742        | ZetaFormat::V0120GitMergeMarkers
1743        | ZetaFormat::V0131GitMergeMarkersPrefix
1744        | ZetaFormat::V0211Prefill
1745        | ZetaFormat::V0211SeedCoder
1746        | ZetaFormat::v0226Hashline
1747        | ZetaFormat::V0304VariableEdit
1748        | ZetaFormat::V0304SeedNoEdits
1749        | ZetaFormat::V0331SeedCoderModelPy
1750        | ZetaFormat::V0306SeedMultiRegions
1751        | ZetaFormat::V0615HashRegions => {
1752            let (mut result, first_hunk_offset) = if empty_patch {
1753                (old_editable.clone(), None)
1754            } else {
1755                udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?
1756            };
1757
1758            if let Some(cursor) = cursor_offset {
1759                let hunk_start = if !empty_patch {
1760                    first_hunk_offset.unwrap_or(0)
1761                } else {
1762                    0
1763                };
1764                let offset = (hunk_start + cursor).min(result.len());
1765                result.insert_str(offset, CURSOR_MARKER);
1766            }
1767
1768            if !result.is_empty() && !result.ends_with('\n') {
1769                result.push('\n');
1770            }
1771
1772            if let Some(end_marker) = output_end_marker_for_format(format) {
1773                result.push_str(end_marker);
1774            }
1775
1776            Ok(result)
1777        }
1778    }
1779}
1780
1781/// Compute the cursor position within the new text after diff application.
1782fn cursor_in_new_text(
1783    cursor_offset: Option<usize>,
1784    first_hunk_offset: Option<usize>,
1785    new_text: &str,
1786) -> Option<usize> {
1787    cursor_offset.map(|cursor| {
1788        let hunk_start = first_hunk_offset.unwrap_or(0);
1789        (hunk_start + cursor).min(new_text.len())
1790    })
1791}
1792
1793#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1794pub struct ParsedOutput {
1795    /// Text that should replace the editable region
1796    pub new_editable_region: String,
1797    /// The byte range within `cursor_excerpt` that this replacement applies to
1798    pub range_in_excerpt: Range<usize>,
1799    /// Byte offset of the cursor marker within `new_editable_region`, if present
1800    pub cursor_offset_in_new_editable_region: Option<usize>,
1801}
1802
1803#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1804pub struct CursorPosition {
1805    pub path: String,
1806    pub row: usize,
1807    pub column: usize,
1808    pub offset: usize,
1809    pub editable_region_offset: usize,
1810}
1811
1812pub fn parsed_output_from_editable_region(
1813    range_in_excerpt: Range<usize>,
1814    mut new_editable_region: String,
1815) -> ParsedOutput {
1816    let cursor_offset_in_new_editable_region = new_editable_region.find(CURSOR_MARKER);
1817    if let Some(offset) = cursor_offset_in_new_editable_region {
1818        new_editable_region.replace_range(offset..offset + CURSOR_MARKER.len(), "");
1819    }
1820
1821    ParsedOutput {
1822        new_editable_region,
1823        range_in_excerpt,
1824        cursor_offset_in_new_editable_region,
1825    }
1826}
1827
1828/// Parse model output for the given zeta format
1829pub fn parse_zeta2_model_output(
1830    output: &str,
1831    format: ZetaFormat,
1832    prompt_inputs: &Zeta2PromptInput,
1833) -> Result<ParsedOutput> {
1834    let output = match output_end_marker_for_format(format) {
1835        Some(marker) => output.strip_suffix(marker).unwrap_or(output),
1836        None => output,
1837    };
1838
1839    let (context, editable_range_in_context, context_range, cursor_offset) =
1840        resolve_cursor_region(prompt_inputs, format);
1841    let context_start = context_range.start;
1842    let old_editable_region = &context[editable_range_in_context.clone()];
1843    let cursor_offset_in_editable = cursor_offset.saturating_sub(editable_range_in_context.start);
1844
1845    let (range_in_context, output) = match format {
1846        ZetaFormat::v0226Hashline => (
1847            editable_range_in_context,
1848            if hashline::output_has_edit_commands(output) {
1849                hashline::apply_edit_commands(old_editable_region, output)
1850            } else {
1851                output.to_string()
1852            },
1853        ),
1854        ZetaFormat::V0304VariableEdit => v0304_variable_edit::apply_variable_edit(context, output)?,
1855        ZetaFormat::V0304SeedNoEdits => (
1856            editable_range_in_context,
1857            if output.starts_with(seed_coder::NO_EDITS) {
1858                old_editable_region.to_string()
1859            } else {
1860                output.to_string()
1861            },
1862        ),
1863        ZetaFormat::V0306SeedMultiRegions => (
1864            editable_range_in_context,
1865            if output.starts_with(seed_coder::NO_EDITS) {
1866                old_editable_region.to_string()
1867            } else {
1868                multi_region::apply_marker_span(old_editable_region, output)?
1869            },
1870        ),
1871        ZetaFormat::V0316SeedMultiRegions => (
1872            editable_range_in_context,
1873            multi_region::apply_marker_span_v0316(old_editable_region, output)?,
1874        ),
1875        ZetaFormat::V0318SeedMultiRegions
1876        | ZetaFormat::V0420Diagnostics
1877        | ZetaFormat::V0608QwenMultiRegions => (
1878            editable_range_in_context,
1879            multi_region::apply_marker_span_v0318(old_editable_region, output)?,
1880        ),
1881        ZetaFormat::V0317SeedMultiRegions => (
1882            editable_range_in_context,
1883            multi_region::apply_marker_span_v0317(
1884                old_editable_region,
1885                output,
1886                Some(cursor_offset_in_editable),
1887            )?,
1888        ),
1889        ZetaFormat::V0327SingleFile => (
1890            editable_range_in_context,
1891            multi_region::apply_marker_span_v0318(old_editable_region, output)?,
1892        ),
1893        ZetaFormat::V0615HashRegions => {
1894            anyhow::bail!(
1895                "V0615HashRegions output addresses related-file context; use parse_zeta2_model_output_as_patch"
1896            )
1897        }
1898        _ => (editable_range_in_context, output.to_string()),
1899    };
1900
1901    let range_in_excerpt =
1902        range_in_context.start + context_start..range_in_context.end + context_start;
1903
1904    Ok(parsed_output_from_editable_region(range_in_excerpt, output))
1905}
1906
1907pub fn parse_zeta2_model_output_as_patch(
1908    output: &str,
1909    format: ZetaFormat,
1910    prompt_inputs: &Zeta2PromptInput,
1911) -> Result<String> {
1912    if format == ZetaFormat::V0615HashRegions {
1913        return hashed_regions::parse_output_as_patch(prompt_inputs, output, CURSOR_MARKER);
1914    }
1915
1916    let parsed = parse_zeta2_model_output(output, format, prompt_inputs)?;
1917    parsed_output_to_patch(prompt_inputs, parsed)
1918}
1919
1920pub fn parse_zeta3_model_output_as_patch(
1921    output: &str,
1922    format: ZetaFormat,
1923    input: &Zeta3PromptInput,
1924) -> Result<String> {
1925    match format {
1926        ZetaFormat::V0318SeedMultiRegions => {}
1927        _ => anyhow::bail!("unsupported Zeta3 output format: {format}"),
1928    }
1929
1930    let output = match output_end_marker_for_format(format) {
1931        Some(marker) => output.strip_suffix(marker).unwrap_or(output),
1932        None => output,
1933    };
1934    let (current_excerpt, cursor_offset_in_excerpt) = zeta3_current_file_excerpt(input)
1935        .ok_or_else(|| anyhow!("Zeta3 input is missing current-file editable context at cursor"))?;
1936    let (context, editable_range_in_context, context_range, _) = resolve_zeta3_cursor_region(
1937        current_excerpt.text.as_ref(),
1938        cursor_offset_in_excerpt,
1939        &input.syntax_ranges,
1940        format,
1941    );
1942    let old_editable_region = &context[editable_range_in_context.clone()];
1943    let range_in_excerpt = editable_range_in_context.start + context_range.start
1944        ..editable_range_in_context.end + context_range.start;
1945    let parsed = parsed_output_from_editable_region(
1946        range_in_excerpt,
1947        multi_region::apply_marker_span_v0318(old_editable_region, output)?,
1948    );
1949
1950    parsed_output_to_patch_for_excerpt(
1951        input.cursor_path.as_ref(),
1952        current_excerpt.text.as_ref(),
1953        current_excerpt.row_range.start,
1954        parsed,
1955    )
1956}
1957
1958pub fn cursor_position_from_parsed_output(
1959    prompt_inputs: &Zeta2PromptInput,
1960    parsed: &ParsedOutput,
1961) -> Option<CursorPosition> {
1962    let cursor_offset = parsed.cursor_offset_in_new_editable_region?;
1963    let editable_region_offset = parsed.range_in_excerpt.start;
1964    let excerpt = prompt_inputs.cursor_excerpt.as_ref();
1965
1966    let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
1967
1968    let new_editable_region = &parsed.new_editable_region;
1969    let prefix_end = cursor_offset.min(new_editable_region.len());
1970    let new_region_prefix = &new_editable_region[..prefix_end];
1971
1972    let row = editable_region_start_line + new_region_prefix.matches('\n').count();
1973
1974    let column = match new_region_prefix.rfind('\n') {
1975        Some(last_newline) => cursor_offset - last_newline - 1,
1976        None => {
1977            let content_prefix = &excerpt[..editable_region_offset];
1978            let content_column = match content_prefix.rfind('\n') {
1979                Some(last_newline) => editable_region_offset - last_newline - 1,
1980                None => editable_region_offset,
1981            };
1982            content_column + cursor_offset
1983        }
1984    };
1985
1986    Some(CursorPosition {
1987        path: prompt_inputs.cursor_path.to_string_lossy().into_owned(),
1988        row,
1989        column,
1990        offset: editable_region_offset + cursor_offset,
1991        editable_region_offset: cursor_offset,
1992    })
1993}
1994
1995pub fn parsed_output_to_patch(
1996    prompt_inputs: &Zeta2PromptInput,
1997    parsed: ParsedOutput,
1998) -> Result<String> {
1999    parsed_output_to_patch_for_excerpt(
2000        prompt_inputs.cursor_path.as_ref(),
2001        prompt_inputs.cursor_excerpt.as_ref(),
2002        0,
2003        parsed,
2004    )
2005}
2006
2007fn parsed_output_to_patch_for_excerpt(
2008    path: &Path,
2009    excerpt: &str,
2010    excerpt_start_row: u32,
2011    parsed: ParsedOutput,
2012) -> Result<String> {
2013    let range_in_excerpt = parsed.range_in_excerpt;
2014    let old_text = excerpt[range_in_excerpt.clone()].to_string();
2015    let mut new_text = parsed.new_editable_region;
2016
2017    let mut old_text_normalized = old_text;
2018    if !new_text.is_empty() && !new_text.ends_with('\n') {
2019        new_text.push('\n');
2020    }
2021    if !old_text_normalized.is_empty() && !old_text_normalized.ends_with('\n') {
2022        old_text_normalized.push('\n');
2023    }
2024
2025    let editable_region_offset = range_in_excerpt.start;
2026    let editable_region_start_line =
2027        excerpt_start_row + excerpt[..editable_region_offset].matches('\n').count() as u32;
2028    let editable_region_lines = old_text_normalized.lines().count() as u32;
2029
2030    let diff = udiff::unified_diff_with_context(
2031        &old_text_normalized,
2032        &new_text,
2033        editable_region_start_line,
2034        editable_region_start_line,
2035        editable_region_lines,
2036    );
2037
2038    let path = path.to_string_lossy().trim_start_matches('/').to_string();
2039    let formatted_diff = format!("--- a/{path}\n+++ b/{path}\n{diff}");
2040
2041    Ok(udiff::encode_cursor_in_patch(
2042        &formatted_diff,
2043        parsed.cursor_offset_in_new_editable_region,
2044    ))
2045}
2046
2047pub fn excerpt_range_for_format(
2048    format: ZetaFormat,
2049    ranges: &ExcerptRanges,
2050) -> (Range<usize>, Range<usize>) {
2051    excerpt_ranges_for_format(format, ranges)
2052}
2053
2054pub fn resolve_cursor_region(
2055    input: &Zeta2PromptInput,
2056    format: ZetaFormat,
2057) -> (&str, Range<usize>, Range<usize>, usize) {
2058    let (editable_range, context_range) = if format == ZetaFormat::V0327SingleFile {
2059        let (editable_tokens, _) = token_limits_for_format(format);
2060        let context_range = 0..input.cursor_excerpt.len();
2061        let editable_range = multi_region::compute_v0327_editable_range(
2062            &input.cursor_excerpt,
2063            input.cursor_offset_in_excerpt,
2064            editable_tokens,
2065        );
2066        (editable_range, context_range)
2067    } else if let Some(syntax_ranges) = &input.syntax_ranges {
2068        let (editable_tokens, context_tokens) = token_limits_for_format(format);
2069        compute_editable_and_context_ranges(
2070            &input.cursor_excerpt,
2071            input.cursor_offset_in_excerpt,
2072            syntax_ranges,
2073            editable_tokens,
2074            context_tokens,
2075        )
2076    } else {
2077        excerpt_range_for_format(format, &input.excerpt_ranges)
2078    };
2079
2080    adjust_cursor_region(
2081        &input.cursor_excerpt,
2082        input.cursor_offset_in_excerpt,
2083        editable_range,
2084        context_range,
2085    )
2086}
2087
2088fn adjust_cursor_region(
2089    cursor_excerpt: &str,
2090    cursor_offset: usize,
2091    editable_range: Range<usize>,
2092    context_range: Range<usize>,
2093) -> (&str, Range<usize>, Range<usize>, usize) {
2094    let context_start = context_range.start;
2095    let context_text = &cursor_excerpt[context_range.clone()];
2096    let adjusted_editable =
2097        (editable_range.start - context_start)..(editable_range.end - context_start);
2098    let adjusted_cursor = cursor_offset - context_start;
2099
2100    (
2101        context_text,
2102        adjusted_editable,
2103        context_range,
2104        adjusted_cursor,
2105    )
2106}
2107
2108pub fn get_prefill(input: &Zeta2PromptInput, format: ZetaFormat) -> String {
2109    let (context, editable_range, _, _) = resolve_cursor_region(input, format);
2110    get_prefill_for_format(format, context, &editable_range)
2111}
2112
2113pub fn format_edit_history_within_budget(
2114    events: &[Arc<Event>],
2115    file_marker: &str,
2116    edit_history_name: &str,
2117    max_tokens: usize,
2118    max_edit_event_count: usize,
2119) -> String {
2120    let header = format!("{}{}\n", file_marker, edit_history_name);
2121    let header_tokens = estimate_tokens(header.len());
2122    if header_tokens >= max_tokens {
2123        return String::new();
2124    }
2125
2126    let mut event_strings: Vec<String> = Vec::new();
2127    let mut total_tokens = header_tokens;
2128
2129    for event in events.iter().rev().take(max_edit_event_count) {
2130        let mut event_str = String::new();
2131        write_event(&mut event_str, event);
2132        let event_tokens = estimate_tokens(event_str.len());
2133
2134        if total_tokens + event_tokens > max_tokens {
2135            break;
2136        }
2137        total_tokens += event_tokens;
2138        event_strings.push(event_str);
2139    }
2140
2141    if event_strings.is_empty() {
2142        return String::new();
2143    }
2144
2145    let mut result = header;
2146    for event_str in event_strings.iter().rev() {
2147        result.push_str(event_str);
2148    }
2149    result
2150}
2151
2152fn excerpt_rendered_tokens(excerpt: &RelatedExcerpt, file_max_row: u32) -> usize {
2153    let needs_newline = !excerpt.text.ends_with('\n');
2154    let needs_ellipsis = excerpt.row_range.end < file_max_row;
2155    let len = excerpt.text.len()
2156        + if needs_newline { "\n".len() } else { 0 }
2157        + if needs_ellipsis { "...\n".len() } else { 0 };
2158    estimate_tokens(len)
2159}
2160
2161pub fn format_related_files_within_budget(
2162    related_files: &[RelatedFile],
2163    file_prefix: &str,
2164    file_suffix: &str,
2165    max_tokens: usize,
2166) -> String {
2167    struct ExcerptCandidate {
2168        file_ix: usize,
2169        excerpt_ix: usize,
2170        order: usize,
2171    }
2172
2173    let mut excerpt_candidates: Vec<ExcerptCandidate> = related_files
2174        .iter()
2175        .enumerate()
2176        .flat_map(|(file_ix, file)| {
2177            file.excerpts
2178                .iter()
2179                .enumerate()
2180                .map(move |(excerpt_ix, e)| ExcerptCandidate {
2181                    file_ix,
2182                    excerpt_ix,
2183                    order: e.order,
2184                })
2185        })
2186        .collect();
2187
2188    // Pre-compute file header strings and their token costs.
2189    let file_headers: Vec<String> = related_files
2190        .iter()
2191        .map(|file| {
2192            let path_str = file.path.to_string_lossy();
2193            format!("{}{}\n", file_prefix, path_str)
2194        })
2195        .collect();
2196
2197    // Sort the excerpts by their order and determine how many fit within the budget.
2198    let mut total_tokens = 0;
2199    let mut included_excerpt_count = 0_usize;
2200    let mut included_file_indices = vec![false; related_files.len()];
2201    excerpt_candidates.sort_by_key(|e| (e.order, e.file_ix, e.excerpt_ix));
2202    for candidate in &excerpt_candidates {
2203        let file = &related_files[candidate.file_ix];
2204        let excerpt = &file.excerpts[candidate.excerpt_ix];
2205        let file_already_included = included_file_indices[candidate.file_ix];
2206        let header_cost = if file_already_included {
2207            0
2208        } else {
2209            estimate_tokens(file_headers[candidate.file_ix].len() + file_suffix.len())
2210        };
2211        let excerpt_cost = excerpt_rendered_tokens(excerpt, file.max_row);
2212        if total_tokens + header_cost + excerpt_cost > max_tokens {
2213            break;
2214        }
2215        total_tokens += header_cost + excerpt_cost;
2216        if !file_already_included {
2217            included_file_indices[candidate.file_ix] = true;
2218        }
2219        included_excerpt_count += 1;
2220    }
2221
2222    excerpt_candidates.truncate(included_excerpt_count);
2223    excerpt_candidates.sort_unstable_by_key(|c| (c.file_ix, c.excerpt_ix));
2224
2225    // Render all of the files that fit within the token budget, in the original order.
2226    let mut result = String::new();
2227    let mut last_file_ix = None;
2228    for (candidate_ix, candidate) in excerpt_candidates.iter().enumerate() {
2229        if last_file_ix != Some(candidate.file_ix) {
2230            if last_file_ix.is_some() {
2231                result.push_str(file_suffix);
2232            }
2233            result.push_str(&file_headers[candidate.file_ix]);
2234            last_file_ix = Some(candidate.file_ix);
2235        }
2236        let file = &related_files[candidate.file_ix];
2237        let excerpt = &file.excerpts[candidate.excerpt_ix];
2238        result.push_str(&excerpt.text);
2239        if !result.ends_with('\n') {
2240            result.push('\n');
2241        }
2242        let next_excerpt_start = excerpt_candidates
2243            .get(candidate_ix + 1)
2244            .filter(|next| next.file_ix == candidate.file_ix)
2245            .map(|next| file.excerpts[next.excerpt_ix].row_range.start);
2246        if rows_omitted_after_excerpt(excerpt, next_excerpt_start, file.max_row) {
2247            result.push_str("...\n");
2248        }
2249    }
2250
2251    result
2252}
2253
2254/// Whether rows are omitted between this excerpt and the next rendered
2255/// excerpt of the same file (or the end of the file), in which case an
2256/// ellipsis line should be rendered.
2257pub fn rows_omitted_after_excerpt(
2258    excerpt: &RelatedExcerpt,
2259    next_excerpt_start: Option<u32>,
2260    file_max_row: u32,
2261) -> bool {
2262    match next_excerpt_start {
2263        Some(next_start) => excerpt.row_range.end < next_start,
2264        None => excerpt.row_range.end < file_max_row,
2265    }
2266}
2267
2268pub fn write_related_files(
2269    prompt: &mut String,
2270    related_files: &[RelatedFile],
2271) -> Vec<Range<usize>> {
2272    let mut ranges = Vec::new();
2273    for file in related_files {
2274        let start = prompt.len();
2275        let path_str = file.path.to_string_lossy();
2276        write!(prompt, "<|file_sep|>{}\n", path_str).ok();
2277        for (excerpt_ix, excerpt) in file.excerpts.iter().enumerate() {
2278            prompt.push_str(&excerpt.text);
2279            if !prompt.ends_with('\n') {
2280                prompt.push('\n');
2281            }
2282            let next_excerpt_start = file
2283                .excerpts
2284                .get(excerpt_ix + 1)
2285                .map(|next| next.row_range.start);
2286            if rows_omitted_after_excerpt(excerpt, next_excerpt_start, file.max_row) {
2287                prompt.push_str("...\n");
2288            }
2289        }
2290        let end = prompt.len();
2291        ranges.push(start..end);
2292    }
2293    ranges
2294}
2295
2296mod v0112_middle_at_end {
2297    use super::*;
2298
2299    pub fn special_tokens() -> &'static [&'static str] {
2300        &[
2301            "<|fim_prefix|>",
2302            "<|fim_suffix|>",
2303            "<|fim_middle|>",
2304            "<|file_sep|>",
2305            CURSOR_MARKER,
2306        ]
2307    }
2308
2309    pub fn write_cursor_excerpt_section(
2310        prompt: &mut String,
2311        path: &Path,
2312        context: &str,
2313        editable_range: &Range<usize>,
2314        cursor_offset: usize,
2315    ) {
2316        let path_str = path.to_string_lossy();
2317        write!(prompt, "<|file_sep|>{}\n", path_str).ok();
2318
2319        prompt.push_str("<|fim_prefix|>\n");
2320        prompt.push_str(&context[..editable_range.start]);
2321
2322        prompt.push_str("<|fim_suffix|>\n");
2323        prompt.push_str(&context[editable_range.end..]);
2324        if !prompt.ends_with('\n') {
2325            prompt.push('\n');
2326        }
2327
2328        prompt.push_str("<|fim_middle|>current\n");
2329        prompt.push_str(&context[editable_range.start..cursor_offset]);
2330        prompt.push_str(CURSOR_MARKER);
2331        prompt.push_str(&context[cursor_offset..editable_range.end]);
2332        if !prompt.ends_with('\n') {
2333            prompt.push('\n');
2334        }
2335
2336        prompt.push_str("<|fim_middle|>updated\n");
2337    }
2338}
2339
2340mod v0113_ordered {
2341    use super::*;
2342
2343    pub fn special_tokens() -> &'static [&'static str] {
2344        &[
2345            "<|fim_prefix|>",
2346            "<|fim_suffix|>",
2347            "<|fim_middle|>",
2348            "<|file_sep|>",
2349            CURSOR_MARKER,
2350        ]
2351    }
2352
2353    pub fn write_cursor_excerpt_section(
2354        prompt: &mut String,
2355        path: &Path,
2356        context: &str,
2357        editable_range: &Range<usize>,
2358        cursor_offset: usize,
2359    ) {
2360        let path_str = path.to_string_lossy();
2361        write!(prompt, "<|file_sep|>{}\n", path_str).ok();
2362
2363        prompt.push_str("<|fim_prefix|>\n");
2364        prompt.push_str(&context[..editable_range.start]);
2365        if !prompt.ends_with('\n') {
2366            prompt.push('\n');
2367        }
2368
2369        prompt.push_str("<|fim_middle|>current\n");
2370        prompt.push_str(&context[editable_range.start..cursor_offset]);
2371        prompt.push_str(CURSOR_MARKER);
2372        prompt.push_str(&context[cursor_offset..editable_range.end]);
2373        if !prompt.ends_with('\n') {
2374            prompt.push('\n');
2375        }
2376
2377        prompt.push_str("<|fim_suffix|>\n");
2378        prompt.push_str(&context[editable_range.end..]);
2379        if !prompt.ends_with('\n') {
2380            prompt.push('\n');
2381        }
2382
2383        prompt.push_str("<|fim_middle|>updated\n");
2384    }
2385}
2386
2387mod v0114180_editable_region {
2388    use super::*;
2389
2390    pub fn special_tokens() -> &'static [&'static str] {
2391        v0113_ordered::special_tokens()
2392    }
2393}
2394
2395pub mod v0120_git_merge_markers {
2396    //! A prompt that uses git-style merge conflict markers to represent the editable region.
2397    //!
2398    //! Example prompt:
2399    //!
2400    //! <|file_sep|>path/to/target_file.py
2401    //! <|fim_prefix|>
2402    //! code before editable region
2403    //! <|fim_suffix|>
2404    //! code after editable region
2405    //! <|fim_middle|>
2406    //! <<<<<<< CURRENT
2407    //! code that
2408    //! needs to<|user_cursor|>
2409    //! be rewritten
2410    //! =======
2411    //!
2412    //! Expected output (should be generated by the model):
2413    //!
2414    //! updated
2415    //! code with
2416    //! changes applied
2417    //! >>>>>>> UPDATED
2418
2419    use super::*;
2420
2421    pub const START_MARKER: &str = "<<<<<<< CURRENT\n";
2422    pub const SEPARATOR: &str = "=======\n";
2423    pub const END_MARKER: &str = ">>>>>>> UPDATED\n";
2424
2425    pub fn special_tokens() -> &'static [&'static str] {
2426        &[
2427            "<|fim_prefix|>",
2428            "<|fim_suffix|>",
2429            "<|fim_middle|>",
2430            "<|file_sep|>",
2431            START_MARKER,
2432            SEPARATOR,
2433            END_MARKER,
2434            CURSOR_MARKER,
2435        ]
2436    }
2437
2438    pub fn write_cursor_excerpt_section(
2439        prompt: &mut String,
2440        path: &Path,
2441        context: &str,
2442        editable_range: &Range<usize>,
2443        cursor_offset: usize,
2444    ) {
2445        let path_str = path.to_string_lossy();
2446        write!(prompt, "<|file_sep|>{}\n", path_str).ok();
2447
2448        prompt.push_str("<|fim_prefix|>");
2449        prompt.push_str(&context[..editable_range.start]);
2450
2451        prompt.push_str("<|fim_suffix|>");
2452        prompt.push_str(&context[editable_range.end..]);
2453        if !prompt.ends_with('\n') {
2454            prompt.push('\n');
2455        }
2456
2457        prompt.push_str("<|fim_middle|>");
2458        prompt.push_str(START_MARKER);
2459        prompt.push_str(&context[editable_range.start..cursor_offset]);
2460        prompt.push_str(CURSOR_MARKER);
2461        prompt.push_str(&context[cursor_offset..editable_range.end]);
2462        if !prompt.ends_with('\n') {
2463            prompt.push('\n');
2464        }
2465        prompt.push_str(SEPARATOR);
2466    }
2467}
2468
2469pub mod v0131_git_merge_markers_prefix {
2470    //! A prompt that uses git-style merge conflict markers to represent the editable region.
2471    //!
2472    //! Example prompt:
2473    //!
2474    //! <|file_sep|>path/to/target_file.py
2475    //! <|fim_prefix|>
2476    //! code before editable region
2477    //! <<<<<<< CURRENT
2478    //! code that
2479    //! needs to<|user_cursor|>
2480    //! be rewritten
2481    //! =======
2482    //! <|fim_suffix|>
2483    //! code after editable region
2484    //! <|fim_middle|>
2485    //!
2486    //! Expected output (should be generated by the model):
2487    //!
2488    //! updated
2489    //! code with
2490    //! changes applied
2491    //! >>>>>>> UPDATED
2492
2493    use super::*;
2494
2495    pub const START_MARKER: &str = "<<<<<<< CURRENT\n";
2496    pub const SEPARATOR: &str = "=======\n";
2497    pub const END_MARKER: &str = ">>>>>>> UPDATED\n";
2498
2499    pub fn special_tokens() -> &'static [&'static str] {
2500        &[
2501            "<|fim_prefix|>",
2502            "<|fim_suffix|>",
2503            "<|fim_middle|>",
2504            "<|file_sep|>",
2505            START_MARKER,
2506            SEPARATOR,
2507            END_MARKER,
2508            CURSOR_MARKER,
2509        ]
2510    }
2511
2512    pub fn write_cursor_excerpt_section(
2513        prompt: &mut String,
2514        path: &Path,
2515        context: &str,
2516        editable_range: &Range<usize>,
2517        cursor_offset: usize,
2518    ) {
2519        let path_str = path.to_string_lossy();
2520        write!(prompt, "<|file_sep|>{}\n", path_str).ok();
2521
2522        prompt.push_str("<|fim_prefix|>");
2523        prompt.push_str(&context[..editable_range.start]);
2524        prompt.push_str(START_MARKER);
2525        prompt.push_str(&context[editable_range.start..cursor_offset]);
2526        prompt.push_str(CURSOR_MARKER);
2527        prompt.push_str(&context[cursor_offset..editable_range.end]);
2528        if !prompt.ends_with('\n') {
2529            prompt.push('\n');
2530        }
2531        prompt.push_str(SEPARATOR);
2532
2533        prompt.push_str("<|fim_suffix|>");
2534        prompt.push_str(&context[editable_range.end..]);
2535        if !prompt.ends_with('\n') {
2536            prompt.push('\n');
2537        }
2538
2539        prompt.push_str("<|fim_middle|>");
2540    }
2541}
2542
2543pub mod v0211_prefill {
2544    use super::*;
2545
2546    pub fn special_tokens() -> &'static [&'static str] {
2547        v0131_git_merge_markers_prefix::special_tokens()
2548    }
2549
2550    pub fn get_prefill(context: &str, editable_range: &Range<usize>) -> String {
2551        let editable_region = &context[editable_range.start..editable_range.end];
2552
2553        let prefill_len = (editable_region.len() as f64 * PREFILL_RATIO) as usize;
2554        let prefill_len = editable_region.floor_char_boundary(prefill_len);
2555
2556        // Find a token boundary to avoid splitting tokens in the prefill.
2557        // In Qwen2.5-Coder, \n is always the END of a token (e.g. `;\n`,
2558        // ` {\n`), and \n\n / \n\n\n are single tokens, so we must include
2559        // the \n and consume any consecutive \n characters after it.
2560        let prefill = &editable_region[..prefill_len];
2561        match prefill.rfind('\n') {
2562            Some(pos) => {
2563                let mut end = pos + 1;
2564                while end < editable_region.len()
2565                    && editable_region.as_bytes().get(end) == Some(&b'\n')
2566                {
2567                    end += 1;
2568                }
2569                editable_region[..end].to_string()
2570            }
2571            // No newline found. Fall back to splitting before the last space
2572            // (word-level boundary)
2573            None => match prefill.rfind(' ') {
2574                Some(pos) => prefill[..pos].to_string(),
2575                None => prefill.to_string(),
2576            },
2577        }
2578    }
2579}
2580
2581pub mod hashline {
2582
2583    use std::fmt::Display;
2584
2585    pub const END_MARKER: &str = "<|fim_middle|>updated";
2586    pub const START_MARKER: &str = "<|fim_middle|>current";
2587
2588    use super::*;
2589
2590    const SET_COMMAND_MARKER: &str = "<|set|>";
2591    const INSERT_COMMAND_MARKER: &str = "<|insert|>";
2592    pub const NO_EDITS_COMMAND_MARKER: &str = "<|no_edits|>";
2593
2594    pub fn special_tokens() -> &'static [&'static str] {
2595        return &[
2596            SET_COMMAND_MARKER,
2597            "<|set_range|>",
2598            INSERT_COMMAND_MARKER,
2599            NO_EDITS_COMMAND_MARKER,
2600            CURSOR_MARKER,
2601            "<|file_sep|>",
2602            "<|fim_prefix|>",
2603            "<|fim_suffix|>",
2604            "<|fim_middle|>",
2605        ];
2606    }
2607
2608    /// A parsed line reference like `3:c3` (line index 3 with hash 0xc3).
2609    #[derive(Debug, Clone, PartialEq, Eq)]
2610    struct LineRef {
2611        index: usize,
2612        hash: u8,
2613    }
2614
2615    impl Display for LineRef {
2616        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2617            write!(f, "{}:{:02x}", self.index, self.hash)
2618        }
2619    }
2620
2621    pub fn hash_line(line: &[u8]) -> u8 {
2622        let mut h: u8 = 0;
2623        for &byte in line {
2624            h = h.wrapping_add(byte);
2625        }
2626        return h;
2627    }
2628
2629    /// Write the hashline-encoded editable region into `out`. Each line of
2630    /// `editable_text` is prefixed with `{line_index}:{hash}|` and the cursor
2631    /// marker is inserted at `cursor_offset_in_editable` (byte offset relative
2632    /// to the start of `editable_text`).
2633    pub fn write_hashline_editable_region(
2634        out: &mut String,
2635        editable_text: &str,
2636        cursor_offset_in_editable: usize,
2637    ) {
2638        let mut offset = 0;
2639        for (i, line) in editable_text.lines().enumerate() {
2640            let (head, cursor, tail) = if cursor_offset_in_editable > offset
2641                && cursor_offset_in_editable < offset + line.len()
2642            {
2643                (
2644                    &line[..cursor_offset_in_editable - offset],
2645                    CURSOR_MARKER,
2646                    &line[cursor_offset_in_editable - offset..],
2647                )
2648            } else {
2649                (line, "", "")
2650            };
2651            write!(
2652                out,
2653                "\n{}|{head}{cursor}{tail}",
2654                LineRef {
2655                    index: i,
2656                    hash: hash_line(line.as_bytes())
2657                }
2658            )
2659            .unwrap();
2660            offset += line.len() + 1;
2661        }
2662    }
2663
2664    pub fn write_cursor_excerpt_section(
2665        prompt: &mut String,
2666        path: &Path,
2667        context: &str,
2668        editable_range: &Range<usize>,
2669        cursor_offset: usize,
2670    ) {
2671        let path_str = path.to_string_lossy();
2672        write!(prompt, "<|file_sep|>{}\n", path_str).ok();
2673
2674        prompt.push_str("<|fim_prefix|>\n");
2675        prompt.push_str(&context[..editable_range.start]);
2676        prompt.push_str(START_MARKER);
2677
2678        let cursor_offset_in_editable = cursor_offset.saturating_sub(editable_range.start);
2679        let editable_region = &context[editable_range.clone()];
2680        write_hashline_editable_region(prompt, editable_region, cursor_offset_in_editable);
2681
2682        if !prompt.ends_with('\n') {
2683            prompt.push('\n');
2684        }
2685
2686        prompt.push_str("<|fim_suffix|>\n");
2687        prompt.push_str(&context[editable_range.end..]);
2688        if !prompt.ends_with('\n') {
2689            prompt.push('\n');
2690        }
2691
2692        prompt.push_str(END_MARKER);
2693        prompt.push('\n');
2694    }
2695
2696    /// A single edit command parsed from the model output.
2697    #[derive(Debug)]
2698    enum EditCommand<'a> {
2699        /// Replace a range of lines (inclusive on both ends). Single-line set is
2700        /// represented by `start == end`.
2701        Set {
2702            start: LineRef,
2703            end: LineRef,
2704            content: &'a str,
2705        },
2706        /// Insert new lines after the given line, or before the first line if
2707        /// `after` is `None`.
2708        Insert {
2709            after: Option<LineRef>,
2710            content: &'a str,
2711        },
2712    }
2713
2714    /// Parse a line reference like `3:c3` into a `LineRef`.
2715    fn parse_line_ref(s: &str) -> Option<LineRef> {
2716        let (idx_str, hash_str) = s.split_once(':')?;
2717        let index = idx_str.parse::<usize>().ok()?;
2718        let hash = u8::from_str_radix(hash_str, 16).ok()?;
2719        Some(LineRef { index, hash })
2720    }
2721
2722    /// Parse the model output into a list of `EditCommand`s.
2723    fn parse_edit_commands(model_output: &str) -> Vec<EditCommand<'_>> {
2724        let mut commands = Vec::new();
2725        let mut offset = 0usize;
2726
2727        while offset < model_output.len() {
2728            let next_nl = model_output[offset..]
2729                .find('\n')
2730                .map(|i| offset + i)
2731                .unwrap_or(model_output.len());
2732            let line = &model_output[offset..next_nl];
2733            let line_end = if next_nl < model_output.len() {
2734                next_nl + 1
2735            } else {
2736                next_nl
2737            };
2738
2739            let trimmed = line.trim();
2740            let (is_set, specifier) = if let Some(spec) = trimmed.strip_prefix(SET_COMMAND_MARKER) {
2741                (true, spec)
2742            } else if let Some(spec) = trimmed.strip_prefix(INSERT_COMMAND_MARKER) {
2743                (false, spec)
2744            } else {
2745                offset = line_end;
2746                continue;
2747            };
2748
2749            let mut content_end = line_end;
2750            let mut scan = line_end;
2751
2752            while scan < model_output.len() {
2753                let body_nl = model_output[scan..]
2754                    .find('\n')
2755                    .map(|i| scan + i)
2756                    .unwrap_or(model_output.len());
2757                let body_line = &model_output[scan..body_nl];
2758                if body_line.trim().starts_with(SET_COMMAND_MARKER)
2759                    || body_line.trim().starts_with(INSERT_COMMAND_MARKER)
2760                {
2761                    break;
2762                }
2763                scan = if body_nl < model_output.len() {
2764                    body_nl + 1
2765                } else {
2766                    body_nl
2767                };
2768                content_end = scan;
2769            }
2770
2771            let content = &model_output[line_end..content_end];
2772
2773            if is_set {
2774                if let Some((start_str, end_str)) = specifier.split_once('-') {
2775                    if let (Some(start), Some(end)) =
2776                        (parse_line_ref(start_str), parse_line_ref(end_str))
2777                    {
2778                        commands.push(EditCommand::Set {
2779                            start,
2780                            end,
2781                            content,
2782                        });
2783                    }
2784                } else if let Some(target) = parse_line_ref(specifier) {
2785                    commands.push(EditCommand::Set {
2786                        start: target.clone(),
2787                        end: target,
2788                        content,
2789                    });
2790                }
2791            } else {
2792                let after = parse_line_ref(specifier);
2793                commands.push(EditCommand::Insert { after, content });
2794            }
2795
2796            offset = scan;
2797        }
2798
2799        commands
2800    }
2801
2802    /// Returns `true` if the model output contains `<|set|>` or `<|insert|>` commands
2803    /// (as opposed to being a plain full-replacement output).
2804    /// Strip the `{line_num}:{hash}|` prefixes from each line of a hashline-encoded
2805    /// editable region, returning the plain text content.
2806    pub fn strip_hashline_prefixes(region: &str) -> String {
2807        let mut decoded: String = region
2808            .lines()
2809            .map(|line| line.find('|').map_or(line, |pos| &line[pos + 1..]))
2810            .collect::<Vec<_>>()
2811            .join("\n");
2812        if region.ends_with('\n') {
2813            decoded.push('\n');
2814        }
2815        decoded
2816    }
2817
2818    pub fn output_has_edit_commands(model_output: &str) -> bool {
2819        model_output.contains(SET_COMMAND_MARKER)
2820            || model_output.contains(INSERT_COMMAND_MARKER)
2821            || model_output.contains(NO_EDITS_COMMAND_MARKER)
2822    }
2823
2824    /// Apply `<|set|>` and `<|insert|>` edit commands from the model output to the
2825    /// original editable region text.
2826    ///
2827    /// `editable_region` is the original text of the editable region (without hash
2828    /// prefixes). `model_output` is the raw model response containing edit commands.
2829    ///
2830    /// Returns the full replacement text for the editable region.
2831    pub fn apply_edit_commands(editable_region: &str, model_output: &str) -> String {
2832        if model_output
2833            .trim_start()
2834            .starts_with(NO_EDITS_COMMAND_MARKER)
2835        {
2836            return editable_region.to_string();
2837        }
2838
2839        let original_lines: Vec<&str> = editable_region.lines().collect();
2840        let old_hashes: Vec<u8> = original_lines
2841            .iter()
2842            .map(|line| hash_line(line.as_bytes()))
2843            .collect();
2844
2845        let commands = parse_edit_commands(model_output);
2846
2847        // For set operations: indexed by start line → Some((end line index, content))
2848        // For insert operations: indexed by line index → vec of content to insert after
2849        // Insert-before-first is tracked separately.
2850        let mut set_ops: Vec<Option<(usize, &str)>> = vec![None; original_lines.len()];
2851        let mut insert_before_first: Vec<&str> = Vec::new();
2852        let mut insert_after: Vec<Vec<&str>> = vec![Vec::new(); original_lines.len()];
2853
2854        for command in &commands {
2855            match command {
2856                EditCommand::Set {
2857                    start,
2858                    end,
2859                    content,
2860                } => {
2861                    if start.index < old_hashes.len()
2862                        && end.index < old_hashes.len()
2863                        && start.index <= end.index
2864                        && old_hashes[start.index] == start.hash
2865                        && old_hashes[end.index] == end.hash
2866                    {
2867                        set_ops[start.index] = Some((end.index, *content));
2868                    }
2869                }
2870                EditCommand::Insert { after, content } => match after {
2871                    None => insert_before_first.push(*content),
2872                    Some(line_ref) => {
2873                        if line_ref.index < old_hashes.len()
2874                            && old_hashes[line_ref.index] == line_ref.hash
2875                        {
2876                            insert_after[line_ref.index].push(*content);
2877                        }
2878                    }
2879                },
2880            }
2881        }
2882
2883        let mut result = String::new();
2884
2885        // Emit any insertions before the first line
2886        for content in &insert_before_first {
2887            result.push_str(content);
2888            if !content.ends_with('\n') {
2889                result.push('\n');
2890            }
2891        }
2892
2893        let mut i = 0;
2894        while i < original_lines.len() {
2895            if let Some((end_index, replacement)) = set_ops[i].as_ref() {
2896                // Replace lines i..=end_index with the replacement content
2897                result.push_str(replacement);
2898                if !replacement.is_empty() && !replacement.ends_with('\n') {
2899                    result.push('\n');
2900                }
2901                // Emit any insertions after the end of this set range
2902                if *end_index < insert_after.len() {
2903                    for content in &insert_after[*end_index] {
2904                        result.push_str(content);
2905                        if !content.ends_with('\n') {
2906                            result.push('\n');
2907                        }
2908                    }
2909                }
2910                i = end_index + 1;
2911            } else {
2912                // Keep the original line
2913                result.push_str(original_lines[i]);
2914                result.push('\n');
2915                // Emit any insertions after this line
2916                for content in &insert_after[i] {
2917                    result.push_str(content);
2918                    if !content.ends_with('\n') {
2919                        result.push('\n');
2920                    }
2921                }
2922                i += 1;
2923            }
2924        }
2925
2926        // Preserve trailing newline behavior: if the original ended with a
2927        // newline the result already has one; if it didn't, trim the extra one
2928        // we added.
2929        if !editable_region.ends_with('\n') && result.ends_with('\n') {
2930            result.pop();
2931        }
2932
2933        result
2934    }
2935
2936    /// Convert a unified diff patch into hashline edit commands.
2937    ///
2938    /// Parses the unified diff `patch` directly to determine which lines of
2939    /// `old_text` are deleted/replaced and what new lines are added, then emits
2940    /// `<|set|>` and `<|insert|>` edit commands referencing old lines by their
2941    /// `{index}:{hash}` identifiers.
2942    ///
2943    /// `cursor_offset` is an optional byte offset into the first hunk's new
2944    /// text (context + additions) where the cursor marker should be placed.
2945    pub fn patch_to_edit_commands(
2946        old_text: &str,
2947        patch: &str,
2948        cursor_offset: Option<usize>,
2949    ) -> Result<String> {
2950        let old_lines: Vec<&str> = old_text.lines().collect();
2951        let old_hashes: Vec<u8> = old_lines
2952            .iter()
2953            .map(|line| hash_line(line.as_bytes()))
2954            .collect();
2955
2956        let mut result = String::new();
2957        let mut first_hunk = true;
2958
2959        struct Hunk<'a> {
2960            line_range: Range<usize>,
2961            new_text_lines: Vec<&'a str>,
2962            cursor_line_offset_in_new_text: Option<(usize, usize)>,
2963        }
2964
2965        // Parse the patch line by line. We only care about hunk headers,
2966        // context, deletions, and additions.
2967        let mut old_line_index: usize = 0;
2968        let mut current_hunk: Option<Hunk> = None;
2969        // Byte offset tracking within the hunk's new text for cursor placement.
2970        let mut new_text_byte_offset: usize = 0;
2971        // The line index of the last old line seen before/in the current hunk
2972        // (used for insert-after reference).
2973        let mut last_old_line_before_hunk: Option<usize> = None;
2974
2975        fn flush_hunk(
2976            hunk: Hunk,
2977            last_old_line: Option<usize>,
2978            result: &mut String,
2979            old_hashes: &[u8],
2980        ) {
2981            if hunk.line_range.is_empty() {
2982                // Pure insertion — reference the old line to insert after when in bounds.
2983                if let Some(after) = last_old_line
2984                    && let Some(&hash) = old_hashes.get(after)
2985                {
2986                    write!(
2987                        result,
2988                        "{INSERT_COMMAND_MARKER}{}\n",
2989                        LineRef { index: after, hash }
2990                    )
2991                    .unwrap();
2992                } else {
2993                    result.push_str(INSERT_COMMAND_MARKER);
2994                    result.push('\n');
2995                }
2996            } else {
2997                let start = hunk.line_range.start;
2998                let end_exclusive = hunk.line_range.end;
2999                let deleted_line_count = end_exclusive.saturating_sub(start);
3000
3001                if deleted_line_count == 1 {
3002                    if let Some(&hash) = old_hashes.get(start) {
3003                        write!(
3004                            result,
3005                            "{SET_COMMAND_MARKER}{}\n",
3006                            LineRef { index: start, hash }
3007                        )
3008                        .unwrap();
3009                    } else {
3010                        result.push_str(SET_COMMAND_MARKER);
3011                        result.push('\n');
3012                    }
3013                } else {
3014                    let end_inclusive = end_exclusive - 1;
3015                    match (
3016                        old_hashes.get(start).copied(),
3017                        old_hashes.get(end_inclusive).copied(),
3018                    ) {
3019                        (Some(start_hash), Some(end_hash)) => {
3020                            write!(
3021                                result,
3022                                "{SET_COMMAND_MARKER}{}-{}\n",
3023                                LineRef {
3024                                    index: start,
3025                                    hash: start_hash
3026                                },
3027                                LineRef {
3028                                    index: end_inclusive,
3029                                    hash: end_hash
3030                                }
3031                            )
3032                            .unwrap();
3033                        }
3034                        _ => {
3035                            result.push_str(SET_COMMAND_MARKER);
3036                            result.push('\n');
3037                        }
3038                    }
3039                }
3040            }
3041            for (line_offset, line) in hunk.new_text_lines.iter().enumerate() {
3042                if let Some((cursor_line_offset, char_offset)) = hunk.cursor_line_offset_in_new_text
3043                    && line_offset == cursor_line_offset
3044                {
3045                    result.push_str(&line[..char_offset]);
3046                    result.push_str(CURSOR_MARKER);
3047                    result.push_str(&line[char_offset..]);
3048                    continue;
3049                }
3050
3051                result.push_str(line);
3052            }
3053        }
3054
3055        for raw_line in patch.split_inclusive('\n') {
3056            if raw_line.starts_with("@@") {
3057                // Flush any pending change hunk from a previous patch hunk.
3058                if let Some(hunk) = current_hunk.take() {
3059                    flush_hunk(hunk, last_old_line_before_hunk, &mut result, &old_hashes);
3060                }
3061
3062                // Parse hunk header: @@ -old_start[,old_count] +new_start[,new_count] @@
3063                // We intentionally do not trust old_start as a direct local index into `old_text`,
3064                // because some patches are produced against a larger file region and carry
3065                // non-local line numbers. We keep indexing local by advancing from parsed patch lines.
3066                if first_hunk {
3067                    new_text_byte_offset = 0;
3068                    first_hunk = false;
3069                }
3070                continue;
3071            }
3072
3073            if raw_line.starts_with("---") || raw_line.starts_with("+++") {
3074                continue;
3075            }
3076            if raw_line.starts_with("\\ No newline") {
3077                continue;
3078            }
3079
3080            if raw_line.starts_with('-') {
3081                // Extend or start a change hunk with this deleted old line.
3082                match &mut current_hunk {
3083                    Some(Hunk {
3084                        line_range: range, ..
3085                    }) => range.end = old_line_index + 1,
3086                    None => {
3087                        current_hunk = Some(Hunk {
3088                            line_range: old_line_index..old_line_index + 1,
3089                            new_text_lines: Vec::new(),
3090                            cursor_line_offset_in_new_text: None,
3091                        });
3092                    }
3093                }
3094                old_line_index += 1;
3095            } else if let Some(added_content) = raw_line.strip_prefix('+') {
3096                // Place cursor marker if cursor_offset falls within this line.
3097                let mut cursor_line_offset = None;
3098                if let Some(cursor_off) = cursor_offset
3099                    && (first_hunk
3100                        || cursor_off >= new_text_byte_offset
3101                            && cursor_off <= new_text_byte_offset + added_content.len())
3102                {
3103                    let line_offset = added_content.floor_char_boundary(
3104                        cursor_off
3105                            .saturating_sub(new_text_byte_offset)
3106                            .min(added_content.len()),
3107                    );
3108                    cursor_line_offset = Some(line_offset);
3109                }
3110
3111                new_text_byte_offset += added_content.len();
3112
3113                let hunk = current_hunk.get_or_insert(Hunk {
3114                    line_range: old_line_index..old_line_index,
3115                    new_text_lines: vec![],
3116                    cursor_line_offset_in_new_text: None,
3117                });
3118                hunk.new_text_lines.push(added_content);
3119                hunk.cursor_line_offset_in_new_text = cursor_line_offset
3120                    .map(|offset_in_line| (hunk.new_text_lines.len() - 1, offset_in_line));
3121            } else {
3122                // Context line (starts with ' ' or is empty).
3123                if let Some(hunk) = current_hunk.take() {
3124                    flush_hunk(hunk, last_old_line_before_hunk, &mut result, &old_hashes);
3125                }
3126                last_old_line_before_hunk = Some(old_line_index);
3127                old_line_index += 1;
3128                let content = raw_line.strip_prefix(' ').unwrap_or(raw_line);
3129                new_text_byte_offset += content.len();
3130            }
3131        }
3132
3133        // Flush final group.
3134        if let Some(hunk) = current_hunk.take() {
3135            flush_hunk(hunk, last_old_line_before_hunk, &mut result, &old_hashes);
3136        }
3137
3138        // Trim a single trailing newline.
3139        if result.ends_with('\n') {
3140            result.pop();
3141        }
3142
3143        if result.is_empty() {
3144            return Ok(NO_EDITS_COMMAND_MARKER.to_string());
3145        }
3146
3147        Ok(result)
3148    }
3149
3150    #[cfg(test)]
3151    mod tests {
3152        use super::*;
3153        use indoc::indoc;
3154
3155        #[test]
3156        fn test_format_cursor_region() {
3157            struct Case {
3158                name: &'static str,
3159                context: &'static str,
3160                editable_range: Range<usize>,
3161                cursor_offset: usize,
3162                expected: &'static str,
3163            }
3164
3165            let cases = [
3166                Case {
3167                    name: "basic_cursor_placement",
3168                    context: "hello world\n",
3169                    editable_range: 0..12,
3170                    cursor_offset: 5,
3171                    expected: indoc! {"
3172                    <|file_sep|>test.rs
3173                    <|fim_prefix|>
3174                    <|fim_middle|>current
3175                    0:5c|hello<|user_cursor|> world
3176                    <|fim_suffix|>
3177                    <|fim_middle|>updated
3178                    "},
3179                },
3180                Case {
3181                    name: "multiline_cursor_on_second_line",
3182                    context: "aaa\nbbb\nccc\n",
3183                    editable_range: 0..12,
3184                    cursor_offset: 5, // byte 5 → 1 byte into "bbb"
3185                    expected: indoc! {"
3186                    <|file_sep|>test.rs
3187                    <|fim_prefix|>
3188                    <|fim_middle|>current
3189                    0:23|aaa
3190                    1:26|b<|user_cursor|>bb
3191                    2:29|ccc
3192                    <|fim_suffix|>
3193                    <|fim_middle|>updated
3194                    "},
3195                },
3196                Case {
3197                    name: "no_trailing_newline_in_context",
3198                    context: "line1\nline2",
3199                    editable_range: 0..11,
3200                    cursor_offset: 3,
3201                    expected: indoc! {"
3202                    <|file_sep|>test.rs
3203                    <|fim_prefix|>
3204                    <|fim_middle|>current
3205                    0:d9|lin<|user_cursor|>e1
3206                    1:da|line2
3207                    <|fim_suffix|>
3208                    <|fim_middle|>updated
3209                    "},
3210                },
3211                Case {
3212                    name: "leading_newline_in_editable_region",
3213                    context: "\nabc\n",
3214                    editable_range: 0..5,
3215                    cursor_offset: 2, // byte 2 = 'a' in "abc" (after leading \n)
3216                    expected: indoc! {"
3217                    <|file_sep|>test.rs
3218                    <|fim_prefix|>
3219                    <|fim_middle|>current
3220                    0:00|
3221                    1:26|a<|user_cursor|>bc
3222                    <|fim_suffix|>
3223                    <|fim_middle|>updated
3224                    "},
3225                },
3226                Case {
3227                    name: "with_suffix",
3228                    context: "abc\ndef",
3229                    editable_range: 0..4, // editable region = "abc\n", suffix = "def"
3230                    cursor_offset: 2,
3231                    expected: indoc! {"
3232                    <|file_sep|>test.rs
3233                    <|fim_prefix|>
3234                    <|fim_middle|>current
3235                    0:26|ab<|user_cursor|>c
3236                    <|fim_suffix|>
3237                    def
3238                    <|fim_middle|>updated
3239                    "},
3240                },
3241                Case {
3242                    name: "unicode_two_byte_chars",
3243                    context: "héllo\n",
3244                    editable_range: 0..7,
3245                    cursor_offset: 3, // byte 3 = after "hé" (h=1 byte, é=2 bytes), before "llo"
3246                    expected: indoc! {"
3247                    <|file_sep|>test.rs
3248                    <|fim_prefix|>
3249                    <|fim_middle|>current
3250                    0:1b|hé<|user_cursor|>llo
3251                    <|fim_suffix|>
3252                    <|fim_middle|>updated
3253                    "},
3254                },
3255                Case {
3256                    name: "unicode_three_byte_chars",
3257                    context: "日本語\n",
3258                    editable_range: 0..10,
3259                    cursor_offset: 6, // byte 6 = after "日本" (3+3 bytes), before "語"
3260                    expected: indoc! {"
3261                    <|file_sep|>test.rs
3262                    <|fim_prefix|>
3263                    <|fim_middle|>current
3264                    0:80|日本<|user_cursor|>語
3265                    <|fim_suffix|>
3266                    <|fim_middle|>updated
3267                    "},
3268                },
3269                Case {
3270                    name: "unicode_four_byte_chars",
3271                    context: "a🌍b\n",
3272                    editable_range: 0..7,
3273                    cursor_offset: 5, // byte 5 = after "a🌍" (1+4 bytes), before "b"
3274                    expected: indoc! {"
3275                    <|file_sep|>test.rs
3276                    <|fim_prefix|>
3277                    <|fim_middle|>current
3278                    0:6b|a🌍<|user_cursor|>b
3279                    <|fim_suffix|>
3280                    <|fim_middle|>updated
3281                    "},
3282                },
3283                Case {
3284                    name: "cursor_at_start_of_region_not_placed",
3285                    context: "abc\n",
3286                    editable_range: 0..4,
3287                    cursor_offset: 0, // cursor_offset(0) > offset(0) is false → cursor not placed
3288                    expected: indoc! {"
3289                    <|file_sep|>test.rs
3290                    <|fim_prefix|>
3291                    <|fim_middle|>current
3292                    0:26|abc
3293                    <|fim_suffix|>
3294                    <|fim_middle|>updated
3295                    "},
3296                },
3297                Case {
3298                    name: "cursor_at_end_of_line_not_placed",
3299                    context: "abc\ndef\n",
3300                    editable_range: 0..8,
3301                    cursor_offset: 3, // byte 3 = the \n after "abc" → falls between lines, not placed
3302                    expected: indoc! {"
3303                    <|file_sep|>test.rs
3304                    <|fim_prefix|>
3305                    <|fim_middle|>current
3306                    0:26|abc
3307                    1:2f|def
3308                    <|fim_suffix|>
3309                    <|fim_middle|>updated
3310                    "},
3311                },
3312                Case {
3313                    name: "cursor_offset_relative_to_context_not_editable_region",
3314                    // cursor_offset is relative to `context`, so when editable_range.start > 0,
3315                    // write_cursor_excerpt_section must subtract it before comparing against
3316                    // per-line offsets within the editable region.
3317                    context: "pre\naaa\nbbb\nsuf\n",
3318                    editable_range: 4..12, // editable region = "aaa\nbbb\n"
3319                    cursor_offset: 9,      // byte 9 in context = second 'b' in "bbb"
3320                    expected: indoc! {"
3321                    <|file_sep|>test.rs
3322                    <|fim_prefix|>
3323                    pre
3324                    <|fim_middle|>current
3325                    0:23|aaa
3326                    1:26|b<|user_cursor|>bb
3327                    <|fim_suffix|>
3328                    suf
3329                    <|fim_middle|>updated
3330                    "},
3331                },
3332            ];
3333
3334            for case in &cases {
3335                let mut prompt = String::new();
3336                hashline::write_cursor_excerpt_section(
3337                    &mut prompt,
3338                    Path::new("test.rs"),
3339                    case.context,
3340                    &case.editable_range,
3341                    case.cursor_offset,
3342                );
3343                assert_eq!(prompt, case.expected, "failed case: {}", case.name);
3344            }
3345        }
3346
3347        #[test]
3348        fn test_apply_edit_commands() {
3349            struct Case {
3350                name: &'static str,
3351                original: &'static str,
3352                model_output: &'static str,
3353                expected: &'static str,
3354            }
3355
3356            let cases = vec![
3357                Case {
3358                    name: "set_single_line",
3359                    original: indoc! {"
3360                    let mut total = 0;
3361                    for product in products {
3362                        total += ;
3363                    }
3364                    total
3365                "},
3366                    model_output: indoc! {"
3367                    <|set|>2:87
3368                        total += product.price;
3369                "},
3370                    expected: indoc! {"
3371                    let mut total = 0;
3372                    for product in products {
3373                        total += product.price;
3374                    }
3375                    total
3376                "},
3377                },
3378                Case {
3379                    name: "set_range",
3380                    original: indoc! {"
3381                    fn foo() {
3382                        let x = 1;
3383                        let y = 2;
3384                        let z = 3;
3385                    }
3386                "},
3387                    model_output: indoc! {"
3388                    <|set|>1:46-3:4a
3389                        let sum = 6;
3390                "},
3391                    expected: indoc! {"
3392                    fn foo() {
3393                        let sum = 6;
3394                    }
3395                "},
3396                },
3397                Case {
3398                    name: "insert_after_line",
3399                    original: indoc! {"
3400                    fn main() {
3401                        let x = 1;
3402                    }
3403                "},
3404                    model_output: indoc! {"
3405                    <|insert|>1:46
3406                        let y = 2;
3407                "},
3408                    expected: indoc! {"
3409                    fn main() {
3410                        let x = 1;
3411                        let y = 2;
3412                    }
3413                "},
3414                },
3415                Case {
3416                    name: "insert_before_first",
3417                    original: indoc! {"
3418                    let x = 1;
3419                    let y = 2;
3420                "},
3421                    model_output: indoc! {"
3422                    <|insert|>
3423                    use std::io;
3424                "},
3425                    expected: indoc! {"
3426                    use std::io;
3427                    let x = 1;
3428                    let y = 2;
3429                "},
3430                },
3431                Case {
3432                    name: "set_with_cursor_marker",
3433                    original: indoc! {"
3434                    fn main() {
3435                        println!();
3436                    }
3437                "},
3438                    model_output: indoc! {"
3439                    <|set|>1:34
3440                        eprintln!(\"<|user_cursor|>\");
3441                "},
3442                    expected: indoc! {"
3443                    fn main() {
3444                        eprintln!(\"<|user_cursor|>\");
3445                    }
3446                "},
3447                },
3448                Case {
3449                    name: "multiple_set_commands",
3450                    original: indoc! {"
3451                    aaa
3452                    bbb
3453                    ccc
3454                    ddd
3455                "},
3456                    model_output: indoc! {"
3457                    <|set|>0:23
3458                    AAA
3459                    <|set|>2:29
3460                    CCC
3461                "},
3462                    expected: indoc! {"
3463                    AAA
3464                    bbb
3465                    CCC
3466                    ddd
3467                "},
3468                },
3469                Case {
3470                    name: "set_range_multiline_replacement",
3471                    original: indoc! {"
3472                    fn handle_submit() {
3473                    }
3474
3475                    fn handle_keystroke() {
3476                "},
3477                    model_output: indoc! {"
3478                    <|set|>0:3f-1:7d
3479                    fn handle_submit(modal_state: &mut ModalState) {
3480                        <|user_cursor|>
3481                    }
3482                "},
3483                    expected: indoc! {"
3484                    fn handle_submit(modal_state: &mut ModalState) {
3485                        <|user_cursor|>
3486                    }
3487
3488                    fn handle_keystroke() {
3489                "},
3490                },
3491                Case {
3492                    name: "no_edit_commands_returns_original",
3493                    original: indoc! {"
3494                    hello
3495                    world
3496                "},
3497                    model_output: "some random text with no commands",
3498                    expected: indoc! {"
3499                    hello
3500                    world
3501                "},
3502                },
3503                Case {
3504                    name: "no_edits_command_returns_original",
3505                    original: indoc! {"
3506                    hello
3507                    world
3508                "},
3509                    model_output: "<|no_edits|>",
3510                    expected: indoc! {"
3511                    hello
3512                    world
3513                "},
3514                },
3515                Case {
3516                    name: "wrong_hash_set_ignored",
3517                    original: indoc! {"
3518                    aaa
3519                    bbb
3520                "},
3521                    model_output: indoc! {"
3522                    <|set|>0:ff
3523                    ZZZ
3524                "},
3525                    expected: indoc! {"
3526                    aaa
3527                    bbb
3528                "},
3529                },
3530                Case {
3531                    name: "insert_and_set_combined",
3532                    original: indoc! {"
3533                    alpha
3534                    beta
3535                    gamma
3536                "},
3537                    model_output: indoc! {"
3538                    <|set|>0:06
3539                    ALPHA
3540                    <|insert|>1:9c
3541                    beta_extra
3542                "},
3543                    expected: indoc! {"
3544                    ALPHA
3545                    beta
3546                    beta_extra
3547                    gamma
3548                "},
3549                },
3550                Case {
3551                    name: "no_trailing_newline_preserved",
3552                    original: "hello\nworld",
3553                    model_output: indoc! {"
3554                    <|set|>0:14
3555                    HELLO
3556                "},
3557                    expected: "HELLO\nworld",
3558                },
3559                Case {
3560                    name: "set_range_hash_mismatch_in_end_bound",
3561                    original: indoc! {"
3562                    one
3563                    two
3564                    three
3565                "},
3566                    model_output: indoc! {"
3567                    <|set|>0:42-2:ff
3568                    ONE_TWO_THREE
3569                "},
3570                    expected: indoc! {"
3571                    one
3572                    two
3573                    three
3574                "},
3575                },
3576                Case {
3577                    name: "set_range_start_greater_than_end_ignored",
3578                    original: indoc! {"
3579                    a
3580                    b
3581                    c
3582                "},
3583                    model_output: indoc! {"
3584                    <|set|>2:63-1:62
3585                    X
3586                "},
3587                    expected: indoc! {"
3588                    a
3589                    b
3590                    c
3591                "},
3592                },
3593                Case {
3594                    name: "insert_out_of_bounds_ignored",
3595                    original: indoc! {"
3596                    x
3597                    y
3598                "},
3599                    model_output: indoc! {"
3600                    <|insert|>99:aa
3601                    z
3602                "},
3603                    expected: indoc! {"
3604                    x
3605                    y
3606                "},
3607                },
3608                Case {
3609                    name: "set_out_of_bounds_ignored",
3610                    original: indoc! {"
3611                    x
3612                    y
3613                "},
3614                    model_output: indoc! {"
3615                    <|set|>99:aa
3616                    z
3617                "},
3618                    expected: indoc! {"
3619                    x
3620                    y
3621                "},
3622                },
3623                Case {
3624                    name: "malformed_set_command_ignored",
3625                    original: indoc! {"
3626                    alpha
3627                    beta
3628                "},
3629                    model_output: indoc! {"
3630                    <|set|>not-a-line-ref
3631                    UPDATED
3632                "},
3633                    expected: indoc! {"
3634                    alpha
3635                    beta
3636                "},
3637                },
3638                Case {
3639                    name: "malformed_insert_hash_treated_as_before_first",
3640                    original: indoc! {"
3641                    alpha
3642                    beta
3643                "},
3644                    model_output: indoc! {"
3645                    <|insert|>1:nothex
3646                    preamble
3647                "},
3648                    expected: indoc! {"
3649                    preamble
3650                    alpha
3651                    beta
3652                "},
3653                },
3654                Case {
3655                    name: "set_then_insert_same_target_orders_insert_after_replacement",
3656                    original: indoc! {"
3657                    cat
3658                    dog
3659                "},
3660                    model_output: indoc! {"
3661                    <|set|>0:38
3662                    CAT
3663                    <|insert|>0:38
3664                    TAIL
3665                "},
3666                    expected: indoc! {"
3667                    CAT
3668                    TAIL
3669                    dog
3670                "},
3671                },
3672                Case {
3673                    name: "overlapping_set_ranges_last_wins",
3674                    original: indoc! {"
3675                    a
3676                    b
3677                    c
3678                    d
3679                "},
3680                    model_output: indoc! {"
3681                    <|set|>0:61-2:63
3682                    FIRST
3683                    <|set|>1:62-3:64
3684                    SECOND
3685                "},
3686                    expected: indoc! {"
3687                    FIRST
3688                    d
3689                "},
3690                },
3691                Case {
3692                    name: "insert_before_first_and_after_line",
3693                    original: indoc! {"
3694                        a
3695                        b
3696                    "},
3697                    model_output: indoc! {"
3698                        <|insert|>
3699                        HEAD
3700                        <|insert|>0:61
3701                        MID
3702                    "},
3703                    expected: indoc! {"
3704                        HEAD
3705                        a
3706                        MID
3707                        b
3708                    "},
3709                },
3710            ];
3711
3712            for case in &cases {
3713                let result = hashline::apply_edit_commands(case.original, &case.model_output);
3714                assert_eq!(result, case.expected, "failed case: {}", case.name);
3715            }
3716        }
3717
3718        #[test]
3719        fn test_output_has_edit_commands() {
3720            assert!(hashline::output_has_edit_commands(&format!(
3721                "{}0:ab\nnew",
3722                SET_COMMAND_MARKER
3723            )));
3724            assert!(hashline::output_has_edit_commands(&format!(
3725                "{}0:ab\nnew",
3726                INSERT_COMMAND_MARKER
3727            )));
3728            assert!(hashline::output_has_edit_commands(&format!(
3729                "some text\n{}1:cd\nstuff",
3730                SET_COMMAND_MARKER
3731            )));
3732            assert!(!hashline::output_has_edit_commands("just plain text"));
3733            assert!(!hashline::output_has_edit_commands("NO_EDITS"));
3734            assert!(hashline::output_has_edit_commands("<|no_edits|>"));
3735        }
3736
3737        // ---- hashline::patch_to_edit_commands round-trip tests ----
3738
3739        #[test]
3740        fn test_patch_to_edit_commands() {
3741            struct Case {
3742                name: &'static str,
3743                old: &'static str,
3744                patch: &'static str,
3745                expected_new: &'static str,
3746            }
3747
3748            let cases = [
3749                Case {
3750                    name: "single_line_replacement",
3751                    old: indoc! {"
3752                    let mut total = 0;
3753                    for product in products {
3754                        total += ;
3755                    }
3756                    total
3757                "},
3758                    patch: indoc! {"
3759                    @@ -1,5 +1,5 @@
3760                     let mut total = 0;
3761                     for product in products {
3762                    -    total += ;
3763                    +    total += product.price;
3764                     }
3765                     total
3766                "},
3767                    expected_new: indoc! {"
3768                    let mut total = 0;
3769                    for product in products {
3770                        total += product.price;
3771                    }
3772                    total
3773                "},
3774                },
3775                Case {
3776                    name: "multiline_replacement",
3777                    old: indoc! {"
3778                    fn foo() {
3779                        let x = 1;
3780                        let y = 2;
3781                        let z = 3;
3782                    }
3783                "},
3784                    patch: indoc! {"
3785                    @@ -1,5 +1,3 @@
3786                     fn foo() {
3787                    -    let x = 1;
3788                    -    let y = 2;
3789                    -    let z = 3;
3790                    +    let sum = 1 + 2 + 3;
3791                     }
3792                "},
3793                    expected_new: indoc! {"
3794                    fn foo() {
3795                        let sum = 1 + 2 + 3;
3796                    }
3797                "},
3798                },
3799                Case {
3800                    name: "insertion",
3801                    old: indoc! {"
3802                    fn main() {
3803                        let x = 1;
3804                    }
3805                "},
3806                    patch: indoc! {"
3807                    @@ -1,3 +1,4 @@
3808                     fn main() {
3809                         let x = 1;
3810                    +    let y = 2;
3811                     }
3812                "},
3813                    expected_new: indoc! {"
3814                    fn main() {
3815                        let x = 1;
3816                        let y = 2;
3817                    }
3818                "},
3819                },
3820                Case {
3821                    name: "insertion_before_first",
3822                    old: indoc! {"
3823                    let x = 1;
3824                    let y = 2;
3825                "},
3826                    patch: indoc! {"
3827                    @@ -1,2 +1,3 @@
3828                    +use std::io;
3829                     let x = 1;
3830                     let y = 2;
3831                "},
3832                    expected_new: indoc! {"
3833                    use std::io;
3834                    let x = 1;
3835                    let y = 2;
3836                "},
3837                },
3838                Case {
3839                    name: "deletion",
3840                    old: indoc! {"
3841                    aaa
3842                    bbb
3843                    ccc
3844                    ddd
3845                "},
3846                    patch: indoc! {"
3847                    @@ -1,4 +1,2 @@
3848                     aaa
3849                    -bbb
3850                    -ccc
3851                     ddd
3852                "},
3853                    expected_new: indoc! {"
3854                    aaa
3855                    ddd
3856                "},
3857                },
3858                Case {
3859                    name: "multiple_changes",
3860                    old: indoc! {"
3861                    alpha
3862                    beta
3863                    gamma
3864                    delta
3865                    epsilon
3866                "},
3867                    patch: indoc! {"
3868                    @@ -1,5 +1,5 @@
3869                    -alpha
3870                    +ALPHA
3871                     beta
3872                     gamma
3873                    -delta
3874                    +DELTA
3875                     epsilon
3876                "},
3877                    expected_new: indoc! {"
3878                    ALPHA
3879                    beta
3880                    gamma
3881                    DELTA
3882                    epsilon
3883                "},
3884                },
3885                Case {
3886                    name: "replace_with_insertion",
3887                    old: indoc! {r#"
3888                    fn handle() {
3889                        modal_state.close();
3890                        modal_state.dismiss();
3891                "#},
3892                    patch: indoc! {r#"
3893                    @@ -1,3 +1,4 @@
3894                     fn handle() {
3895                         modal_state.close();
3896                    +    eprintln!("");
3897                         modal_state.dismiss();
3898                "#},
3899                    expected_new: indoc! {r#"
3900                    fn handle() {
3901                        modal_state.close();
3902                        eprintln!("");
3903                        modal_state.dismiss();
3904                "#},
3905                },
3906                Case {
3907                    name: "complete_replacement",
3908                    old: indoc! {"
3909                    aaa
3910                    bbb
3911                    ccc
3912                "},
3913                    patch: indoc! {"
3914                    @@ -1,3 +1,3 @@
3915                    -aaa
3916                    -bbb
3917                    -ccc
3918                    +xxx
3919                    +yyy
3920                    +zzz
3921                "},
3922                    expected_new: indoc! {"
3923                    xxx
3924                    yyy
3925                    zzz
3926                "},
3927                },
3928                Case {
3929                    name: "add_function_body",
3930                    old: indoc! {"
3931                    fn foo() {
3932                        modal_state.dismiss();
3933                    }
3934
3935                    fn
3936
3937                    fn handle_keystroke() {
3938                "},
3939                    patch: indoc! {"
3940                    @@ -1,6 +1,8 @@
3941                     fn foo() {
3942                         modal_state.dismiss();
3943                     }
3944
3945                    -fn
3946                    +fn handle_submit() {
3947                    +    todo()
3948                    +}
3949
3950                     fn handle_keystroke() {
3951                "},
3952                    expected_new: indoc! {"
3953                    fn foo() {
3954                        modal_state.dismiss();
3955                    }
3956
3957                    fn handle_submit() {
3958                        todo()
3959                    }
3960
3961                    fn handle_keystroke() {
3962                "},
3963                },
3964                Case {
3965                    name: "with_cursor_offset",
3966                    old: indoc! {r#"
3967                    fn main() {
3968                        println!();
3969                    }
3970                "#},
3971                    patch: indoc! {r#"
3972                        @@ -1,3 +1,3 @@
3973                        fn main() {
3974                        -    println!();
3975                        +    eprintln!("");
3976                        }
3977                    "#},
3978                    expected_new: indoc! {r#"
3979                        fn main() {
3980                            eprintln!("<|user_cursor|>");
3981                        }
3982                    "#},
3983                },
3984                Case {
3985                    name: "non_local_hunk_header_pure_insertion_repro",
3986                    old: indoc! {"
3987                        aaa
3988                        bbb
3989                    "},
3990                    patch: indoc! {"
3991                        @@ -20,2 +20,3 @@
3992                        aaa
3993                        +xxx
3994                        bbb
3995                    "},
3996                    expected_new: indoc! {"
3997                        aaa
3998                        xxx
3999                        bbb
4000                    "},
4001                },
4002                Case {
4003                    name: "empty_patch_produces_no_edits_marker",
4004                    old: indoc! {"
4005                        aaa
4006                        bbb
4007                    "},
4008                    patch: "@@ -20,2 +20,3 @@\n",
4009                    expected_new: indoc! {"
4010                        aaa
4011                        bbb
4012                    "},
4013                },
4014            ];
4015
4016            for case in &cases {
4017                // The cursor_offset for patch_to_edit_commands is relative to
4018                // the first hunk's new text (context + additions). We compute
4019                // it by finding where the marker sits in the expected output
4020                // (which mirrors the new text of the hunk).
4021                let cursor_offset = case.expected_new.find(CURSOR_MARKER);
4022
4023                let commands =
4024                    hashline::patch_to_edit_commands(case.old, case.patch, cursor_offset)
4025                        .unwrap_or_else(|e| panic!("failed case {}: {e}", case.name));
4026
4027                assert!(
4028                    hashline::output_has_edit_commands(&commands),
4029                    "case {}: expected edit commands, got: {commands:?}",
4030                    case.name,
4031                );
4032
4033                let applied = hashline::apply_edit_commands(case.old, &commands);
4034                assert_eq!(applied, case.expected_new, "case {}", case.name);
4035            }
4036        }
4037    }
4038}
4039
4040pub mod qwen {
4041    use super::*;
4042
4043    pub const FIM_PREFIX: &str = "<|fim_prefix|>";
4044    pub const FIM_SUFFIX: &str = "<|fim_suffix|>";
4045    pub const FIM_MIDDLE: &str = "<|fim_middle|>";
4046    pub const FILE_MARKER: &str = "<|file_sep|>";
4047    pub const END_MARKER: &str = "<|im_end|>";
4048
4049    pub fn assemble_fim_prompt(
4050        context: &str,
4051        editable_range: &Range<usize>,
4052        cursor_prefix_section: &str,
4053        events: &[Arc<Event>],
4054        related_files: &[RelatedFile],
4055        max_tokens: usize,
4056    ) -> String {
4057        let suffix_section = build_suffix_section(context, editable_range);
4058
4059        let cursor_prefix_tokens = estimate_tokens(cursor_prefix_section.len() + FIM_PREFIX.len());
4060        let suffix_tokens =
4061            estimate_tokens(suffix_section.len() + FIM_SUFFIX.len() + FIM_MIDDLE.len());
4062        let budget_after_cursor = max_tokens.saturating_sub(cursor_prefix_tokens + suffix_tokens);
4063
4064        let edit_history_section = super::format_edit_history_within_budget(
4065            events,
4066            FILE_MARKER,
4067            "edit_history",
4068            budget_after_cursor,
4069            max_edit_event_count_for_format(&ZetaFormat::V0608QwenMultiRegions),
4070        );
4071        let edit_history_tokens = estimate_tokens(edit_history_section.len() + "\n".len());
4072        let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens);
4073
4074        let related_files_section = super::format_related_files_within_budget(
4075            related_files,
4076            FILE_MARKER,
4077            "",
4078            budget_after_edit_history,
4079        );
4080
4081        let mut prompt = String::new();
4082        prompt.push_str(FIM_PREFIX);
4083        prompt.push_str(&related_files_section);
4084        if !related_files_section.is_empty() {
4085            prompt.push('\n');
4086        }
4087        prompt.push_str(&edit_history_section);
4088        if !edit_history_section.is_empty() {
4089            prompt.push('\n');
4090        }
4091        prompt.push_str(cursor_prefix_section);
4092        prompt.push_str(&suffix_section);
4093        prompt.push_str(FIM_MIDDLE);
4094        prompt
4095    }
4096
4097    fn build_suffix_section(context: &str, editable_range: &Range<usize>) -> String {
4098        let mut section = String::new();
4099        section.push_str(FIM_SUFFIX);
4100        section.push_str(&context[editable_range.end..]);
4101        if !section.ends_with('\n') {
4102            section.push('\n');
4103        }
4104        section
4105    }
4106}
4107
4108pub mod seed_coder {
4109    //! Seed-Coder prompt format using SPM (Suffix-Prefix-Middle) FIM mode.
4110    //!
4111    //! Seed-Coder uses different FIM tokens and order than Qwen:
4112    //! - SPM order: suffix comes FIRST, then prefix, then middle
4113    //! - Tokens: `<[fim-suffix]>`, `<[fim-prefix]>`, `<[fim-middle]>`
4114    //! - File markers: StarCoder-style `<filename>path` (single token + path)
4115    //!
4116    //! All context (related files, edit history) goes in the PREFIX section.
4117    //! The suffix contains only code after the editable region.
4118    //!
4119    //! Example prompt:
4120    //!
4121    //! <[fim-suffix]>
4122    //! code after editable region
4123    //! <[fim-prefix]><filename>related/file.py
4124    //! related file content
4125    //!
4126    //! <filename>edit_history
4127    //! --- a/some_file.py
4128    //! +++ b/some_file.py
4129    //! -old
4130    //! +new
4131    //!
4132    //! <filename>path/to/target_file.py
4133    //! code before editable region
4134    //! <<<<<<< CURRENT
4135    //! code that
4136    //! needs to<|user_cursor|>
4137    //! be rewritten
4138    //! =======
4139    //! <[fim-middle]>
4140    //!
4141    //! Expected output (model generates):
4142    //!
4143    //! updated
4144    //! code with
4145    //! changes applied
4146    //! >>>>>>> UPDATED
4147
4148    use super::*;
4149
4150    pub const FIM_SUFFIX: &str = "<[fim-suffix]>";
4151    pub const FIM_PREFIX: &str = "<[fim-prefix]>";
4152    pub const FIM_MIDDLE: &str = "<[fim-middle]>";
4153    pub const FILE_MARKER: &str = "<filename>";
4154
4155    pub const START_MARKER: &str = "<<<<<<< CURRENT\n";
4156    pub const SEPARATOR: &str = "=======\n";
4157    pub const END_MARKER: &str = ">>>>>>> UPDATED\n";
4158
4159    pub const NO_EDITS: &str = "NO_EDITS\n";
4160
4161    pub fn special_tokens() -> &'static [&'static str] {
4162        &[
4163            FIM_SUFFIX,
4164            FIM_PREFIX,
4165            FIM_MIDDLE,
4166            FILE_MARKER,
4167            START_MARKER,
4168            SEPARATOR,
4169            END_MARKER,
4170            CURSOR_MARKER,
4171        ]
4172    }
4173
4174    pub fn write_cursor_excerpt_section(
4175        prompt: &mut String,
4176        path: &Path,
4177        context: &str,
4178        editable_range: &Range<usize>,
4179        cursor_offset: usize,
4180    ) {
4181        let section = build_cursor_prefix_section(path, context, editable_range, cursor_offset);
4182        prompt.push_str(&section);
4183    }
4184
4185    pub fn format_prompt_with_budget(
4186        path: &Path,
4187        context: &str,
4188        editable_range: &Range<usize>,
4189        cursor_offset: usize,
4190        events: &[Arc<Event>],
4191        related_files: &[RelatedFile],
4192        diagnostics: &[ActiveBufferDiagnostic],
4193        max_tokens: usize,
4194    ) -> String {
4195        let cursor_prefix_section =
4196            build_cursor_prefix_section(path, context, editable_range, cursor_offset);
4197        assemble_fim_prompt(
4198            context,
4199            editable_range,
4200            &cursor_prefix_section,
4201            events,
4202            related_files,
4203            diagnostics,
4204            None,
4205            max_tokens,
4206        )
4207    }
4208
4209    pub fn assemble_fim_prompt(
4210        context: &str,
4211        editable_range: &Range<usize>,
4212        cursor_prefix_section: &str,
4213        events: &[Arc<Event>],
4214        related_files: &[RelatedFile],
4215        diagnostics: &[ActiveBufferDiagnostic],
4216        cursor_buffer_row: Option<u32>,
4217        max_tokens: usize,
4218    ) -> String {
4219        let suffix_section = build_suffix_section(context, editable_range);
4220
4221        let suffix_tokens = estimate_tokens(suffix_section.len() + FIM_PREFIX.len());
4222        let cursor_prefix_tokens = estimate_tokens(cursor_prefix_section.len() + FIM_MIDDLE.len());
4223        let budget_after_cursor = max_tokens.saturating_sub(suffix_tokens + cursor_prefix_tokens);
4224
4225        let edit_history_section = super::format_edit_history_within_budget(
4226            events,
4227            FILE_MARKER,
4228            "edit_history",
4229            budget_after_cursor,
4230            max_edit_event_count_for_format(&ZetaFormat::V0211SeedCoder),
4231        );
4232        let edit_history_tokens = estimate_tokens(edit_history_section.len() + "\n".len());
4233        let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens);
4234
4235        let diagnostics_section = super::format_active_buffer_diagnostics_with_budget(
4236            diagnostics,
4237            cursor_buffer_row,
4238            budget_after_edit_history,
4239        );
4240        let diagnostics_tokens = estimate_tokens(diagnostics_section.len() + "\n".len());
4241        let budget_after_diagnostics = budget_after_edit_history.saturating_sub(diagnostics_tokens);
4242
4243        let related_files_section = super::format_related_files_within_budget(
4244            related_files,
4245            FILE_MARKER,
4246            "",
4247            budget_after_diagnostics,
4248        );
4249
4250        let mut prompt = String::new();
4251        prompt.push_str(&suffix_section);
4252        prompt.push_str(FIM_PREFIX);
4253        prompt.push_str(&diagnostics_section);
4254        if !diagnostics_section.is_empty() {
4255            prompt.push('\n');
4256        }
4257        prompt.push_str(&related_files_section);
4258        if !related_files_section.is_empty() {
4259            prompt.push('\n');
4260        }
4261        prompt.push_str(&edit_history_section);
4262        if !edit_history_section.is_empty() {
4263            prompt.push('\n');
4264        }
4265        prompt.push_str(cursor_prefix_section);
4266        prompt.push_str(FIM_MIDDLE);
4267
4268        prompt
4269    }
4270
4271    pub(crate) fn build_suffix_section(context: &str, editable_range: &Range<usize>) -> String {
4272        let mut section = String::new();
4273        section.push_str(FIM_SUFFIX);
4274        section.push_str(&context[editable_range.end..]);
4275        if !section.ends_with('\n') {
4276            section.push('\n');
4277        }
4278        section
4279    }
4280
4281    fn build_cursor_prefix_section(
4282        path: &Path,
4283        context: &str,
4284        editable_range: &Range<usize>,
4285        cursor_offset: usize,
4286    ) -> String {
4287        let mut section = String::new();
4288        let path_str = path.to_string_lossy();
4289        write!(section, "{}{}\n", FILE_MARKER, path_str).ok();
4290
4291        section.push_str(&context[..editable_range.start]);
4292        section.push_str(START_MARKER);
4293        section.push_str(&context[editable_range.start..cursor_offset]);
4294        section.push_str(CURSOR_MARKER);
4295        section.push_str(&context[cursor_offset..editable_range.end]);
4296        if !section.ends_with('\n') {
4297            section.push('\n');
4298        }
4299        section.push_str(SEPARATOR);
4300        section
4301    }
4302
4303    /// Format patch as containing no changes if it's empty; otherwise return None.
4304    pub(crate) fn no_edits(patch: &str) -> Option<String> {
4305        // Count lines in the patch
4306        let empty_patch = patch.lines().count() <= 3;
4307        if empty_patch {
4308            Some(format!("{NO_EDITS}{END_MARKER}"))
4309        } else {
4310            None
4311        }
4312    }
4313}
4314
4315pub mod v0304_variable_edit {
4316    //! A prompt format with no fixed editable region. The entire context is shown
4317    //! to the model, and it chooses which text to replace by outputting surrounding
4318    //! context lines with `<|fim_middle|>` and `<|fim_suffix|>` delimiting the new
4319    //! text.
4320    //!
4321    //! Example prompt:
4322    //!
4323    //! <|file_sep|>path/to/file.py
4324    //! zero
4325    //! one
4326    //! two
4327    //! three<|user_cursor|>
4328    //! four
4329    //! five
4330    //! <|fim_prefix|>
4331    //
4332    //! Expected output (model generates):
4333    //!
4334    //! two
4335    //! <|fim_middle|>
4336    //! THREE
4337    //! <|fim_suffix|>
4338    //! four
4339    //!
4340    //! The output means: find "two\n...\nfour" in the context, and replace
4341    //! everything between "two\n" and "four" with "THREE\n".
4342
4343    use super::*;
4344
4345    pub fn special_tokens() -> &'static [&'static str] {
4346        &[
4347            "<|fim_prefix|>",
4348            "<|fim_suffix|>",
4349            "<|fim_middle|>",
4350            "<|file_sep|>",
4351            CURSOR_MARKER,
4352        ]
4353    }
4354
4355    pub fn write_cursor_excerpt_section(
4356        prompt: &mut String,
4357        path: &Path,
4358        context: &str,
4359        cursor_offset: usize,
4360    ) {
4361        let path_str = path.to_string_lossy();
4362        write!(prompt, "<|file_sep|>{}\n", path_str).ok();
4363
4364        prompt.push_str(&context[..cursor_offset]);
4365        prompt.push_str(CURSOR_MARKER);
4366        prompt.push_str(&context[cursor_offset..]);
4367        if !prompt.ends_with('\n') {
4368            prompt.push('\n');
4369        }
4370        prompt.push_str("<|fim_prefix|>\n")
4371    }
4372
4373    /// Apply a variable-edit model output to the original context text.
4374    ///
4375    /// The model output has the form:
4376    ///
4377    /// - prefix context lines
4378    /// - `<|fim_middle|>`
4379    /// - new text
4380    /// - `<|fim_suffix|>`
4381    /// - suffix context lines
4382    ///
4383    /// We locate the prefix/suffix context lines in the original text and replace
4384    /// everything between them with the new text.
4385    pub fn apply_variable_edit(
4386        context: &str,
4387        model_output: &str,
4388    ) -> Result<(Range<usize>, String)> {
4389        let (prefix_context, rest) = model_output
4390            .split_once("<|fim_middle|>\n")
4391            .or_else(|| model_output.split_once("<|fim_middle|>"))
4392            .ok_or_else(|| anyhow::anyhow!("missing <|fim_middle|> in model output"))?;
4393
4394        let (new_text, suffix_context) = rest
4395            .split_once("<|fim_suffix|>\n")
4396            .or_else(|| rest.split_once("<|fim_suffix|>"))
4397            .unwrap_or((rest, ""));
4398
4399        let suffix_context = if prefix_context.is_empty() && !suffix_context.is_empty() {
4400            suffix_context.strip_prefix('\n').unwrap_or(suffix_context)
4401        } else {
4402            suffix_context
4403        };
4404
4405        let prefix_offset = find_substring_at_line_boundary(context, prefix_context)
4406            .ok_or_else(|| anyhow!("could not locate prefix lines"))?
4407            + prefix_context.len();
4408        let suffix_offset = if suffix_context.is_empty() {
4409            context.len()
4410        } else {
4411            find_substring_at_line_boundary(&context[prefix_offset..], suffix_context)
4412                .ok_or_else(|| anyhow!("could not locate suffix lines"))?
4413                + prefix_offset
4414        };
4415
4416        let edit_range = prefix_offset..suffix_offset;
4417        return Ok((edit_range, new_text.to_string()));
4418    }
4419
4420    fn find_substring_at_line_boundary(haystack: &str, needle: &str) -> Option<usize> {
4421        if needle.is_empty() {
4422            return Some(0);
4423        }
4424
4425        haystack.match_indices(needle).find_map(|(offset, _)| {
4426            let matched_line_start = offset == 0 || haystack[..offset].ends_with('\n');
4427            matched_line_start.then_some(offset)
4428        })
4429    }
4430
4431    /// Convert a unified diff patch into the variable-edit output format.
4432    ///
4433    /// Parses `patch` as a unified diff against `old_text` and produces model
4434    /// output with context lines surrounding `<|fim_middle|>` / `<|fim_suffix|>`
4435    /// delimiters. The diff is resolved by content matching rather than line
4436    /// numbers.
4437    pub fn patch_to_variable_edit_output(
4438        old_text: &str,
4439        patch: &str,
4440        cursor_offset: Option<usize>,
4441    ) -> Result<String> {
4442        // Parse the unified diff into hunks. Each hunk has an `old_context`
4443        // string (context + deleted lines interleaved in order) and a list of
4444        // edits expressed as byte ranges within that context plus replacement
4445        // text.
4446        let hunks = parse_hunks(patch);
4447        if hunks.is_empty() {
4448            return Ok(String::new());
4449        }
4450
4451        // Apply each hunk by finding its old_context in the text and
4452        // performing the edits. We search forward from where the previous
4453        // hunk ended so that hunks are applied in order.
4454        let mut new_text = old_text.to_string();
4455        let mut search_from: usize = 0;
4456        let mut first_hunk_pos: Option<usize> = None;
4457
4458        for hunk in &hunks {
4459            let context_pos = new_text[search_from..]
4460                .find(&hunk.old_context)
4461                .map(|pos| pos + search_from)
4462                .ok_or_else(|| anyhow::anyhow!("could not locate hunk context in text"))?;
4463
4464            if first_hunk_pos.is_none() {
4465                first_hunk_pos = Some(context_pos);
4466            }
4467
4468            // Apply edits in reverse order so byte offsets remain valid.
4469            for edit in hunk.edits.iter().rev() {
4470                let abs_start = context_pos + edit.range.start;
4471                let abs_end = context_pos + edit.range.end;
4472                new_text.replace_range(abs_start..abs_end, &edit.text);
4473            }
4474
4475            // Advance past this hunk's region in the (now modified) text.
4476            let new_region_len: usize =
4477                hunk.edits.iter().fold(hunk.old_context.len(), |len, edit| {
4478                    len + edit.text.len() - (edit.range.end - edit.range.start)
4479                });
4480            search_from = context_pos + new_region_len;
4481        }
4482
4483        // Now we have old_text and new_text. Find the changed line range by
4484        // comparing them.
4485        let old_lines: Vec<&str> = old_text.lines().collect();
4486        let new_lines: Vec<&str> = new_text.lines().collect();
4487
4488        // Find first differing line.
4489        let first_changed_row = old_lines
4490            .iter()
4491            .zip(new_lines.iter())
4492            .position(|(a, b)| a != b)
4493            .unwrap_or_else(|| old_lines.len().min(new_lines.len()));
4494
4495        // Find last differing line (from the end).
4496        let max_suffix = old_lines.len().min(new_lines.len()) - first_changed_row;
4497        let common_suffix = old_lines
4498            .iter()
4499            .rev()
4500            .zip(new_lines.iter().rev())
4501            .take(max_suffix)
4502            .take_while(|(a, b)| a == b)
4503            .count();
4504
4505        let old_end = old_lines.len() - common_suffix;
4506        let new_end = new_lines.len() - common_suffix;
4507
4508        if first_changed_row == old_end && first_changed_row == new_end {
4509            return Ok(String::new());
4510        }
4511
4512        // Build the replacement text from new_lines[first_diff..new_end].
4513        let mut merged_new_text = String::new();
4514        for line in &new_lines[first_changed_row..new_end] {
4515            merged_new_text.push_str(line);
4516            merged_new_text.push('\n');
4517        }
4518
4519        // cursor_offset is relative to the first hunk's new content in
4520        // new_text. Translate it to an offset within merged_new_text, which
4521        // only contains lines first_diff..new_end of new_text.
4522        if let Some(hunk_offset) = cursor_offset {
4523            let hunk_start = first_hunk_pos.unwrap_or(0);
4524            let absolute_pos = hunk_start + hunk_offset;
4525
4526            // Byte offset where first_diff starts in new_text.
4527            let merged_start: usize = new_lines[..first_changed_row]
4528                .iter()
4529                .map(|line| line.len() + 1)
4530                .sum();
4531
4532            if absolute_pos >= merged_start {
4533                let relative_offset = absolute_pos - merged_start;
4534                if relative_offset <= merged_new_text.len() {
4535                    merged_new_text.insert_str(relative_offset, CURSOR_MARKER);
4536                }
4537            }
4538        }
4539
4540        // Build output with 2 lines of context above and below.
4541        let context_lines_count = 2;
4542        let mut prefix_start = first_changed_row.saturating_sub(context_lines_count);
4543        let mut suffix_end = (old_end + context_lines_count).min(old_lines.len());
4544
4545        fn count_matches(line_range: Range<usize>, lines: &[&str]) -> usize {
4546            let pattern = &lines[line_range];
4547            let pattern_len = pattern.len();
4548
4549            let mut count = 0;
4550            for offset in 0..=lines.len() - pattern_len {
4551                if &lines[offset..offset + pattern_len] == pattern {
4552                    count += 1;
4553                }
4554            }
4555            count
4556        }
4557
4558        // Expand prefix and suffix until they are unique
4559        while prefix_start > 0 {
4560            if count_matches(prefix_start..first_changed_row, &old_lines) > 1 {
4561                prefix_start -= 1;
4562            } else {
4563                break;
4564            }
4565        }
4566        while suffix_end < old_lines.len() {
4567            if count_matches(old_end..suffix_end, &old_lines) > 1 {
4568                suffix_end += 1;
4569            } else {
4570                break;
4571            }
4572        }
4573
4574        let mut output = String::new();
4575        for line in &old_lines[prefix_start..first_changed_row] {
4576            output.push_str(line);
4577            output.push('\n');
4578        }
4579        output.push_str("<|fim_middle|>\n");
4580        output.push_str(&merged_new_text);
4581        output.push_str("<|fim_suffix|>\n");
4582        for line in &old_lines[old_end..suffix_end] {
4583            output.push_str(line);
4584            output.push('\n');
4585        }
4586
4587        Ok(output)
4588    }
4589
4590    struct ParsedHunk {
4591        old_context: String,
4592        edits: Vec<ParsedEdit>,
4593    }
4594
4595    struct ParsedEdit {
4596        range: Range<usize>,
4597        text: String,
4598    }
4599
4600    /// Parse a unified diff into content-based hunks. Each hunk contains an
4601    /// `old_context` string (context lines + deleted lines, which together
4602    /// form the text that should be found in the original) and a list of edits
4603    /// expressed as byte ranges within that context.
4604    fn parse_hunks(patch: &str) -> Vec<ParsedHunk> {
4605        let mut hunks = Vec::new();
4606        let mut current: Option<ParsedHunk> = None;
4607
4608        for line in patch.lines() {
4609            if line.starts_with("@@") {
4610                if let Some(hunk) = current.take() {
4611                    if !hunk.old_context.is_empty() || !hunk.edits.is_empty() {
4612                        hunks.push(hunk);
4613                    }
4614                }
4615                current = Some(ParsedHunk {
4616                    old_context: String::new(),
4617                    edits: Vec::new(),
4618                });
4619            } else if line.starts_with("---") || line.starts_with("+++") {
4620                continue;
4621            } else if let Some(hunk) = &mut current {
4622                if let Some(added) = line.strip_prefix('+') {
4623                    let pos = hunk.old_context.len();
4624                    if let Some(last_edit) = hunk.edits.last_mut() {
4625                        if last_edit.range.end == pos {
4626                            writeln!(&mut last_edit.text, "{added}").ok();
4627                            continue;
4628                        }
4629                    }
4630                    hunk.edits.push(ParsedEdit {
4631                        range: pos..pos,
4632                        text: format!("{added}\n"),
4633                    });
4634                } else if let Some(removed) = line.strip_prefix('-') {
4635                    let start = hunk.old_context.len();
4636                    writeln!(&mut hunk.old_context, "{removed}").ok();
4637                    let end = hunk.old_context.len();
4638                    if let Some(last_edit) = hunk.edits.last_mut() {
4639                        if last_edit.range.end == start {
4640                            last_edit.range.end = end;
4641                            continue;
4642                        }
4643                    }
4644                    hunk.edits.push(ParsedEdit {
4645                        range: start..end,
4646                        text: String::new(),
4647                    });
4648                } else {
4649                    let ctx = line.strip_prefix(' ').unwrap_or(line);
4650                    writeln!(&mut hunk.old_context, "{ctx}").ok();
4651                }
4652            }
4653        }
4654
4655        if let Some(hunk) = current {
4656            if !hunk.old_context.is_empty() || !hunk.edits.is_empty() {
4657                hunks.push(hunk);
4658            }
4659        }
4660
4661        hunks
4662    }
4663
4664    #[cfg(test)]
4665    mod tests {
4666        use super::*;
4667        use indoc::indoc;
4668
4669        #[test]
4670        fn test_apply_variable_edit() {
4671            struct Case {
4672                name: &'static str,
4673                original: &'static str,
4674                model_output: &'static str,
4675                expected: &'static str,
4676            }
4677
4678            let cases = [
4679                Case {
4680                    name: "simple_single_line_replacement",
4681                    original: indoc! {"
4682                        zero
4683                        one
4684                        two
4685                        three
4686                        four
4687                        five
4688                    "},
4689                    model_output: indoc! {"
4690                        two
4691                        <|fim_middle|>
4692                        THREE
4693                        <|fim_suffix|>
4694                        four
4695                    "},
4696                    expected: indoc! {"
4697                        zero
4698                        one
4699                        two
4700                        THREE
4701                        four
4702                        five
4703                    "},
4704                },
4705                Case {
4706                    name: "multi_line_replacement",
4707                    original: indoc! {"
4708                        a
4709                        b
4710                        c
4711                        d
4712                        e
4713                    "},
4714                    model_output: indoc! {"
4715                        a
4716                        <|fim_middle|>
4717                        B
4718                        C
4719                        D
4720                        <|fim_suffix|>
4721                        e
4722                    "},
4723                    expected: indoc! {"
4724                        a
4725                        B
4726                        C
4727                        D
4728                        e
4729                    "},
4730                },
4731                Case {
4732                    name: "insertion_between_existing_lines",
4733                    original: indoc! {"
4734                        a
4735                        b
4736                        c
4737                    "},
4738                    model_output: indoc! {"
4739                        a
4740                        <|fim_middle|>
4741                        X
4742                        <|fim_suffix|>
4743                        b
4744                    "},
4745                    expected: indoc! {"
4746                        a
4747                        X
4748                        b
4749                        c
4750                    "},
4751                },
4752                Case {
4753                    name: "deletion",
4754                    original: indoc! {"
4755                        a
4756                        b
4757                        c
4758                        d
4759                    "},
4760                    model_output: indoc! {"
4761                        a
4762                        <|fim_middle|>
4763                        <|fim_suffix|>
4764                        c
4765                    "},
4766                    expected: indoc! {"
4767                        a
4768                        c
4769                        d
4770                    "},
4771                },
4772                Case {
4773                    name: "replacement_at_start_no_prefix_context",
4774                    original: indoc! {"
4775                        a
4776                        b
4777                        c
4778                    "},
4779                    model_output: indoc! {"
4780                        <|fim_middle|>
4781                        X
4782                        <|fim_suffix|>
4783                        b
4784                    "},
4785                    expected: indoc! {"
4786                        X
4787                        b
4788                        c
4789                    "},
4790                },
4791                Case {
4792                    name: "replacement_at_end_no_suffix_context",
4793                    original: indoc! {"
4794                        a
4795                        b
4796                        c
4797                    "},
4798                    model_output: indoc! {"
4799                        b
4800                        <|fim_middle|>
4801                        Z
4802                        <|fim_suffix|>
4803                    "},
4804                    expected: indoc! {"
4805                        a
4806                        b
4807                        Z
4808                    "},
4809                },
4810                Case {
4811                    name: "context_with_trailing_newline_is_preserved",
4812                    original: indoc! {"
4813                        a
4814                        b
4815                        c
4816                    "},
4817                    model_output: indoc! {"
4818                        a
4819                        <|fim_middle|>
4820                        B
4821                        <|fim_suffix|>
4822                        c
4823                    "},
4824                    expected: indoc! {"
4825                        a
4826                        B
4827                        c
4828                    "},
4829                },
4830                Case {
4831                    name: "cursor_marker_passes_through_untouched",
4832                    original: indoc! {"
4833                        a
4834                        b
4835                        c
4836                    "},
4837                    model_output: indoc! {"
4838                        a
4839                        <|fim_middle|>
4840                        B<|user_cursor|>B
4841                        <|fim_suffix|>
4842                        c
4843                    "},
4844                    expected: indoc! {"
4845                        a
4846                        B<|user_cursor|>B
4847                        c
4848                    "},
4849                },
4850                Case {
4851                    name: "multiple_prefix_context_lines",
4852                    original: indoc! {"
4853                        a
4854                        b
4855                        c
4856                        d
4857                        e
4858                    "},
4859                    model_output: indoc! {"
4860                        b
4861                        c
4862                        <|fim_middle|>
4863                        D
4864                        <|fim_suffix|>
4865                        e
4866                    "},
4867                    expected: indoc! {"
4868                        a
4869                        b
4870                        c
4871                        D
4872                        e
4873                    "},
4874                },
4875            ];
4876
4877            for case in cases {
4878                let (edit_range, replacement) =
4879                    apply_variable_edit(case.original, case.model_output).unwrap();
4880                let mut edited = case.original.to_string();
4881                edited.replace_range(edit_range, &replacement);
4882                assert_eq!(edited, case.expected, "{}", case.name);
4883            }
4884        }
4885
4886        #[test]
4887        fn test_patch_to_variable_edit() {
4888            struct Case {
4889                name: &'static str,
4890                old: &'static str,
4891                patch: &'static str,
4892                cursor_offset: Option<usize>,
4893                expected_variable_edit: &'static str,
4894                expected_after_apply: &'static str,
4895            }
4896
4897            let cases = [
4898                Case {
4899                    name: "simple_replacement",
4900                    old: indoc! {"
4901                        zero
4902                        one
4903                        two
4904                        three
4905                        four
4906                        five
4907                    "},
4908                    patch: indoc! {"
4909                        @@ -3,3 +3,3 @@
4910                         two
4911                        -three
4912                        +THREE
4913                         four
4914                    "},
4915                    cursor_offset: None,
4916                    expected_variable_edit: indoc! {"
4917                        one
4918                        two
4919                        <|fim_middle|>
4920                        THREE
4921                        <|fim_suffix|>
4922                        four
4923                        five
4924                    "},
4925                    expected_after_apply: indoc! {"
4926                        zero
4927                        one
4928                        two
4929                        THREE
4930                        four
4931                        five
4932                    "},
4933                },
4934                Case {
4935                    name: "insertion",
4936                    old: indoc! {"
4937                        a
4938                        b
4939                        c
4940                        d
4941                        e
4942                    "},
4943                    patch: indoc! {"
4944                        @@ -2,0 +3,1 @@
4945                         b
4946                        +X
4947                         c
4948                    "},
4949                    cursor_offset: None,
4950                    expected_variable_edit: indoc! {"
4951                        a
4952                        b
4953                        <|fim_middle|>
4954                        X
4955                        <|fim_suffix|>
4956                        c
4957                        d
4958                    "},
4959                    expected_after_apply: indoc! {"
4960                        a
4961                        b
4962                        X
4963                        c
4964                        d
4965                        e
4966                    "},
4967                },
4968                Case {
4969                    name: "deletion",
4970                    old: indoc! {"
4971                        a
4972                        b
4973                        c
4974                        d
4975                        e
4976                    "},
4977                    patch: indoc! {"
4978                        @@ -2,3 +2,2 @@
4979                         b
4980                        -c
4981                         d
4982                    "},
4983                    cursor_offset: None,
4984                    expected_variable_edit: indoc! {"
4985                        a
4986                        b
4987                        <|fim_middle|>
4988                        <|fim_suffix|>
4989                        d
4990                        e
4991                    "},
4992                    expected_after_apply: indoc! {"
4993                        a
4994                        b
4995                        d
4996                        e
4997                    "},
4998                },
4999                Case {
5000                    name: "edit_near_start",
5001                    old: indoc! {"
5002                        first
5003                        second
5004                        third
5005                        fourth
5006                    "},
5007                    patch: indoc! {"
5008                        @@ -1,1 +1,1 @@
5009                        -first
5010                        +FIRST
5011                    "},
5012                    cursor_offset: None,
5013                    expected_variable_edit: indoc! {"
5014                        <|fim_middle|>
5015                        FIRST
5016                        <|fim_suffix|>
5017                        second
5018                        third
5019                    "},
5020                    expected_after_apply: indoc! {"
5021                        FIRST
5022                        second
5023                        third
5024                        fourth
5025                    "},
5026                },
5027                Case {
5028                    name: "edit_near_end",
5029                    old: indoc! {"
5030                        first
5031                        second
5032                        third
5033                        fourth
5034                    "},
5035                    patch: indoc! {"
5036                        @@ -4,1 +4,1 @@
5037                        -fourth
5038                        +FOURTH
5039                    "},
5040                    cursor_offset: None,
5041                    expected_variable_edit: indoc! {"
5042                        second
5043                        third
5044                        <|fim_middle|>
5045                        FOURTH
5046                        <|fim_suffix|>
5047                    "},
5048                    expected_after_apply: indoc! {"
5049                        first
5050                        second
5051                        third
5052                        FOURTH
5053                    "},
5054                },
5055                Case {
5056                    name: "cursor_at_start_of_replacement",
5057                    old: indoc! {"
5058                        zero
5059                        one
5060                        two
5061                        three
5062                        four
5063                        five
5064                    "},
5065                    patch: indoc! {"
5066                        @@ -3,3 +3,3 @@
5067                         two
5068                        -three
5069                        +THREE
5070                         four
5071                    "},
5072                    cursor_offset: Some(4),
5073                    expected_variable_edit: indoc! {"
5074                        one
5075                        two
5076                        <|fim_middle|>
5077                        <|user_cursor|>THREE
5078                        <|fim_suffix|>
5079                        four
5080                        five
5081                    "},
5082                    expected_after_apply: indoc! {"
5083                        zero
5084                        one
5085                        two
5086                        <|user_cursor|>THREE
5087                        four
5088                        five
5089                    "},
5090                },
5091                Case {
5092                    name: "cursor_in_middle_of_replacement",
5093                    old: indoc! {"
5094                        zero
5095                        one
5096                        two
5097                        three
5098                        four
5099                        five
5100                    "},
5101                    patch: indoc! {"
5102                        @@ -3,3 +3,3 @@
5103                         two
5104                        -three
5105                        +THREE
5106                         four
5107                    "},
5108                    cursor_offset: Some(6),
5109                    expected_variable_edit: indoc! {"
5110                        one
5111                        two
5112                        <|fim_middle|>
5113                        TH<|user_cursor|>REE
5114                        <|fim_suffix|>
5115                        four
5116                        five
5117                    "},
5118                    expected_after_apply: indoc! {"
5119                        zero
5120                        one
5121                        two
5122                        TH<|user_cursor|>REE
5123                        four
5124                        five
5125                    "},
5126                },
5127                Case {
5128                    name: "expands_context_when_two_lines_not_unique_before_and_after",
5129                    old: indoc! {"
5130                        one
5131                        a
5132                        b
5133                        c
5134                        d
5135                        two
5136                        a
5137                        b
5138                        c
5139                        d
5140                        three
5141                        a
5142                        b
5143                        c
5144                        d
5145                        four
5146                    "},
5147                    patch: indoc! {"
5148                        @@ -4,5 +4,5 @@
5149                         two
5150                         a
5151                         b
5152                        -c
5153                        +C
5154                         d
5155                         three
5156                    "},
5157                    cursor_offset: None,
5158                    expected_variable_edit: indoc! {"
5159                        two
5160                        a
5161                        b
5162                        <|fim_middle|>
5163                        C
5164                        <|fim_suffix|>
5165                        d
5166                        three
5167                    "},
5168                    expected_after_apply: indoc! {"
5169                        one
5170                        a
5171                        b
5172                        c
5173                        d
5174                        two
5175                        a
5176                        b
5177                        C
5178                        d
5179                        three
5180                        a
5181                        b
5182                        c
5183                        d
5184                        four
5185                    "},
5186                },
5187                Case {
5188                    name: "expands_context_when_two_lines_not_unique_before_and_after",
5189                    old: indoc! {"
5190                        {
5191                            {
5192                                one();
5193                            }
5194                        }
5195                        {
5196                            {
5197                                two();
5198                            }
5199                        }
5200                        {
5201                            {
5202                                three();
5203                            }
5204                        }
5205                        {
5206                            {
5207                                four();
5208                            }
5209                        }
5210                    "},
5211                    patch: indoc! {"
5212                        @@ -4,5 +4,5 @@
5213                             {
5214                        -        two();
5215                        +        TWO();
5216                             }
5217                    "},
5218                    cursor_offset: None,
5219                    expected_variable_edit: indoc! {"
5220                                one();
5221                            }
5222                        }
5223                        {
5224                            {
5225                        <|fim_middle|>
5226                                TWO();
5227                        <|fim_suffix|>
5228                            }
5229                        }
5230                        {
5231                            {
5232                                three();
5233                    "},
5234                    expected_after_apply: indoc! {"
5235                        {
5236                            {
5237                                one();
5238                            }
5239                        }
5240                        {
5241                            {
5242                                TWO();
5243                            }
5244                        }
5245                        {
5246                            {
5247                                three();
5248                            }
5249                        }
5250                        {
5251                            {
5252                                four();
5253                            }
5254                        }
5255                    "},
5256                },
5257            ];
5258
5259            for case in cases {
5260                let output =
5261                    patch_to_variable_edit_output(case.old, case.patch, case.cursor_offset)
5262                        .unwrap_or_else(|error| {
5263                            panic!("failed converting patch for {}: {error}", case.name)
5264                        });
5265                assert_eq!(
5266                    output, case.expected_variable_edit,
5267                    "patch->variable_edit mismatch for {}",
5268                    case.name
5269                );
5270
5271                let (edit_range, replacement) = apply_variable_edit(case.old, &output)
5272                    .unwrap_or_else(|error| {
5273                        panic!("failed applying variable_edit for {}: {error}", case.name)
5274                    });
5275                let mut edited_by_variable_edit = case.old.to_string();
5276                edited_by_variable_edit.replace_range(edit_range, &replacement);
5277                assert_eq!(
5278                    edited_by_variable_edit, case.expected_after_apply,
5279                    "variable_edit apply mismatch for {}",
5280                    case.name
5281                );
5282
5283                let (expected_edit_range, expected_replacement) =
5284                    apply_variable_edit(case.old, case.expected_variable_edit).unwrap_or_else(
5285                        |error| {
5286                            panic!(
5287                                "failed applying expected variable_edit for {}: {error}",
5288                                case.name
5289                            )
5290                        },
5291                    );
5292                let mut edited_by_expected_variable_edit = case.old.to_string();
5293                edited_by_expected_variable_edit
5294                    .replace_range(expected_edit_range, &expected_replacement);
5295                assert_eq!(
5296                    edited_by_expected_variable_edit, case.expected_after_apply,
5297                    "expected variable_edit apply mismatch for {}",
5298                    case.name
5299                );
5300            }
5301        }
5302
5303        #[test]
5304        fn test_write_cursor_excerpt_section() {
5305            let path = Path::new("test.rs");
5306            let context = "fn main() {\n    hello();\n}\n";
5307            let cursor_offset = 17;
5308            let mut prompt = String::new();
5309            write_cursor_excerpt_section(&mut prompt, path, context, cursor_offset);
5310            assert_eq!(
5311                prompt,
5312                "<|file_sep|>test.rs\nfn main() {\n    h<|user_cursor|>ello();\n}\n<|fim_prefix|>\n"
5313            );
5314        }
5315    }
5316}
5317
5318/// The zeta1 prompt format
5319pub mod zeta1 {
5320    use super::*;
5321    use std::fmt::Write;
5322
5323    pub const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
5324    pub const START_OF_FILE_MARKER: &str = "<|start_of_file|>";
5325    pub const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>";
5326    pub const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>";
5327
5328    const INSTRUCTION_HEADER: &str = concat!(
5329        "### Instruction:\n",
5330        "You are a code completion assistant and your task is to analyze user edits and then rewrite an ",
5331        "excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking ",
5332        "into account the cursor location.\n\n",
5333        "### User Edits:\n\n"
5334    );
5335    const EXCERPT_HEADER: &str = "\n\n### User Excerpt:\n\n";
5336    const RESPONSE_HEADER: &str = "\n\n### Response:\n";
5337
5338    /// Formats a complete zeta1 prompt from the input events and excerpt.
5339    pub fn format_zeta1_prompt(input_events: &str, input_excerpt: &str) -> String {
5340        let mut prompt = String::with_capacity(
5341            INSTRUCTION_HEADER.len()
5342                + input_events.len()
5343                + EXCERPT_HEADER.len()
5344                + input_excerpt.len()
5345                + RESPONSE_HEADER.len(),
5346        );
5347        prompt.push_str(INSTRUCTION_HEADER);
5348        prompt.push_str(input_events);
5349        prompt.push_str(EXCERPT_HEADER);
5350        prompt.push_str(input_excerpt);
5351        prompt.push_str(RESPONSE_HEADER);
5352        prompt
5353    }
5354
5355    /// Formats a complete zeta1 prompt from a `Zeta2PromptInput` using the given
5356    /// editable and context byte-offset ranges within `cursor_excerpt`.
5357    pub fn format_zeta1_from_input(
5358        input: &Zeta2PromptInput,
5359        editable_range: Range<usize>,
5360        context_range: Range<usize>,
5361    ) -> String {
5362        let events = format_zeta1_events(&input.events);
5363        let excerpt = format_zeta1_excerpt(input, editable_range, context_range);
5364        format_zeta1_prompt(&events, &excerpt)
5365    }
5366
5367    /// Formats events in zeta1 style (oldest first).
5368    fn format_zeta1_events(events: &[Arc<Event>]) -> String {
5369        let mut result = String::new();
5370        for event in
5371            events
5372                .iter()
5373                .skip(events.len().saturating_sub(max_edit_event_count_for_format(
5374                    &ZetaFormat::V0114180EditableRegion,
5375                )))
5376        {
5377            let event_string = format_zeta1_event(event);
5378            if event_string.is_empty() {
5379                continue;
5380            }
5381            if !result.is_empty() {
5382                result.push_str("\n\n");
5383            }
5384            result.push_str(&event_string);
5385        }
5386        result
5387    }
5388
5389    fn format_zeta1_event(event: &Event) -> String {
5390        match event {
5391            Event::BufferChange {
5392                path,
5393                old_path,
5394                diff,
5395                ..
5396            } => {
5397                let mut prompt = String::new();
5398                if old_path != path {
5399                    writeln!(
5400                        prompt,
5401                        "User renamed {} to {}\n",
5402                        old_path.display(),
5403                        path.display()
5404                    )
5405                    .ok();
5406                }
5407                if !diff.is_empty() {
5408                    write!(
5409                        prompt,
5410                        "User edited {}:\n```diff\n{}\n```",
5411                        path.display(),
5412                        diff
5413                    )
5414                    .ok();
5415                }
5416                prompt
5417            }
5418        }
5419    }
5420
5421    /// Formats the excerpt section of a zeta1 prompt using byte-offset ranges
5422    /// within `cursor_excerpt`.
5423    fn format_zeta1_excerpt(
5424        input: &Zeta2PromptInput,
5425        editable_range: Range<usize>,
5426        context_range: Range<usize>,
5427    ) -> String {
5428        let path_str = input.cursor_path.to_string_lossy();
5429        let excerpt = &*input.cursor_excerpt;
5430        let cursor_offset = input.cursor_offset_in_excerpt;
5431
5432        let mut prompt = String::new();
5433        writeln!(&mut prompt, "```{path_str}").ok();
5434
5435        let starts_at_file_beginning =
5436            input.excerpt_start_row == Some(0) && context_range.start == 0;
5437        if starts_at_file_beginning {
5438            writeln!(&mut prompt, "{START_OF_FILE_MARKER}").ok();
5439        }
5440
5441        prompt.push_str(&excerpt[context_range.start..editable_range.start]);
5442
5443        writeln!(&mut prompt, "{EDITABLE_REGION_START_MARKER}").ok();
5444        prompt.push_str(&excerpt[editable_range.start..cursor_offset]);
5445        prompt.push_str(CURSOR_MARKER);
5446        prompt.push_str(&excerpt[cursor_offset..editable_range.end]);
5447        write!(&mut prompt, "\n{EDITABLE_REGION_END_MARKER}").ok();
5448
5449        prompt.push_str(&excerpt[editable_range.end..context_range.end]);
5450        write!(prompt, "\n```").ok();
5451
5452        prompt
5453    }
5454
5455    /// Cleans zeta1 model output by extracting content between editable region
5456    /// markers and converting the zeta1 cursor marker to the universal one.
5457    /// Returns `None` if the output doesn't contain the expected markers.
5458    pub fn clean_zeta1_model_output(output: &str) -> Option<String> {
5459        let content = output.replace(CURSOR_MARKER, "");
5460
5461        let content_start = content
5462            .find(EDITABLE_REGION_START_MARKER)
5463            .map(|pos| pos + EDITABLE_REGION_START_MARKER.len())
5464            .map(|pos| {
5465                if content.as_bytes().get(pos) == Some(&b'\n') {
5466                    pos + 1
5467                } else {
5468                    pos
5469                }
5470            })
5471            .unwrap_or(0);
5472
5473        let content_end = content
5474            .find(EDITABLE_REGION_END_MARKER)
5475            .map(|pos| {
5476                if pos > 0 && content.as_bytes().get(pos - 1) == Some(&b'\n') {
5477                    pos - 1
5478                } else {
5479                    pos
5480                }
5481            })
5482            .unwrap_or(content.len());
5483
5484        if content_start > content_end {
5485            return Some(String::new());
5486        }
5487
5488        let extracted = &content[content_start..content_end];
5489
5490        let cursor_offset = output.find(CURSOR_MARKER).map(|zeta1_cursor_pos| {
5491            let text_before_cursor = output[..zeta1_cursor_pos].replace(CURSOR_MARKER, "");
5492            let text_before_cursor = text_before_cursor
5493                .find(EDITABLE_REGION_START_MARKER)
5494                .map(|pos| {
5495                    let after_marker = pos + EDITABLE_REGION_START_MARKER.len();
5496                    if text_before_cursor.as_bytes().get(after_marker) == Some(&b'\n') {
5497                        after_marker + 1
5498                    } else {
5499                        after_marker
5500                    }
5501                })
5502                .unwrap_or(0);
5503            let offset_in_extracted = zeta1_cursor_pos
5504                .saturating_sub(text_before_cursor)
5505                .min(extracted.len());
5506            offset_in_extracted
5507        });
5508
5509        let mut result = String::with_capacity(extracted.len() + super::CURSOR_MARKER.len());
5510        if let Some(offset) = cursor_offset {
5511            result.push_str(&extracted[..offset]);
5512            result.push_str(super::CURSOR_MARKER);
5513            result.push_str(&extracted[offset..]);
5514        } else {
5515            result.push_str(extracted);
5516        }
5517
5518        Some(result)
5519    }
5520}
5521
5522#[cfg(test)]
5523mod tests {
5524    use super::*;
5525    use indoc::indoc;
5526
5527    fn make_input(
5528        cursor_excerpt: &str,
5529        editable_range: Range<usize>,
5530        cursor_offset: usize,
5531        events: Vec<Event>,
5532        related_files: Vec<RelatedFile>,
5533    ) -> Zeta2PromptInput {
5534        let context_range = 0..cursor_excerpt.len();
5535        Zeta2PromptInput {
5536            cursor_path: Path::new("test.rs").into(),
5537            cursor_excerpt: cursor_excerpt.into(),
5538            cursor_offset_in_excerpt: cursor_offset,
5539            excerpt_start_row: None,
5540            events: events.into_iter().map(Arc::new).collect(),
5541            related_files: Some(related_files),
5542            active_buffer_diagnostics: vec![],
5543            excerpt_ranges: ExcerptRanges {
5544                editable_150: editable_range.clone(),
5545                editable_180: editable_range.clone(),
5546                editable_350: editable_range,
5547                editable_150_context_350: context_range.clone(),
5548                editable_180_context_350: context_range.clone(),
5549                editable_350_context_150: context_range,
5550                ..Default::default()
5551            },
5552            syntax_ranges: None,
5553            in_open_source_repo: false,
5554            can_collect_data: false,
5555            repo_url: None,
5556        }
5557    }
5558
5559    fn make_input_with_context_range(
5560        excerpt: &str,
5561        editable_range: Range<usize>,
5562        context_range: Range<usize>,
5563        cursor_offset: usize,
5564    ) -> Zeta2PromptInput {
5565        Zeta2PromptInput {
5566            cursor_path: Path::new("test.rs").into(),
5567            cursor_excerpt: excerpt.into(),
5568            cursor_offset_in_excerpt: cursor_offset,
5569            excerpt_start_row: None,
5570            events: vec![],
5571            related_files: Some(vec![]),
5572            active_buffer_diagnostics: vec![],
5573            excerpt_ranges: ExcerptRanges {
5574                editable_150: editable_range.clone(),
5575                editable_180: editable_range.clone(),
5576                editable_350: editable_range,
5577                editable_150_context_350: context_range.clone(),
5578                editable_180_context_350: context_range.clone(),
5579                editable_350_context_150: context_range,
5580                ..Default::default()
5581            },
5582            syntax_ranges: None,
5583            in_open_source_repo: false,
5584            can_collect_data: false,
5585            repo_url: None,
5586        }
5587    }
5588
5589    fn make_event(path: &str, diff: &str) -> Event {
5590        Event::BufferChange {
5591            path: Path::new(path).into(),
5592            old_path: Path::new(path).into(),
5593            diff: diff.to_string(),
5594            old_range: 0..diff.len(),
5595            new_range: 0..diff.len(),
5596            predicted: false,
5597            in_open_source_repo: false,
5598        }
5599    }
5600
5601    fn make_related_file(path: &str, content: &str) -> RelatedFile {
5602        RelatedFile {
5603            path: Path::new(path).into(),
5604            max_row: content.lines().count() as u32,
5605            excerpts: vec![RelatedExcerpt {
5606                row_range: 0..content.lines().count() as u32,
5607                text: content.into(),
5608                order: 0,
5609                context_source: ContextSource::Lsp,
5610            }],
5611            in_open_source_repo: false,
5612        }
5613    }
5614
5615    fn format_with_budget(input: &Zeta2PromptInput, max_tokens: usize) -> Option<String> {
5616        format_prompt_with_budget_for_format(input, ZetaFormat::V0114180EditableRegion, max_tokens)
5617    }
5618
5619    fn budget_with_margin(requested_tokens: usize) -> usize {
5620        ((requested_tokens as f64) / 0.9).ceil() as usize
5621    }
5622
5623    #[test]
5624    fn test_no_truncation_when_within_budget() {
5625        let input = make_input(
5626            "prefix\neditable\nsuffix",
5627            7..15,
5628            10,
5629            vec![make_event("a.rs", "-old\n+new\n")],
5630            vec![make_related_file("related.rs", "fn helper() {}\n")],
5631        );
5632
5633        assert_eq!(
5634            format_with_budget(&input, 10000).unwrap(),
5635            indoc! {r#"
5636                <|file_sep|>related.rs
5637                fn helper() {}
5638                <|file_sep|>edit history
5639                --- a/a.rs
5640                +++ b/a.rs
5641                -old
5642                +new
5643                <|file_sep|>test.rs
5644                <|fim_prefix|>
5645                prefix
5646                <|fim_middle|>current
5647                edi<|user_cursor|>table
5648                <|fim_suffix|>
5649
5650                suffix
5651                <|fim_middle|>updated
5652            "#}
5653            .to_string()
5654        );
5655    }
5656
5657    #[test]
5658    fn test_truncation_drops_edit_history_when_budget_tight() {
5659        let input = make_input(
5660            "code",
5661            0..4,
5662            2,
5663            vec![make_event("a.rs", "-x\n+y\n")],
5664            vec![
5665                make_related_file("r1.rs", "aaaaaaa\n"),
5666                make_related_file("r2.rs", "bbbbbbb\n"),
5667            ],
5668        );
5669
5670        assert_eq!(
5671            format_with_budget(&input, 10000).unwrap(),
5672            indoc! {r#"
5673                <|file_sep|>r1.rs
5674                aaaaaaa
5675                <|file_sep|>r2.rs
5676                bbbbbbb
5677                <|file_sep|>edit history
5678                --- a/a.rs
5679                +++ b/a.rs
5680                -x
5681                +y
5682                <|file_sep|>test.rs
5683                <|fim_prefix|>
5684                <|fim_middle|>current
5685                co<|user_cursor|>de
5686                <|fim_suffix|>
5687                <|fim_middle|>updated
5688            "#}
5689            .to_string()
5690        );
5691
5692        assert_eq!(
5693            format_with_budget(&input, budget_with_margin(55)),
5694            Some(
5695                indoc! {r#"
5696                <|file_sep|>edit history
5697                --- a/a.rs
5698                +++ b/a.rs
5699                -x
5700                +y
5701                <|file_sep|>test.rs
5702                <|fim_prefix|>
5703                <|fim_middle|>current
5704                co<|user_cursor|>de
5705                <|fim_suffix|>
5706                <|fim_middle|>updated
5707            "#}
5708                .to_string()
5709            )
5710        );
5711    }
5712
5713    #[test]
5714    fn test_truncation_includes_partial_excerpts() {
5715        let input = make_input(
5716            "x",
5717            0..1,
5718            0,
5719            vec![],
5720            vec![RelatedFile {
5721                path: Path::new("big.rs").into(),
5722                max_row: 30,
5723                in_open_source_repo: false,
5724                excerpts: vec![
5725                    RelatedExcerpt {
5726                        row_range: 0..10,
5727                        text: "first excerpt\n".into(),
5728                        order: 0,
5729                        context_source: ContextSource::Lsp,
5730                    },
5731                    RelatedExcerpt {
5732                        row_range: 11..20,
5733                        text: "second excerpt\n".into(),
5734                        order: 0,
5735                        context_source: ContextSource::Lsp,
5736                    },
5737                    RelatedExcerpt {
5738                        row_range: 21..30,
5739                        text: "third excerpt\n".into(),
5740                        order: 0,
5741                        context_source: ContextSource::Lsp,
5742                    },
5743                ],
5744            }],
5745        );
5746
5747        assert_eq!(
5748            format_with_budget(&input, 10000).unwrap(),
5749            indoc! {r#"
5750                <|file_sep|>big.rs
5751                first excerpt
5752                ...
5753                second excerpt
5754                ...
5755                third excerpt
5756                <|file_sep|>test.rs
5757                <|fim_prefix|>
5758                <|fim_middle|>current
5759                <|user_cursor|>x
5760                <|fim_suffix|>
5761                <|fim_middle|>updated
5762            "#}
5763            .to_string()
5764        );
5765
5766        assert_eq!(
5767            format_with_budget(&input, budget_with_margin(50)).unwrap(),
5768            indoc! {r#"
5769                <|file_sep|>big.rs
5770                first excerpt
5771                ...
5772                <|file_sep|>test.rs
5773                <|fim_prefix|>
5774                <|fim_middle|>current
5775                <|user_cursor|>x
5776                <|fim_suffix|>
5777                <|fim_middle|>updated
5778            "#}
5779            .to_string()
5780        );
5781    }
5782
5783    #[test]
5784    fn test_contiguous_excerpts_render_without_ellipsis() {
5785        // Excerpts whose row ranges touch (end == next start) are contiguous
5786        // segments of the same region and must render seamlessly, without an
5787        // ellipsis line between them.
5788        let input = make_input(
5789            "x",
5790            0..1,
5791            0,
5792            vec![],
5793            vec![RelatedFile {
5794                path: Path::new("big.rs").into(),
5795                max_row: 30,
5796                in_open_source_repo: false,
5797                excerpts: vec![
5798                    RelatedExcerpt {
5799                        row_range: 0..10,
5800                        text: "first segment\n".into(),
5801                        order: 1,
5802                        context_source: ContextSource::GitLog,
5803                    },
5804                    RelatedExcerpt {
5805                        row_range: 10..20,
5806                        text: "second segment\n".into(),
5807                        order: 0,
5808                        context_source: ContextSource::OracleSnippet,
5809                    },
5810                ],
5811            }],
5812        );
5813
5814        assert_eq!(
5815            format_with_budget(&input, 10000).unwrap(),
5816            indoc! {r#"
5817                <|file_sep|>big.rs
5818                first segment
5819                second segment
5820                ...
5821                <|file_sep|>test.rs
5822                <|fim_prefix|>
5823                <|fim_middle|>current
5824                <|user_cursor|>x
5825                <|fim_suffix|>
5826                <|fim_middle|>updated
5827            "#}
5828            .to_string()
5829        );
5830    }
5831
5832    #[test]
5833    fn test_truncation_prioritizes_lower_order_excerpts() {
5834        // Two files: file_a has a high-order excerpt, file_b has a low-order one.
5835        // With tight budget, only the lower-order excerpt from file_b should be included.
5836        let input = make_input(
5837            "x",
5838            0..1,
5839            0,
5840            vec![],
5841            vec![
5842                RelatedFile {
5843                    path: Path::new("file_a.rs").into(),
5844                    max_row: 10,
5845                    in_open_source_repo: false,
5846                    excerpts: vec![RelatedExcerpt {
5847                        row_range: 0..10,
5848                        text: "low priority content\n".into(),
5849                        order: 5,
5850                        context_source: ContextSource::Lsp,
5851                    }],
5852                },
5853                RelatedFile {
5854                    path: Path::new("file_b.rs").into(),
5855                    max_row: 10,
5856                    in_open_source_repo: false,
5857                    excerpts: vec![RelatedExcerpt {
5858                        row_range: 0..10,
5859                        text: "high priority content\n".into(),
5860                        order: 1,
5861                        context_source: ContextSource::Lsp,
5862                    }],
5863                },
5864            ],
5865        );
5866
5867        // With large budget, both files included; rendered in stable lexicographic order.
5868        assert_eq!(
5869            format_with_budget(&input, 10000).unwrap(),
5870            indoc! {r#"
5871                <|file_sep|>file_a.rs
5872                low priority content
5873                <|file_sep|>file_b.rs
5874                high priority content
5875                <|file_sep|>test.rs
5876                <|fim_prefix|>
5877                <|fim_middle|>current
5878                <|user_cursor|>x
5879                <|fim_suffix|>
5880                <|fim_middle|>updated
5881            "#}
5882            .to_string()
5883        );
5884
5885        // With tight budget, only file_b (lower order) fits.
5886        // Cursor section is ~37 tokens, so budget 52 leaves ~15 for related files.
5887        // file_b header (7) + excerpt (7) = 14 tokens, which fits.
5888        // file_a would need another 14 tokens, which doesn't fit.
5889        assert_eq!(
5890            format_with_budget(&input, budget_with_margin(52)).unwrap(),
5891            indoc! {r#"
5892                <|file_sep|>file_b.rs
5893                high priority content
5894                <|file_sep|>test.rs
5895                <|fim_prefix|>
5896                <|fim_middle|>current
5897                <|user_cursor|>x
5898                <|fim_suffix|>
5899                <|fim_middle|>updated
5900            "#}
5901            .to_string()
5902        );
5903    }
5904
5905    #[test]
5906    fn test_truncation_drops_high_order_excerpts_within_file() {
5907        // A single file has excerpts at order 1 and order 3. With a tight budget,
5908        // only the order-1 excerpts are included while the order-3 excerpt is
5909        // dropped — even though they belong to the same file. This also preserves
5910        // the parent invariant: parent outline items have order ≤ their best
5911        // child, so they're always included when any child is.
5912        let input = make_input(
5913            "x",
5914            0..1,
5915            0,
5916            vec![],
5917            vec![RelatedFile {
5918                path: Path::new("mod.rs").into(),
5919                max_row: 30,
5920                in_open_source_repo: false,
5921                excerpts: vec![
5922                    RelatedExcerpt {
5923                        row_range: 0..5,
5924                        text: "mod header\n".into(),
5925                        order: 1,
5926                        context_source: ContextSource::Lsp,
5927                    },
5928                    RelatedExcerpt {
5929                        row_range: 6..15,
5930                        text: "important fn\n".into(),
5931                        order: 1,
5932                        context_source: ContextSource::Lsp,
5933                    },
5934                    RelatedExcerpt {
5935                        row_range: 16..30,
5936                        text: "less important fn\n".into(),
5937                        order: 3,
5938                        context_source: ContextSource::Lsp,
5939                    },
5940                ],
5941            }],
5942        );
5943
5944        // With large budget, all three excerpts included.
5945        assert_eq!(
5946            format_with_budget(&input, 10000).unwrap(),
5947            indoc! {r#"
5948                <|file_sep|>mod.rs
5949                mod header
5950                ...
5951                important fn
5952                ...
5953                less important fn
5954                <|file_sep|>test.rs
5955                <|fim_prefix|>
5956                <|fim_middle|>current
5957                <|user_cursor|>x
5958                <|fim_suffix|>
5959                <|fim_middle|>updated
5960            "#}
5961            .to_string()
5962        );
5963
5964        // With tight budget, only order<=1 excerpts included (header + important fn).
5965        assert_eq!(
5966            format_with_budget(&input, budget_with_margin(55)).unwrap(),
5967            indoc! {r#"
5968                <|file_sep|>mod.rs
5969                mod header
5970                ...
5971                important fn
5972                ...
5973                <|file_sep|>test.rs
5974                <|fim_prefix|>
5975                <|fim_middle|>current
5976                <|user_cursor|>x
5977                <|fim_suffix|>
5978                <|fim_middle|>updated
5979            "#}
5980            .to_string()
5981        );
5982    }
5983
5984    #[test]
5985    fn test_truncation_drops_older_events_first() {
5986        let input = make_input(
5987            "x",
5988            0..1,
5989            0,
5990            vec![make_event("old.rs", "-1\n"), make_event("new.rs", "-2\n")],
5991            vec![],
5992        );
5993
5994        assert_eq!(
5995            format_with_budget(&input, 10000).unwrap(),
5996            indoc! {r#"
5997                <|file_sep|>edit history
5998                --- a/old.rs
5999                +++ b/old.rs
6000                -1
6001                --- a/new.rs
6002                +++ b/new.rs
6003                -2
6004                <|file_sep|>test.rs
6005                <|fim_prefix|>
6006                <|fim_middle|>current
6007                <|user_cursor|>x
6008                <|fim_suffix|>
6009                <|fim_middle|>updated
6010            "#}
6011            .to_string()
6012        );
6013
6014        assert_eq!(
6015            format_with_budget(&input, 60).unwrap(),
6016            indoc! {r#"
6017                <|file_sep|>edit history
6018                --- a/new.rs
6019                +++ b/new.rs
6020                -2
6021                <|file_sep|>test.rs
6022                <|fim_prefix|>
6023                <|fim_middle|>current
6024                <|user_cursor|>x
6025                <|fim_suffix|>
6026                <|fim_middle|>updated
6027            "#}
6028            .to_string()
6029        );
6030    }
6031
6032    #[test]
6033    fn test_cursor_excerpt_always_included_with_minimal_budget() {
6034        let input = make_input(
6035            "fn main() {}",
6036            0..12,
6037            3,
6038            vec![make_event("a.rs", "-old\n+new\n")],
6039            vec![make_related_file("related.rs", "helper\n")],
6040        );
6041
6042        assert!(format_with_budget(&input, 30).is_none())
6043    }
6044
6045    #[track_caller]
6046    fn format_seed_coder(input: &Zeta2PromptInput) -> String {
6047        format_prompt_with_budget_for_format(input, ZetaFormat::V0211SeedCoder, 10000)
6048            .expect("seed coder prompt formatting should succeed")
6049    }
6050
6051    #[track_caller]
6052    fn format_seed_coder_with_budget(input: &Zeta2PromptInput, max_tokens: usize) -> String {
6053        format_prompt_with_budget_for_format(input, ZetaFormat::V0211SeedCoder, max_tokens)
6054            .expect("seed coder prompt formatting should succeed")
6055    }
6056
6057    #[test]
6058    fn test_seed_coder_alias_matches_v0211_seed_coder() {
6059        let input = make_input(
6060            "prefix\neditable\nsuffix",
6061            7..15,
6062            10,
6063            vec![make_event("a.rs", "-old\n+new\n")],
6064            vec![make_related_file("related.rs", "fn helper() {}\n")],
6065        );
6066
6067        assert_eq!(
6068            format_prompt_with_budget_for_format(&input, ZetaFormat::V0211SeedCoder, 10000),
6069            format_prompt_with_budget_for_format(&input, ZetaFormat::V0331SeedCoderModelPy, 10000)
6070        );
6071        assert_eq!(
6072            ZetaFormat::parse("V0331SeedCoderModelPy").unwrap(),
6073            ZetaFormat::V0331SeedCoderModelPy
6074        );
6075    }
6076
6077    #[test]
6078    fn test_seed_coder_basic_format() {
6079        let input = make_input(
6080            "prefix\neditable\nsuffix",
6081            7..15,
6082            10,
6083            vec![make_event("a.rs", "-old\n+new\n")],
6084            vec![make_related_file("related.rs", "fn helper() {}\n")],
6085        );
6086
6087        assert_eq!(
6088            format_seed_coder(&input),
6089            indoc! {r#"
6090                <[fim-suffix]>
6091                suffix
6092                <[fim-prefix]><filename>related.rs
6093                fn helper() {}
6094
6095                <filename>edit_history
6096                --- a/a.rs
6097                +++ b/a.rs
6098                -old
6099                +new
6100
6101                <filename>test.rs
6102                prefix
6103                <<<<<<< CURRENT
6104                edi<|user_cursor|>table
6105                =======
6106                <[fim-middle]>"#}
6107        );
6108    }
6109
6110    #[test]
6111    fn test_qwen36_multi_region_uses_qwen_psm_fim_format() {
6112        let input = make_input(
6113            "prefix\neditable\nsuffix",
6114            7..15,
6115            10,
6116            vec![make_event("a.rs", "-old\n+new\n")],
6117            vec![make_related_file("related.rs", "fn helper() {}\n")],
6118        );
6119
6120        assert_eq!(
6121            format_prompt_with_budget_for_format(&input, ZetaFormat::V0608QwenMultiRegions, 10000)
6122                .expect("qwen prompt formatting should succeed"),
6123            indoc! {r#"
6124                <|fim_prefix|><|file_sep|>related.rs
6125                fn helper() {}
6126
6127                <|file_sep|>edit_history
6128                --- a/a.rs
6129                +++ b/a.rs
6130                -old
6131                +new
6132
6133                <|file_sep|>test.rs
6134                prefix
6135                <|marker_1|>edi<|user_cursor|>table<|marker_2|>
6136                <|fim_suffix|>
6137                suffix
6138                <|fim_middle|>"#}
6139        );
6140    }
6141
6142    #[test]
6143    fn test_v0420_formats_diagnostics_before_related_files() {
6144        let mut input = make_input(
6145            "prefix\neditable\nsuffix",
6146            7..15,
6147            10,
6148            vec![],
6149            vec![make_related_file("related.rs", "fn helper() {}\n")],
6150        );
6151        input.active_buffer_diagnostics = vec![
6152            ActiveBufferDiagnostic {
6153                severity: Some(1),
6154                message: "missing semicolon".to_string(),
6155                snippet: "let value = 1".to_string(),
6156                snippet_buffer_row_range: 1..2,
6157                diagnostic_range_in_snippet: 12..13,
6158            },
6159            ActiveBufferDiagnostic {
6160                severity: Some(2),
6161                message: "file-level warning".to_string(),
6162                snippet: String::new(),
6163                snippet_buffer_row_range: 0..0,
6164                diagnostic_range_in_snippet: 0..0,
6165            },
6166        ];
6167
6168        let prompt =
6169            format_prompt_with_budget_for_format(&input, ZetaFormat::V0420Diagnostics, 10000)
6170                .expect("v0420 prompt formatting should succeed");
6171
6172        assert_eq!(
6173            prompt,
6174            indoc! {r#"
6175                <[fim-suffix]>
6176                suffix
6177                <[fim-prefix]><filename>diagnostics
6178                *missing semicolon*:
6179                ```
6180                let value = 1
6181                ```
6182                *file-level warning*
6183
6184                <filename>related.rs
6185                fn helper() {}
6186
6187                <filename>test.rs
6188                prefix
6189                <|marker_1|>edi<|user_cursor|>table<|marker_2|>
6190                <[fim-middle]>"#}
6191        );
6192    }
6193
6194    #[test]
6195    fn test_v0317_formats_prompt_with_many_related_files() {
6196        let related_files = (0..900)
6197            .map(|index| {
6198                make_related_file(
6199                    &format!("related_{index}.rs"),
6200                    "fn helper() {\n    let value = 1;\n}\n",
6201                )
6202            })
6203            .collect();
6204
6205        let input = make_input(
6206            "code",
6207            0..4,
6208            2,
6209            vec![make_event("a.rs", "-x\n+y\n")],
6210            related_files,
6211        );
6212
6213        let prompt =
6214            format_prompt_with_budget_for_format(&input, ZetaFormat::V0317SeedMultiRegions, 4096);
6215
6216        assert!(prompt.is_some());
6217        let prompt = prompt.expect("v0317 should produce a prompt under high related-file count");
6218        assert!(prompt.contains("test.rs"));
6219        assert!(prompt.contains(CURSOR_MARKER));
6220    }
6221
6222    #[test]
6223    fn test_v0327_formats_single_file_prompt_without_related_files() {
6224        let excerpt = indoc! {"
6225            line01
6226            line02
6227            line03
6228            line04
6229            line05
6230            line06
6231            line07
6232            line08
6233            line09
6234            line10
6235            line11
6236            line12
6237            line13
6238            line14
6239            line15
6240            line16
6241            line17
6242            line18
6243            line19
6244            line20
6245        "};
6246        let cursor_offset = excerpt.find("line10").expect("cursor line exists");
6247        let input = make_input(
6248            excerpt,
6249            0..excerpt.len(),
6250            cursor_offset,
6251            vec![make_event("a.rs", "-x\n+y\n")],
6252            vec![make_related_file("related.rs", "fn helper() {}\n")],
6253        );
6254
6255        let prompt =
6256            format_prompt_with_budget_for_format(&input, ZetaFormat::V0327SingleFile, 4096)
6257                .expect("v0327 prompt should fit");
6258
6259        assert!(prompt.contains("line01"));
6260        assert!(prompt.contains("line20"));
6261        assert!(prompt.contains("<filename>edit_history"));
6262        assert!(prompt.contains("<filename>test.rs"));
6263        assert!(prompt.contains(CURSOR_MARKER));
6264        assert!(!prompt.contains("related.rs"));
6265        assert!(!prompt.contains("fn helper() {}"));
6266    }
6267
6268    #[test]
6269    fn test_v0327_resolve_cursor_region_uses_full_excerpt_context() {
6270        let excerpt = (0..80)
6271            .map(|index| format!("l{index:02}\n"))
6272            .collect::<String>();
6273        let cursor_offset = excerpt.find("l40").expect("cursor line exists");
6274        let input = make_input(&excerpt, 0..excerpt.len(), cursor_offset, vec![], vec![]);
6275
6276        let (context, editable_range, context_range, adjusted_cursor) =
6277            resolve_cursor_region(&input, ZetaFormat::V0327SingleFile);
6278
6279        assert_eq!(context, excerpt);
6280        assert_eq!(context_range, 0..excerpt.len());
6281        assert_eq!(adjusted_cursor, cursor_offset);
6282        assert!(editable_range.start < adjusted_cursor);
6283        assert!(editable_range.end > adjusted_cursor);
6284        assert!(editable_range.end < excerpt.len());
6285    }
6286
6287    #[test]
6288    fn test_v0615_formats_hashed_markers_for_rendered_related_context() {
6289        let current_text = "fn main() {\n    let value = 1;\n}\n";
6290        let cursor_offset = current_text.find("let value").unwrap();
6291        let input = Zeta2PromptInput {
6292            cursor_path: Path::new("test.rs").into(),
6293            cursor_excerpt: current_text.into(),
6294            cursor_offset_in_excerpt: cursor_offset,
6295            excerpt_start_row: Some(0),
6296            events: Vec::new(),
6297            related_files: Some(vec![
6298                RelatedFile {
6299                    path: Path::new("test.rs").into(),
6300                    max_row: 3,
6301                    excerpts: vec![RelatedExcerpt {
6302                        row_range: 0..3,
6303                        text: current_text.into(),
6304                        order: 0,
6305                        context_source: ContextSource::CurrentFile,
6306                    }],
6307                    in_open_source_repo: false,
6308                },
6309                RelatedFile {
6310                    path: Path::new("helper.rs").into(),
6311                    max_row: 3,
6312                    excerpts: vec![RelatedExcerpt {
6313                        row_range: 10..13,
6314                        text: "fn helper() {\n    one();\n}\n".into(),
6315                        order: 1,
6316                        context_source: ContextSource::EditHistory,
6317                    }],
6318                    in_open_source_repo: false,
6319                },
6320                RelatedFile {
6321                    path: Path::new("readonly.rs").into(),
6322                    max_row: 1,
6323                    excerpts: vec![RelatedExcerpt {
6324                        row_range: 0..1,
6325                        text: "pub fn readonly() {}\n".into(),
6326                        order: 2,
6327                        context_source: ContextSource::Lsp,
6328                    }],
6329                    in_open_source_repo: false,
6330                },
6331            ]),
6332            active_buffer_diagnostics: vec![],
6333            excerpt_ranges: ExcerptRanges::default(),
6334            syntax_ranges: None,
6335            in_open_source_repo: false,
6336            can_collect_data: false,
6337            repo_url: None,
6338        };
6339
6340        let prompt = format_zeta_prompt(&input, ZetaFormat::V0615HashRegions).unwrap();
6341        let marker_table = hashed_regions::build_marker_table(&input);
6342        let marker_count: usize = marker_table
6343            .iter()
6344            .map(|snippet| snippet.markers.len())
6345            .sum();
6346
6347        assert!(prompt.starts_with("<[fim-suffix]>\n<[fim-prefix]><filename>test.rs\n"));
6348        assert!(prompt.ends_with("<[fim-middle]>"));
6349        assert!(prompt.contains(CURSOR_MARKER));
6350        assert!(prompt.contains("<filename>helper.rs\n"));
6351        assert!(prompt.contains("<filename>readonly.rs\n"));
6352        assert!(prompt.contains("<filename>readonly.rs\n<|marker_"));
6353        assert!(prompt.contains("pub fn readonly() {}"));
6354        assert_eq!(
6355            prompt.matches(hashed_regions::MARKER_TAG_PREFIX).count(),
6356            marker_count
6357        );
6358        assert!(!prompt.contains(seed_coder::START_MARKER));
6359    }
6360
6361    #[test]
6362    fn test_v0615_parse_related_file_jump_as_patch() {
6363        let current_text = "fn main() {\n    helper();\n}\n";
6364        let helper_text = "fn helper() {\n    one();\n}\n";
6365        let input = Zeta2PromptInput {
6366            cursor_path: Path::new("test.rs").into(),
6367            cursor_excerpt: current_text.into(),
6368            cursor_offset_in_excerpt: current_text.find("helper").unwrap(),
6369            excerpt_start_row: Some(0),
6370            events: Vec::new(),
6371            related_files: Some(vec![
6372                RelatedFile {
6373                    path: Path::new("test.rs").into(),
6374                    max_row: 3,
6375                    excerpts: vec![RelatedExcerpt {
6376                        row_range: 0..3,
6377                        text: current_text.into(),
6378                        order: 0,
6379                        context_source: ContextSource::CurrentFile,
6380                    }],
6381                    in_open_source_repo: false,
6382                },
6383                RelatedFile {
6384                    path: Path::new("helper.rs").into(),
6385                    max_row: 13,
6386                    excerpts: vec![RelatedExcerpt {
6387                        row_range: 10..13,
6388                        text: helper_text.into(),
6389                        order: 1,
6390                        context_source: ContextSource::EditHistory,
6391                    }],
6392                    in_open_source_repo: false,
6393                },
6394            ]),
6395            active_buffer_diagnostics: vec![],
6396            excerpt_ranges: ExcerptRanges::default(),
6397            syntax_ranges: None,
6398            in_open_source_repo: false,
6399            can_collect_data: false,
6400            repo_url: None,
6401        };
6402        let marker_table = hashed_regions::build_marker_table(&input);
6403        let helper_markers = &marker_table[1].markers;
6404        let start_tag = hashed_regions::marker_tag(&helper_markers[0].0);
6405        let end_tag = hashed_regions::marker_tag(&helper_markers[helper_markers.len() - 1].0);
6406        let output = format!(
6407            "{start_tag}\nfn helper() {{\n    two();\n}}\n{end_tag}{}",
6408            hashed_regions::V0615_END_MARKER
6409        );
6410
6411        let patch =
6412            parse_zeta2_model_output_as_patch(&output, ZetaFormat::V0615HashRegions, &input)
6413                .unwrap();
6414
6415        assert!(patch.contains("--- a/helper.rs"), "patch: {patch}");
6416        assert!(patch.contains("@@ -11,"), "patch: {patch}");
6417        assert!(patch.contains("-    one();"), "patch: {patch}");
6418        assert!(patch.contains("+    two();"), "patch: {patch}");
6419    }
6420
6421    #[test]
6422    fn test_v0615_expected_output_round_trips_to_patch() {
6423        let current_text = "fn main() {\n    helper();\n}\n";
6424        let helper_text = "fn helper() {\n    one();\n}\n";
6425        let input = Zeta2PromptInput {
6426            cursor_path: Path::new("test.rs").into(),
6427            cursor_excerpt: current_text.into(),
6428            cursor_offset_in_excerpt: current_text.find("helper").unwrap(),
6429            excerpt_start_row: Some(0),
6430            events: Vec::new(),
6431            related_files: Some(vec![
6432                RelatedFile {
6433                    path: Path::new("test.rs").into(),
6434                    max_row: 3,
6435                    excerpts: vec![RelatedExcerpt {
6436                        row_range: 0..3,
6437                        text: current_text.into(),
6438                        order: 0,
6439                        context_source: ContextSource::CurrentFile,
6440                    }],
6441                    in_open_source_repo: false,
6442                },
6443                RelatedFile {
6444                    path: Path::new("helper.rs").into(),
6445                    max_row: 13,
6446                    excerpts: vec![RelatedExcerpt {
6447                        row_range: 10..13,
6448                        text: helper_text.into(),
6449                        order: 1,
6450                        context_source: ContextSource::EditHistory,
6451                    }],
6452                    in_open_source_repo: false,
6453                },
6454            ]),
6455            active_buffer_diagnostics: vec![],
6456            excerpt_ranges: ExcerptRanges::default(),
6457            syntax_ranges: None,
6458            in_open_source_repo: false,
6459            can_collect_data: false,
6460            repo_url: None,
6461        };
6462        let patch = indoc! {"
6463            --- a/helper.rs
6464            +++ b/helper.rs
6465            @@ -11,3 +11,3 @@
6466             fn helper() {
6467            -    one();
6468            +    two();
6469             }
6470        "};
6471
6472        let output =
6473            format_expected_output(&input, ZetaFormat::V0615HashRegions, patch, None).unwrap();
6474        let parsed_patch =
6475            parse_zeta2_model_output_as_patch(&output, ZetaFormat::V0615HashRegions, &input)
6476                .unwrap();
6477
6478        assert!(output.contains(hashed_regions::MARKER_TAG_PREFIX));
6479        assert!(
6480            parsed_patch.contains("-    one();"),
6481            "patch: {parsed_patch}"
6482        );
6483        assert!(
6484            parsed_patch.contains("+    two();"),
6485            "patch: {parsed_patch}"
6486        );
6487    }
6488
6489    #[test]
6490    fn test_zeta3_prompt_matches_zeta2_seed_multi_region_prompt() {
6491        let excerpt = "fn main() {\n    helper();\n}\n";
6492        let cursor_offset = excerpt.find("helper").expect("cursor text exists") + "help".len();
6493        let cursor_row_start = excerpt.find("    helper").expect("cursor row exists");
6494        let syntax_ranges = vec![0..excerpt.len()];
6495        let related_file = make_related_file("related.rs", "fn helper() {}\n");
6496        let mut zeta2_input = make_input(
6497            excerpt,
6498            0..excerpt.len(),
6499            cursor_offset,
6500            vec![make_event("test.rs", "-old\n+new\n")],
6501            vec![related_file.clone()],
6502        );
6503        zeta2_input.excerpt_start_row = Some(10);
6504        zeta2_input.excerpt_ranges =
6505            compute_legacy_excerpt_ranges(excerpt, cursor_offset, &syntax_ranges);
6506        zeta2_input.syntax_ranges = Some(syntax_ranges.clone());
6507
6508        let zeta3_input = Zeta3PromptInput {
6509            cursor_path: Path::new("test.rs").into(),
6510            cursor_position: FilePosition {
6511                row: 11,
6512                column: (cursor_offset - cursor_row_start) as u32,
6513            },
6514            events: vec![Arc::new(make_event("test.rs", "-old\n+new\n"))],
6515            editable_context: vec![
6516                RelatedFile {
6517                    path: Path::new("test.rs").into(),
6518                    max_row: 12,
6519                    excerpts: vec![RelatedExcerpt {
6520                        row_range: 10..12,
6521                        text: excerpt.into(),
6522                        order: 0,
6523                        context_source: ContextSource::CurrentFile,
6524                    }],
6525                    in_open_source_repo: false,
6526                },
6527                related_file,
6528            ],
6529            syntax_ranges,
6530            active_buffer_diagnostics: vec![],
6531            in_open_source_repo: false,
6532            can_collect_data: false,
6533            repo_url: None,
6534        };
6535
6536        assert_eq!(
6537            format_zeta3_prompt(&zeta3_input, ZetaFormat::V0318SeedMultiRegions),
6538            format_zeta_prompt(&zeta2_input, ZetaFormat::V0318SeedMultiRegions),
6539        );
6540    }
6541
6542    #[test]
6543    fn test_zeta3_parse_seed_multi_region_output_as_patch() {
6544        let prefix = (0..20)
6545            .map(|index| format!("prefix {index}\n"))
6546            .collect::<String>();
6547        let excerpt = "fn main() {\n    let value = 1;\n    dbg!(value);\n}\n";
6548        let new_excerpt = "fn main() {\n    let value = 2;\n    dbg!(value);\n}\n";
6549        let cursor_offset = excerpt.find("value =").expect("cursor text exists");
6550        let input = Zeta3PromptInput {
6551            cursor_path: Path::new("test.rs").into(),
6552            cursor_position: FilePosition { row: 21, column: 8 },
6553            events: vec![],
6554            editable_context: vec![RelatedFile {
6555                path: Path::new("test.rs").into(),
6556                max_row: 24,
6557                excerpts: vec![RelatedExcerpt {
6558                    row_range: 20..24,
6559                    text: excerpt.into(),
6560                    order: 0,
6561                    context_source: ContextSource::CurrentFile,
6562                }],
6563                in_open_source_repo: false,
6564            }],
6565            syntax_ranges: vec![0..excerpt.len()],
6566            active_buffer_diagnostics: vec![],
6567            in_open_source_repo: false,
6568            can_collect_data: false,
6569            repo_url: None,
6570        };
6571        let output = multi_region::encode_from_old_and_new_v0318(
6572            excerpt,
6573            new_excerpt,
6574            Some(cursor_offset),
6575            CURSOR_MARKER,
6576            multi_region::V0318_END_MARKER,
6577        )
6578        .unwrap();
6579
6580        let patch =
6581            parse_zeta3_model_output_as_patch(&output, ZetaFormat::V0318SeedMultiRegions, &input)
6582                .unwrap();
6583        let full_text = format!("{prefix}{excerpt}");
6584        let expected = format!("{prefix}{new_excerpt}");
6585
6586        assert!(patch.contains(CURSOR_MARKER));
6587        assert_eq!(
6588            udiff::apply_diff_to_string(&patch.replace(CURSOR_MARKER, ""), &full_text).unwrap(),
6589            expected
6590        );
6591    }
6592
6593    #[test]
6594    fn test_seed_coder_no_context() {
6595        let input = make_input("before\nmiddle\nafter", 7..13, 10, vec![], vec![]);
6596
6597        assert_eq!(
6598            format_seed_coder(&input),
6599            indoc! {r#"
6600                <[fim-suffix]>
6601                after
6602                <[fim-prefix]><filename>test.rs
6603                before
6604                <<<<<<< CURRENT
6605                mid<|user_cursor|>dle
6606                =======
6607                <[fim-middle]>"#}
6608        );
6609    }
6610
6611    #[test]
6612    fn test_seed_coder_truncation_drops_context() {
6613        let input = make_input(
6614            "code",
6615            0..4,
6616            2,
6617            vec![make_event("a.rs", "-x\n+y\n")],
6618            vec![make_related_file("r1.rs", "content\n")],
6619        );
6620
6621        // With large budget, everything is included
6622        assert_eq!(
6623            format_seed_coder(&input),
6624            indoc! {r#"
6625                <[fim-suffix]>
6626                <[fim-prefix]><filename>r1.rs
6627                content
6628
6629                <filename>edit_history
6630                --- a/a.rs
6631                +++ b/a.rs
6632                -x
6633                +y
6634
6635                <filename>test.rs
6636                <<<<<<< CURRENT
6637                co<|user_cursor|>de
6638                =======
6639                <[fim-middle]>"#}
6640        );
6641
6642        assert_eq!(
6643            format_prompt_with_budget_for_format(&input, ZetaFormat::V0211SeedCoder, 24),
6644            None
6645        );
6646
6647        assert_eq!(
6648            format_seed_coder_with_budget(&input, 40),
6649            indoc! {r#"
6650                <[fim-suffix]>
6651                <[fim-prefix]><filename>test.rs
6652                <<<<<<< CURRENT
6653                co<|user_cursor|>de
6654                =======
6655                <[fim-middle]>"#
6656            }
6657        )
6658    }
6659
6660    #[test]
6661    fn test_seed_coder_truncation_prioritizes_lower_order() {
6662        let input = make_input(
6663            "code",
6664            0..4,
6665            2,
6666            vec![],
6667            vec![
6668                RelatedFile {
6669                    path: Path::new("low_prio.rs").into(),
6670                    max_row: 5,
6671                    in_open_source_repo: false,
6672                    excerpts: vec![RelatedExcerpt {
6673                        row_range: 0..5,
6674                        text: "low prio\n".into(),
6675                        order: 10,
6676                        context_source: ContextSource::Lsp,
6677                    }],
6678                },
6679                RelatedFile {
6680                    path: Path::new("high_prio.rs").into(),
6681                    max_row: 5,
6682                    in_open_source_repo: false,
6683                    excerpts: vec![RelatedExcerpt {
6684                        row_range: 0..5,
6685                        text: "high prio\n".into(),
6686                        order: 1,
6687                        context_source: ContextSource::Lsp,
6688                    }],
6689                },
6690            ],
6691        );
6692
6693        // With large budget, both included; rendered in stable lexicographic order.
6694        assert_eq!(
6695            format_seed_coder(&input),
6696            indoc! {r#"
6697                <[fim-suffix]>
6698                <[fim-prefix]><filename>low_prio.rs
6699                low prio
6700                <filename>high_prio.rs
6701                high prio
6702
6703                <filename>test.rs
6704                <<<<<<< CURRENT
6705                co<|user_cursor|>de
6706                =======
6707                <[fim-middle]>"#}
6708        );
6709
6710        // With tight budget under the generic heuristic, context is dropped but the
6711        // minimal cursor section still fits.
6712        assert_eq!(
6713            format_prompt_with_budget_for_format(&input, ZetaFormat::V0211SeedCoder, 44),
6714            Some(
6715                indoc! {r#"
6716                    <[fim-suffix]>
6717                    <[fim-prefix]><filename>test.rs
6718                    <<<<<<< CURRENT
6719                    co<|user_cursor|>de
6720                    =======
6721                    <[fim-middle]>"#}
6722                .to_string()
6723            )
6724        );
6725    }
6726
6727    #[test]
6728    fn test_format_zeta1_from_input_basic() {
6729        let excerpt = "fn before() {}\nfn foo() {\n    let x = 1;\n}\nfn after() {}\n";
6730        let input = Zeta2PromptInput {
6731            cursor_path: Path::new("src/main.rs").into(),
6732            cursor_excerpt: excerpt.into(),
6733            cursor_offset_in_excerpt: 30,
6734            excerpt_start_row: Some(0),
6735            events: vec![Arc::new(make_event("other.rs", "-old\n+new\n"))],
6736            related_files: Some(vec![]),
6737            active_buffer_diagnostics: vec![],
6738            excerpt_ranges: ExcerptRanges {
6739                editable_150: 15..41,
6740                editable_180: 15..41,
6741                editable_350: 15..41,
6742                editable_150_context_350: 0..excerpt.len(),
6743                editable_180_context_350: 0..excerpt.len(),
6744                editable_350_context_150: 0..excerpt.len(),
6745                ..Default::default()
6746            },
6747            syntax_ranges: None,
6748            in_open_source_repo: false,
6749            can_collect_data: false,
6750            repo_url: None,
6751        };
6752
6753        let prompt = zeta1::format_zeta1_from_input(&input, 15..41, 0..excerpt.len());
6754
6755        assert_eq!(
6756            prompt,
6757            concat!(
6758                "### Instruction:\n",
6759                "You are a code completion assistant and your task is to analyze user edits and then rewrite an ",
6760                "excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking ",
6761                "into account the cursor location.\n",
6762                "\n",
6763                "### User Edits:\n",
6764                "\n",
6765                "User edited other.rs:\n",
6766                "```diff\n",
6767                "-old\n",
6768                "+new\n",
6769                "\n",
6770                "```\n",
6771                "\n",
6772                "### User Excerpt:\n",
6773                "\n",
6774                "```src/main.rs\n",
6775                "<|start_of_file|>\n",
6776                "fn before() {}\n",
6777                "<|editable_region_start|>\n",
6778                "fn foo() {\n",
6779                "    <|user_cursor_is_here|>let x = 1;\n",
6780                "\n",
6781                "<|editable_region_end|>}\n",
6782                "fn after() {}\n",
6783                "\n",
6784                "```\n",
6785                "\n",
6786                "### Response:\n",
6787            ),
6788        );
6789    }
6790
6791    #[test]
6792    fn test_format_zeta1_from_input_no_start_of_file() {
6793        let excerpt = "fn foo() {\n    let x = 1;\n}\n";
6794        let input = Zeta2PromptInput {
6795            cursor_path: Path::new("src/main.rs").into(),
6796            cursor_excerpt: excerpt.into(),
6797            cursor_offset_in_excerpt: 15,
6798            excerpt_start_row: Some(10),
6799            events: vec![],
6800            related_files: Some(vec![]),
6801            active_buffer_diagnostics: vec![],
6802            excerpt_ranges: ExcerptRanges {
6803                editable_150: 0..28,
6804                editable_180: 0..28,
6805                editable_350: 0..28,
6806                editable_150_context_350: 0..28,
6807                editable_180_context_350: 0..28,
6808                editable_350_context_150: 0..28,
6809                ..Default::default()
6810            },
6811            syntax_ranges: None,
6812            in_open_source_repo: false,
6813            can_collect_data: false,
6814            repo_url: None,
6815        };
6816
6817        let prompt = zeta1::format_zeta1_from_input(&input, 0..28, 0..28);
6818
6819        assert_eq!(
6820            prompt,
6821            concat!(
6822                "### Instruction:\n",
6823                "You are a code completion assistant and your task is to analyze user edits and then rewrite an ",
6824                "excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking ",
6825                "into account the cursor location.\n",
6826                "\n",
6827                "### User Edits:\n",
6828                "\n",
6829                "\n",
6830                "\n",
6831                "### User Excerpt:\n",
6832                "\n",
6833                "```src/main.rs\n",
6834                "<|editable_region_start|>\n",
6835                "fn foo() {\n",
6836                "    <|user_cursor_is_here|>let x = 1;\n",
6837                "}\n",
6838                "\n",
6839                "<|editable_region_end|>\n",
6840                "```\n",
6841                "\n",
6842                "### Response:\n",
6843            ),
6844        );
6845    }
6846
6847    #[test]
6848    fn test_format_zeta1_from_input_with_sub_ranges() {
6849        let excerpt = "// prefix\nfn foo() {\n    let x = 1;\n}\n// suffix\n";
6850        let editable_range = 10..37;
6851        let context_range = 0..excerpt.len();
6852
6853        let input = Zeta2PromptInput {
6854            cursor_path: Path::new("test.rs").into(),
6855            cursor_excerpt: excerpt.into(),
6856            cursor_offset_in_excerpt: 25,
6857            excerpt_start_row: Some(0),
6858            events: vec![],
6859            related_files: Some(vec![]),
6860            active_buffer_diagnostics: vec![],
6861            excerpt_ranges: ExcerptRanges {
6862                editable_150: editable_range.clone(),
6863                editable_180: editable_range.clone(),
6864                editable_350: editable_range.clone(),
6865                editable_150_context_350: context_range.clone(),
6866                editable_180_context_350: context_range.clone(),
6867                editable_350_context_150: context_range.clone(),
6868                ..Default::default()
6869            },
6870            syntax_ranges: None,
6871            in_open_source_repo: false,
6872            can_collect_data: false,
6873            repo_url: None,
6874        };
6875
6876        let prompt = zeta1::format_zeta1_from_input(&input, editable_range, context_range);
6877
6878        assert_eq!(
6879            prompt,
6880            concat!(
6881                "### Instruction:\n",
6882                "You are a code completion assistant and your task is to analyze user edits and then rewrite an ",
6883                "excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking ",
6884                "into account the cursor location.\n",
6885                "\n",
6886                "### User Edits:\n",
6887                "\n",
6888                "\n",
6889                "\n",
6890                "### User Excerpt:\n",
6891                "\n",
6892                "```test.rs\n",
6893                "<|start_of_file|>\n",
6894                "// prefix\n",
6895                "<|editable_region_start|>\n",
6896                "fn foo() {\n",
6897                "    <|user_cursor_is_here|>let x = 1;\n",
6898                "}\n",
6899                "<|editable_region_end|>\n",
6900                "// suffix\n",
6901                "\n",
6902                "```\n",
6903                "\n",
6904                "### Response:\n",
6905            ),
6906        );
6907    }
6908
6909    #[test]
6910    fn test_max_event_count() {
6911        fn make_numbered_event(index: usize) -> Event {
6912            return make_event(
6913                &format!("event-{index}.rs"),
6914                &format!("-old-{index}\n+new-{index}\n"),
6915            );
6916        }
6917        let input = make_input(
6918            "x",
6919            0..1,
6920            0,
6921            (0..3).map(make_numbered_event).collect(),
6922            vec![],
6923        );
6924
6925        let edit_history_section = format_edit_history_within_budget(
6926            &input.events,
6927            "<|file_sep|>",
6928            "edit history",
6929            usize::MAX,
6930            5,
6931        );
6932
6933        assert_eq!(
6934            &edit_history_section,
6935            indoc!(
6936                "
6937                <|file_sep|>edit history
6938                --- a/event-0.rs
6939                +++ b/event-0.rs
6940                -old-0
6941                +new-0
6942                --- a/event-1.rs
6943                +++ b/event-1.rs
6944                -old-1
6945                +new-1
6946                --- a/event-2.rs
6947                +++ b/event-2.rs
6948                -old-2
6949                +new-2
6950            "
6951            )
6952        );
6953
6954        let edit_history_section = format_edit_history_within_budget(
6955            &input.events,
6956            "<|file_sep|>",
6957            "edit history",
6958            usize::MAX,
6959            2,
6960        );
6961
6962        assert_eq!(
6963            &edit_history_section,
6964            indoc!(
6965                "
6966                <|file_sep|>edit history
6967                --- a/event-1.rs
6968                +++ b/event-1.rs
6969                -old-1
6970                +new-1
6971                --- a/event-2.rs
6972                +++ b/event-2.rs
6973                -old-2
6974                +new-2
6975            "
6976            )
6977        );
6978
6979        let edit_history_section = format_edit_history_within_budget(
6980            &input.events,
6981            "<|file_sep|>",
6982            "edit history",
6983            usize::MAX,
6984            0,
6985        );
6986
6987        assert_eq!(&edit_history_section, "");
6988    }
6989
6990    #[test]
6991    fn test_clean_zeta1_model_output_basic() {
6992        let output = indoc! {"
6993            <|editable_region_start|>
6994            fn main() {
6995                println!(\"hello\");
6996            }
6997            <|editable_region_end|>
6998        "};
6999
7000        let cleaned = zeta1::clean_zeta1_model_output(output).unwrap();
7001        assert_eq!(cleaned, "fn main() {\n    println!(\"hello\");\n}");
7002    }
7003
7004    #[test]
7005    fn test_clean_zeta1_model_output_with_cursor() {
7006        let output = indoc! {"
7007            <|editable_region_start|>
7008            fn main() {
7009                <|user_cursor_is_here|>println!(\"hello\");
7010            }
7011            <|editable_region_end|>
7012        "};
7013
7014        let cleaned = zeta1::clean_zeta1_model_output(output).unwrap();
7015        assert_eq!(
7016            cleaned,
7017            "fn main() {\n    <|user_cursor|>println!(\"hello\");\n}"
7018        );
7019    }
7020
7021    #[test]
7022    fn test_clean_zeta1_model_output_no_markers() {
7023        let output = "fn main() {}\n";
7024        let cleaned = zeta1::clean_zeta1_model_output(output).unwrap();
7025        assert_eq!(cleaned, "fn main() {}\n");
7026    }
7027
7028    #[test]
7029    fn test_clean_zeta1_model_output_empty_region() {
7030        let output = "<|editable_region_start|>\n<|editable_region_end|>\n";
7031        let cleaned = zeta1::clean_zeta1_model_output(output).unwrap();
7032        assert_eq!(cleaned, "");
7033    }
7034
7035    fn apply_edit(excerpt: &str, parsed_output: &ParsedOutput) -> String {
7036        let mut result = excerpt.to_string();
7037        result.replace_range(
7038            parsed_output.range_in_excerpt.clone(),
7039            &parsed_output.new_editable_region,
7040        );
7041        result
7042    }
7043
7044    #[test]
7045    fn test_parse_zeta2_model_output() {
7046        let excerpt = "before ctx\nctx start\neditable old\nctx end\nafter ctx\n";
7047        let context_start = excerpt.find("ctx start").unwrap();
7048        let context_end = excerpt.find("after ctx").unwrap();
7049        let editable_start = excerpt.find("editable old").unwrap();
7050        let editable_end = editable_start + "editable old\n".len();
7051        let input = make_input_with_context_range(
7052            excerpt,
7053            editable_start..editable_end,
7054            context_start..context_end,
7055            editable_start,
7056        );
7057
7058        let output = parse_zeta2_model_output(
7059            "editable new\n>>>>>>> UPDATED\n",
7060            ZetaFormat::V0131GitMergeMarkersPrefix,
7061            &input,
7062        )
7063        .unwrap();
7064
7065        assert_eq!(
7066            apply_edit(excerpt, &output),
7067            "before ctx\nctx start\neditable new\nctx end\nafter ctx\n"
7068        );
7069    }
7070
7071    #[test]
7072    fn test_parse_zeta2_model_output_identity() {
7073        let excerpt = "aaa\nbbb\nccc\nddd\neee\n";
7074        let editable_start = excerpt.find("bbb").unwrap();
7075        let editable_end = excerpt.find("ddd").unwrap();
7076        let input = make_input_with_context_range(
7077            excerpt,
7078            editable_start..editable_end,
7079            0..excerpt.len(),
7080            editable_start,
7081        );
7082
7083        let format = ZetaFormat::V0131GitMergeMarkersPrefix;
7084        let output =
7085            parse_zeta2_model_output("bbb\nccc\n>>>>>>> UPDATED\n", format, &input).unwrap();
7086
7087        assert_eq!(apply_edit(excerpt, &output), excerpt);
7088    }
7089
7090    #[test]
7091    fn test_parse_zeta2_model_output_strips_end_marker() {
7092        let excerpt = "hello\nworld\n";
7093        let input = make_input_with_context_range(excerpt, 0..excerpt.len(), 0..excerpt.len(), 0);
7094
7095        let format = ZetaFormat::V0131GitMergeMarkersPrefix;
7096        let output1 =
7097            parse_zeta2_model_output("new content\n>>>>>>> UPDATED\n", format, &input).unwrap();
7098        let output2 = parse_zeta2_model_output("new content\n", format, &input).unwrap();
7099
7100        assert_eq!(apply_edit(excerpt, &output1), apply_edit(excerpt, &output2));
7101        assert_eq!(apply_edit(excerpt, &output1), "new content\n");
7102    }
7103
7104    #[test]
7105    fn test_parsed_output_to_patch_round_trips_through_udiff_application() {
7106        let excerpt = "before ctx\nctx start\neditable old\nctx end\nafter ctx\n";
7107        let context_start = excerpt.find("ctx start").unwrap();
7108        let context_end = excerpt.find("after ctx").unwrap();
7109        let editable_start = excerpt.find("editable old").unwrap();
7110        let editable_end = editable_start + "editable old\n".len();
7111        let input = make_input_with_context_range(
7112            excerpt,
7113            editable_start..editable_end,
7114            context_start..context_end,
7115            editable_start,
7116        );
7117
7118        let parsed = parse_zeta2_model_output(
7119            "editable new\n>>>>>>> UPDATED\n",
7120            ZetaFormat::V0131GitMergeMarkersPrefix,
7121            &input,
7122        )
7123        .unwrap();
7124        let expected = apply_edit(excerpt, &parsed);
7125        let patch = parsed_output_to_patch(&input, parsed).unwrap();
7126        let patched = udiff::apply_diff_to_string(&patch, excerpt).unwrap();
7127
7128        assert_eq!(patched, expected);
7129    }
7130
7131    #[test]
7132    fn test_special_tokens_not_triggered_by_comment_separator() {
7133        // Regression test for https://github.com/zed-industries/zed/issues/52489
7134        let excerpt = "fn main() {\n    // =======\n    println!(\"hello\");\n}\n";
7135        let input = make_input(excerpt, 0..excerpt.len(), 0, vec![], vec![]);
7136        assert!(
7137            !prompt_input_contains_special_tokens(&input, ZetaFormat::V0131GitMergeMarkersPrefix),
7138            "comment containing ======= should not trigger special token detection"
7139        );
7140    }
7141}
7142