use crate::{    FormatPromptArgs, PredictionProvider,    example::{ActualCursor, Example, ExamplePrompt},    headless::EpAppState,    progress::{ExampleProgress, Step},    retrieve_context::{ContextRetrievalType, run_context_retrieval},};use anyhow::{Context as _, Result, anyhow};use gpui::AsyncApp;use std::ops::Range;use std::sync::Arc;use zeta_prompt::{    Zeta2PromptInput, ZetaFormat, format_edit_history_within_budget, format_expected_output,    format_zeta_prompt,    hashed_regions::{self, SnippetMarkers},    max_edit_event_count_for_format, resolve_cursor_region,}; fn resolved_excerpt_ranges_for_format(    input: &zeta_prompt::Zeta2PromptInput,    format: ZetaFormat,) -> (Range<usize>, Range<usize>) {    let (_, editable_range_in_context, context_range, _) = resolve_cursor_region(input, format);    let editable_range = (context_range.start + editable_range_in_context.start)        ..(context_range.start + editable_range_in_context.end);    (editable_range, context_range)} pub async fn run_format_prompt(    example: &mut Example,    args: &FormatPromptArgs,    app_state: Arc<EpAppState>,    example_progress: &ExampleProgress,    cx: AsyncApp,) -> Result<()> {    run_context_retrieval(        example,        app_state.clone(),        example_progress,        vec![ContextRetrievalType::Lsp],        false,        cx.clone(),    )    .await?;     // Teacher-jumps addresses every edit through related-file excerpts and    // hard-errors unless the cursor file is covered by one. Settled-data    // samples carry a `cursor_excerpt` but were not run through current-file    // context retrieval, so normalize the input to synthesize a current-file    // excerpt from it when the cursor isn't already covered (a no-op for proper    // `ep context --type=current-file` runs).    if matches!(        args.provider,        PredictionProvider::TeacherJumps(_) | PredictionProvider::TeacherJumpsNonBatching(_)    ) {        if let Some(prompt_inputs) = example.prompt_inputs.as_mut() {            hashed_regions::ensure_cursor_file_excerpt(prompt_inputs);        }    }     let step_progress = example_progress.start(Step::FormatPrompt);     let prompt_inputs = example        .prompt_inputs        .as_ref()        .context("prompt_inputs must be set after context retrieval")?;     match args.provider {        PredictionProvider::Teacher(_, zeta_format)        | PredictionProvider::TeacherNonBatching(_, zeta_format) => {            step_progress.set_substatus("formatting teacher prompt");             let (editable_range, context_range) =                resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);             let include_diagnostics = matches!(zeta_format, ZetaFormat::V0420Diagnostics);             let prompt = TeacherPrompt::format_prompt(                example,                editable_range,                context_range,                include_diagnostics,            );            example.prompt = Some(ExamplePrompt {                input: prompt,                expected_output: None,                rejected_output: None,                prefill: None,                provider: args.provider,            });        }        PredictionProvider::TeacherJumps(_) | PredictionProvider::TeacherJumpsNonBatching(_) => {            step_progress.set_substatus("formatting teacher jumps prompt");             let prompt = TeacherJumpsPrompt::format_prompt(example, args.related_files_budget)?;            example.prompt = Some(ExamplePrompt {                input: prompt,                expected_output: None,                rejected_output: None,                prefill: None,                provider: args.provider,            });        }        PredictionProvider::Zeta2(zeta_format) => {            step_progress.set_substatus("formatting zeta2 prompt");             let prompt = format_zeta_prompt(prompt_inputs, zeta_format);            let prefill = zeta_prompt::get_prefill(prompt_inputs, zeta_format);            let expected_output = example                .spec                .expected_patches_with_cursor_positions()                .into_iter()                .next()                .and_then(|(expected_patch, expected_cursor_offset)| {                    format_expected_output(                        prompt_inputs,                        zeta_format,                        &expected_patch,                        expected_cursor_offset,                    )                    .ok()                });             let rejected_output = example.spec.rejected_patch.as_ref().and_then(|patch| {                format_expected_output(prompt_inputs, zeta_format, patch, None).ok()            });             example.prompt = prompt.map(|prompt| ExamplePrompt {                input: prompt,                expected_output,                rejected_output,                provider: args.provider,                prefill: Some(prefill),            });        }        _ => {            panic!("Cannot format prompt for {:?}", args.provider);        }    };    Ok(())} pub struct TeacherPrompt; impl TeacherPrompt {    pub(crate) const EDITABLE_REGION_START: &str = "<|editable_region_start|>\n";    pub(crate) const EDITABLE_REGION_END: &str = "\n<|editable_region_end|>";    pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";    pub(crate) const NO_EDITS: &str = "NO_EDITS";     /// Truncate edit history to this number of last lines    const MAX_HISTORY_LINES: usize = 128;     pub fn format_prompt(        example: &Example,        editable_range: Range<usize>,        context_range: Range<usize>,        include_diagnostics: bool,    ) -> String {        let edit_history = Self::format_edit_history(&example.spec.edit_history);        let context = Self::format_context(example);        let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range);        let diagnostics = include_diagnostics            .then(|| Self::format_diagnostics(example))            .map(|diagnostics| format!("# 4. Diagnostics\n\n{diagnostics}"));         let prompt_template = crate::prompt_assets::get_prompt("teacher.md");        let prompt = prompt_template            .replace("{{context}}", &context)            .replace("{{edit_history}}", &edit_history)            .replace("{{diagnostics}}", diagnostics.as_deref().unwrap_or(""))            .replace("{{cursor_excerpt}}", &cursor_excerpt);         prompt    }     pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {        // Check if the model indicated no edits are needed        let no_edits = (String::new(), None);        if let Some(last_codeblock) = extract_last_codeblock(&response) {            if last_codeblock.trim() == Self::NO_EDITS {                return Ok(no_edits);            }        }         if response            .trim_end_matches(&[' ', '\n', '`'])            .ends_with(Self::NO_EDITS)        {            return Ok(no_edits);        }         // Extract updated (new) editable region from the model response.        let new_editable_region = Self::extract_editable_region(&response)?;        let cursor_offset = new_editable_region.find(Self::USER_CURSOR_MARKER);        let mut new_editable_region = new_editable_region.replace(Self::USER_CURSOR_MARKER, "");        let old_editable_region = Self::extract_editable_region(            &example                .prompt                .as_ref()                .context("example prompt missing")?                .input,        )?        .replace(Self::USER_CURSOR_MARKER, "");         let prompt_inputs = example            .prompt_inputs            .as_ref()            .context("example is missing prompt inputs")?;         // Normalize leading newlines: if old starts with newline but new doesn't,        // prepend newline to new to preserve whitespace structure.        // This handles the case where the model drops the leading blank line.        if old_editable_region.starts_with('\n') && !new_editable_region.starts_with('\n') {            new_editable_region.insert(0, '\n');        }         let excerpt = prompt_inputs.cursor_excerpt.as_ref();        let (editable_region_offset, _) = excerpt            .match_indices(&old_editable_region)            .min_by_key(|(index, _)| index.abs_diff(prompt_inputs.cursor_offset_in_excerpt))            .context("editable region not found in prompt content")?;        let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();         let editable_region_lines = old_editable_region.lines().count() as u32;        let diff = language::unified_diff_with_context(            &old_editable_region,            &new_editable_region,            editable_region_start_line as u32,            editable_region_start_line as u32,            editable_region_lines,        );         let diff = indoc::formatdoc! {"            --- a/{path}            +++ b/{path}            {diff}",            path = example.spec.cursor_path.to_string_lossy(),            diff = diff,        };         let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {            ActualCursor::from_editable_region(                &example.spec.cursor_path,                editable_region_cursor_offset,                &new_editable_region,                excerpt,                editable_region_offset,                editable_region_start_line,            )        });         Ok((diff, actual_cursor))    }     fn format_edit_history(edit_history: &str) -> String {        let lines: Vec<&str> = edit_history.lines().collect();         if lines.is_empty() {            return "(No edit history)".to_string();        }         if lines.len() > Self::MAX_HISTORY_LINES {            let truncated = lines[lines.len() - Self::MAX_HISTORY_LINES..].join("\n");            format!("{truncated}\n[...truncated...]")        } else {            lines.join("\n")        }    }     pub fn format_context(example: &Example) -> String {        let related_files = example            .prompt_inputs            .as_ref()            .and_then(|pi| pi.related_files.as_deref());         let Some(related_files) = related_files else {            return "(No context)".to_string();        };         if related_files.is_empty() {            return "(No context)".to_string();        }         let prefix = "`````";        let suffix = "`````\n\n";        let max_tokens = 1024;        zeta_prompt::format_related_files_within_budget(related_files, &prefix, &suffix, max_tokens)    }     fn format_cursor_excerpt(        example: &Example,        editable_range: Range<usize>,        context_range: Range<usize>,    ) -> String {        let mut result = String::new();         let prompt_inputs = example.prompt_inputs.as_ref().unwrap();        let excerpt = prompt_inputs.cursor_excerpt.as_ref();        let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;         let path_str = example.spec.cursor_path.to_string_lossy();        result.push_str(&format!("`````{path_str}\n"));        result.push_str(&excerpt[context_range.start..editable_range.start]);        result.push_str(Self::EDITABLE_REGION_START);        result.push_str(&excerpt[editable_range.start..cursor_offset]);        result.push_str(Self::USER_CURSOR_MARKER);        result.push_str(&excerpt[cursor_offset..editable_range.end]);        result.push_str(Self::EDITABLE_REGION_END);        result.push_str(&excerpt[editable_range.end..context_range.end]);        result.push_str("\n`````");         result    }     pub fn extract_editable_region(text: &str) -> Result<String> {        let start = text            .rfind(Self::EDITABLE_REGION_START)            .map_or(0, |pos| pos + Self::EDITABLE_REGION_START.len());        let end = text.rfind(Self::EDITABLE_REGION_END).unwrap_or(text.len());         if start >= end {            return Err(anyhow!("Invalid editable region markers"));        }         let region = &text[start..end];        Ok(region.strip_suffix('\n').unwrap_or(region).to_string())    }     fn format_diagnostics(example: &Example) -> String {        let Some(prompt_inputs) = example.prompt_inputs.as_ref() else {            return "No Diagnostics".to_string();        };         let cursor_buffer_row = prompt_inputs.excerpt_start_row.map(|excerpt_start_row| {            excerpt_start_row                + prompt_inputs.cursor_excerpt[..prompt_inputs.cursor_offset_in_excerpt]                    .bytes()                    .filter(|byte| *byte == b'\n')                    .count() as u32        });        let diagnostics = zeta_prompt::format_active_buffer_diagnostics_with_budget(            &prompt_inputs.active_buffer_diagnostics,            cursor_buffer_row,            2_000,        );         let diagnostics = diagnostics            .strip_prefix("<filename>diagnostics\n")            .unwrap_or(&diagnostics);         if diagnostics.is_empty() {            "No Diagnostics".to_string()        } else {            diagnostics.to_string()        }    }} /// Teacher prompt for long-range edit prediction ("jumps"). All prompt/// context — the cursor file and every related-file excerpt — is annotated/// with hashed region markers (V0609HashedRegions), and the teacher may/// output a sequence of marker-bounded edits targeting any of it.pub struct TeacherJumpsPrompt; impl TeacherJumpsPrompt {    pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";    pub(crate) const NO_EDITS: &str = "NO_EDITS";     const MAX_HISTORY_TOKENS: usize = 4000;     pub const DEFAULT_RELATED_FILES_BUDGET: usize = 8192;     pub fn format_prompt(example: &Example, related_files_budget: usize) -> Result<String> {        let prompt_inputs = example            .prompt_inputs            .as_ref()            .context("example is missing prompt inputs")?;        let marker_table = hashed_regions::build_marker_table(prompt_inputs);        let cursor = hashed_regions::locate_cursor_in_related_files(prompt_inputs).context(            "cursor position is not covered by any related-file excerpt of the cursor file; \             teacher-jumps requires current-file context retrieval (e.g. `ep context --type=current-file,...`)",        )?;         let edit_history = Self::format_edit_history(&prompt_inputs);        let context = Self::format_context(            prompt_inputs,            &marker_table,            related_files_budget,            cursor.file_ix,        );        let cursor_excerpt =            Self::format_cursor_excerpt(example, prompt_inputs, &marker_table, &cursor)?;         let prompt_template = crate::prompt_assets::get_prompt("teacher_jumps.md");        let prompt = prompt_template            .replace("{{context}}", &context)            .replace("{{edit_history}}", &edit_history)            .replace("{{cursor_excerpt}}", &cursor_excerpt);         Ok(prompt)    }     pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {        let no_edits = (String::new(), None);        if let Some(last_codeblock) = extract_last_codeblock(&response) {            if last_codeblock.trim() == Self::NO_EDITS {                return Ok(no_edits);            }        }         if response.trim().ends_with(Self::NO_EDITS) {            return Ok(no_edits);        }         let prompt_inputs = example            .prompt_inputs            .as_ref()            .context("example is missing prompt inputs")?;         // The teacher emits reasoning plus a sequence of markdown code fences,        // one per edit, each a marker-bounded span. Extract the spans from the        // fences, then hand off to the shared hash-region patch assembler that        // the student parser also uses.        let codeblocks: Vec<String> = extract_all_codeblocks(response)            .into_iter()            .filter(|block| block.contains(hashed_regions::MARKER_TAG_PREFIX))            .collect();        if codeblocks.is_empty() {            return Err(anyhow!(                "no marker-bounded edit codeblocks found in model response"            ));        }         let mut spans = Vec::with_capacity(codeblocks.len());        for codeblock in &codeblocks {            spans.push(hashed_regions::extract_marker_span(codeblock)?);        }         let (patch, cursor) = hashed_regions::build_patch_from_spans(            prompt_inputs,            &spans,            Self::USER_CURSOR_MARKER,        )?;         let actual_cursor = cursor.map(|cursor| {            ActualCursor::from_editable_region(                &cursor.path,                cursor.cursor_offset_in_new_text,                &cursor.new_text,                &cursor.old_text,                0,                cursor.start_row as usize,            )        });         Ok((patch, actual_cursor))    }     fn format_edit_history(prompt_inputs: &Zeta2PromptInput) -> String {        format_edit_history_within_budget(            &prompt_inputs.events,            "",            "",            Self::MAX_HISTORY_TOKENS,            max_edit_event_count_for_format(&ZetaFormat::V0327SingleFile),        )    }     /// Render related files with hashed region markers, within a token    /// budget. Mirrors `zeta_prompt::format_related_files_within_budget`,    /// but inserts marker tags into every included excerpt. The cursor file    /// is skipped: it renders in its own prompt section via    /// `format_cursor_excerpt`, and including it here would duplicate it.    fn format_context(        prompt_inputs: &Zeta2PromptInput,        marker_table: &[SnippetMarkers],        max_tokens: usize,        cursor_file_ix: usize,    ) -> String {        let Some(related_files) = prompt_inputs.related_files.as_deref() else {            return "(No context)".to_string();        };        if related_files.is_empty() {            return "(No context)".to_string();        }         let estimate_tokens = |bytes: usize| bytes / 3;         struct RenderedExcerpt {            file_ix: usize,            excerpt_ix: usize,            order: usize,            rendered: String,        }         let mut candidates = Vec::new();        for (file_ix, file) in related_files.iter().enumerate() {            if file_ix == cursor_file_ix {                continue;

Lines 1–500 of 1678. Load more