use crate::{    FormatPromptArgs, PredictArgs, PredictionProvider, TeacherBackend,    anthropic_client::AnthropicClient,    example::{Example, ExamplePrediction, ExamplePrompt},    format_prompt::{TeacherMultiRegionPrompt, TeacherPrompt, run_format_prompt},    headless::EpAppState,    load_project::run_load_project,    openai_client::OpenAiClient,    parse_output::parse_prediction_output,    paths::{LATEST_EXAMPLE_RUN_DIR, RUN_DIR},    progress::{ExampleProgress, InfoStyle, Progress, Step, StepProgress},    retrieve_context::run_context_retrieval,};use anyhow::Context as _;use cloud_llm_client::predict_edits_v3::{RawCompletionRequest, RawCompletionResponse};use edit_prediction::{DebugEvent, EditPredictionStore, Zeta2RawConfig};use futures::{AsyncReadExt as _, FutureExt as _, StreamExt as _, future::Shared};use gpui::{AppContext as _, AsyncApp, Task};use http_client::{AsyncBody, HttpClient, Method};use reqwest_client::ReqwestClient;use std::{    fs,    sync::{        Arc, Mutex, OnceLock,        atomic::{AtomicUsize, Ordering::SeqCst},    },};use zeta_prompt::ZetaFormat; static ANTHROPIC_CLIENT: OnceLock<AnthropicClient> = OnceLock::new();static OPENAI_CLIENT: OnceLock<OpenAiClient> = OnceLock::new(); pub async fn run_prediction(    example: &mut Example,    args: &PredictArgs,    app_state: Arc<EpAppState>,    example_progress: &ExampleProgress,    mut cx: AsyncApp,) -> anyhow::Result<()> {    let repetition_count = args.repetitions;     if let Some(existing_prediction) = example.predictions.first() {        let has_prediction = existing_prediction.actual_patch.is_some()            || !existing_prediction.actual_output.is_empty();        if has_prediction {            match args.provider {                None => return Ok(()),                Some(provider) if existing_prediction.provider == provider => return Ok(()),                Some(_) => example.predictions.clear(),            }        }    }     let Some(provider) = args.provider else {        anyhow::bail!(            "No existing predictions found. Use --provider to specify which model to use for prediction."        );    };     if matches!(        provider,        PredictionProvider::TeacherMultiRegion(..)            | PredictionProvider::TeacherMultiRegionNonBatching(..)    ) {        anyhow::bail!("Teacher multi-region providers are not supported for prediction.");    }     if let PredictionProvider::Teacher(backend, _)    | PredictionProvider::TeacherNonBatching(backend, _) = provider    {        run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;        run_format_prompt(            example,            &FormatPromptArgs { provider },            app_state.clone(),            example_progress,            cx,        )        .await?;         let step_progress = example_progress.start(Step::Predict);        let batched = matches!(            provider,            PredictionProvider::Teacher(..) | PredictionProvider::TeacherMultiRegion(..)        );        return predict_teacher(            example,            backend,            batched,            repetition_count,            args.cache_only,            &step_progress,        )        .await;    }     if let PredictionProvider::Baseten(format) = provider {        run_format_prompt(            example,            &FormatPromptArgs {                provider: PredictionProvider::Zeta2(format),            },            app_state.clone(),            example_progress,            cx,        )        .await?;         let step_progress = example_progress.start(Step::Predict);        return predict_baseten(example, format, &step_progress).await;    }     run_load_project(example, app_state.clone(), example_progress, cx.clone()).await?;    run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;     let step_progress = example_progress.start(Step::Predict);     if matches!(        provider,        PredictionProvider::Zeta1 | PredictionProvider::Zeta2(_)    ) {        step_progress.set_substatus("authenticating");        static AUTHENTICATED: OnceLock<Shared<Task<()>>> = OnceLock::new();        AUTHENTICATED            .get_or_init(|| {                let client = app_state.client.clone();                cx.spawn(async move |cx| {                    if let Err(e) = client.sign_in_with_optional_connect(true, cx).await {                        eprintln!("Authentication failed: {}", e);                    }                })                .shared()            })            .clone()            .await;    }     let ep_store = cx        .update(|cx| EditPredictionStore::try_global(cx))        .context("EditPredictionStore not initialized")?;     ep_store.update(&mut cx, |store, _cx| {        let model = match provider {            PredictionProvider::Zeta1 => edit_prediction::EditPredictionModel::Zeta,            PredictionProvider::Zeta2(_) => edit_prediction::EditPredictionModel::Zeta,            PredictionProvider::Mercury => edit_prediction::EditPredictionModel::Mercury,            PredictionProvider::Teacher(..)            | PredictionProvider::TeacherMultiRegion(..)            | PredictionProvider::TeacherNonBatching(..)            | PredictionProvider::TeacherMultiRegionNonBatching(..)            | PredictionProvider::Repair            | PredictionProvider::Baseten(_) => {                unreachable!()            }        };        store.set_edit_prediction_model(model);         // If user specified a non-default Zeta2 version, configure raw endpoint.        // ZED_ZETA_MODEL env var is optional.        if let PredictionProvider::Zeta2(format) = provider {            if format != ZetaFormat::default() {                let model_id = std::env::var("ZED_ZETA_MODEL").ok();                let environment = std::env::var("ZED_ZETA_ENVIRONMENT").ok();                store.set_zeta2_raw_config(Zeta2RawConfig {                    model_id,                    environment,                    format,                });            }        }    });    step_progress.set_substatus("configuring model");    let state = example.state.as_ref().context("state must be set")?;    let run_dir = RUN_DIR.join(&example.spec.name);     let updated_example = Arc::new(Mutex::new(example.clone()));    let current_run_ix = Arc::new(AtomicUsize::new(0));     let mut debug_rx = ep_store.update(&mut cx, |store, cx| store.debug_info(&state.project, cx));    let debug_task = cx.background_spawn({        let updated_example = updated_example.clone();        let current_run_ix = current_run_ix.clone();        let run_dir = run_dir.clone();        async move {            while let Some(event) = debug_rx.next().await {                let run_ix = current_run_ix.load(SeqCst);                let mut updated_example = updated_example.lock().unwrap();                 let run_dir = if repetition_count > 1 {                    run_dir.join(format!("{:03}", run_ix))                } else {                    run_dir.clone()                };                 match event {                    DebugEvent::EditPredictionStarted(request) => {                        assert_eq!(updated_example.predictions.len(), run_ix + 1);                         if let Some(prompt) = request.prompt {                            fs::write(run_dir.join("prediction_prompt.md"), &prompt)?;                            if matches!(provider, PredictionProvider::Zeta2(_)) {                                updated_example.prompt.get_or_insert(ExamplePrompt {                                    input: prompt,                                    expected_output: None,                                    rejected_output: None,                                    provider,                                    prefill: None,                                });                            }                        }                    }                    DebugEvent::EditPredictionFinished(request) => {                        assert_eq!(updated_example.predictions.len(), run_ix + 1);                         if let Some(output) = request.model_output {                            fs::write(run_dir.join("prediction_response.md"), &output)?;                            updated_example                                .predictions                                .last_mut()                                .unwrap()                                .actual_output = output;                        }                        if run_ix >= repetition_count {                            break;                        }                    }                    _ => {}                }            }            anyhow::Ok(())        }    });     for ix in 0..repetition_count {        current_run_ix.store(ix, SeqCst);        let run_dir = if repetition_count > 1 {            run_dir.join(format!("{:03}", ix))        } else {            run_dir.clone()        };         if repetition_count > 1 {            step_progress.set_substatus(format!(                "running prediction {}/{}",                ix + 1,                repetition_count            ));        } else {            step_progress.set_substatus("running prediction");        }         fs::create_dir_all(&run_dir)?;        if LATEST_EXAMPLE_RUN_DIR.is_symlink() {            fs::remove_file(&*LATEST_EXAMPLE_RUN_DIR)?;        }        #[cfg(unix)]        std::os::unix::fs::symlink(&run_dir, &*LATEST_EXAMPLE_RUN_DIR)?;        #[cfg(windows)]        std::os::windows::fs::symlink_dir(&run_dir, &*LATEST_EXAMPLE_RUN_DIR)?;         updated_example            .lock()            .unwrap()            .predictions            .push(ExamplePrediction {                actual_patch: None,                actual_output: String::new(),                actual_cursor: None,                error: None,                provider,                cumulative_logprob: None,                avg_logprob: None,            });         step_progress.set_substatus("requesting prediction");        let prediction = ep_store            .update(&mut cx, |store, cx| {                store.request_prediction(                    &state.project,                    &state.buffer,                    state.cursor_position,                    cloud_llm_client::PredictEditsRequestTrigger::Cli,                    cx,                )            })            .await?;         let actual_patch = prediction.and_then(|prediction| {            let prediction = prediction.prediction.ok()?;            prediction                .edit_preview                .as_unified_diff(prediction.snapshot.file(), &prediction.edits)        });         let has_prediction = actual_patch.as_ref().is_some_and(|p| !p.is_empty());         updated_example            .lock()            .unwrap()            .predictions            .last_mut()            .unwrap()            .actual_patch = actual_patch;         if ix == repetition_count - 1 {            let (info, style) = if has_prediction {                ("predicted", InfoStyle::Normal)            } else {                ("no prediction", InfoStyle::Warning)            };            step_progress.set_info(info, style);        }    }     ep_store.update(&mut cx, |store, _| {        store.remove_project(&state.project);    });    debug_task.await?;     *example = Arc::into_inner(updated_example)        .ok_or_else(|| anyhow::anyhow!("Failed to unwrap Arc"))?        .into_inner()        .map_err(|_| anyhow::anyhow!("Failed to unwrap Mutex"))?;    Ok(())} async fn predict_teacher(    example: &mut Example,    backend: TeacherBackend,    batched: bool,    repetition_count: usize,    cache_only: bool,    step_progress: &crate::progress::StepProgress,) -> anyhow::Result<()> {    match backend {        TeacherBackend::Sonnet45 | TeacherBackend::Sonnet46 => {            predict_anthropic(                example,                backend,                batched,                repetition_count,                cache_only,                step_progress,            )            .await        }        TeacherBackend::Gpt52 => {            predict_openai(                example,                backend,                batched,                repetition_count,                cache_only,                step_progress,            )            .await        }    }} async fn predict_anthropic(    example: &mut Example,    backend: TeacherBackend,    batched: bool,    repetition_count: usize,    cache_only: bool,    step_progress: &crate::progress::StepProgress,) -> anyhow::Result<()> {    let llm_model_name = backend.model_name();    let max_tokens = 16384;    let llm_client = ANTHROPIC_CLIENT.get_or_init(|| {        let client = if batched {            AnthropicClient::batch(&crate::paths::LLM_CACHE_DB)        } else {            AnthropicClient::plain()        };        client.expect("Failed to create Anthropic client")    });     let prompt = example.prompt.as_ref().context("Prompt is required")?;     for ix in 0..repetition_count {        if repetition_count > 1 {            step_progress.set_substatus(format!(                "running prediction {}/{}",                ix + 1,                repetition_count            ));        } else {            step_progress.set_substatus("running prediction");        }         let messages = vec![anthropic::Message {            role: anthropic::Role::User,            content: vec![anthropic::RequestContent::Text {                text: prompt.input.clone(),                cache_control: None,            }],        }];         let seed = if repetition_count > 1 { Some(ix) } else { None };        let Some(response) = llm_client            .generate(llm_model_name, max_tokens, messages, seed, cache_only)            .await?        else {            // Request stashed for batched processing            continue;        };         let actual_output = response            .content            .into_iter()            .filter_map(|content| match content {                anthropic::ResponseContent::Text { text } => Some(text),                _ => None,            })            .collect::<Vec<String>>()            .join("\n");         let parser_provider = if batched {            example                .prompt                .as_ref()                .map(|prompt| prompt.provider)                .unwrap_or(PredictionProvider::Teacher(backend, ZetaFormat::default()))        } else {            match example.prompt.as_ref().map(|prompt| prompt.provider) {                Some(PredictionProvider::TeacherMultiRegion(_))                | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {                    PredictionProvider::TeacherMultiRegionNonBatching(backend)                }                _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),            }        };         let (actual_patch, actual_cursor) = match parser_provider {            PredictionProvider::TeacherMultiRegion(_)            | PredictionProvider::TeacherMultiRegionNonBatching(_) => {                TeacherMultiRegionPrompt::parse(example, &actual_output)?            }            _ => TeacherPrompt::parse(example, &actual_output)?,        };         let prediction = ExamplePrediction {            actual_patch: Some(actual_patch),            actual_output,            actual_cursor,            error: None,            provider: if batched {                match example.prompt.as_ref().map(|prompt| prompt.provider) {                    Some(PredictionProvider::TeacherMultiRegion(_)) => {                        PredictionProvider::TeacherMultiRegion(backend)                    }                    _ => PredictionProvider::Teacher(backend, ZetaFormat::default()),                }            } else {                match example.prompt.as_ref().map(|prompt| prompt.provider) {                    Some(PredictionProvider::TeacherMultiRegion(_))                    | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {                        PredictionProvider::TeacherMultiRegionNonBatching(backend)                    }                    _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),                }            },            cumulative_logprob: None,            avg_logprob: None,        };         example.predictions.push(prediction);    }    Ok(())} async fn predict_openai(    example: &mut Example,    backend: TeacherBackend,    batched: bool,    repetition_count: usize,    cache_only: bool,    step_progress: &crate::progress::StepProgress,) -> anyhow::Result<()> {    let llm_model_name = backend.model_name();    let max_tokens = 16384;    let llm_client = OPENAI_CLIENT.get_or_init(|| {        let client = if batched {            OpenAiClient::batch(&crate::paths::LLM_CACHE_DB)        } else {            OpenAiClient::plain()        };        client.expect("Failed to create OpenAI client")    });     let prompt = example.prompt.as_ref().context("Prompt is required")?;     for ix in 0..repetition_count {        if repetition_count > 1 {            step_progress.set_substatus(format!(                "running prediction {}/{}",                ix + 1,                repetition_count

Lines 1–500 of 789. Load more