tenant.openagents/omega

1use crate::{
2    CloudRequestTimeoutError, CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent,
3    EditPredictionId, EditPredictionInputs, EditPredictionModelInput,
4    EditPredictionStartedDebugEvent, EditPredictionStore, PromptHistoryBoundary,
5    ZedUpdateRequiredError, buffer_path_with_id_fallback,
6    cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges},
7    data_collection::CapturedPredictionContext,
8    prediction::EditPredictionResult,
9    udiff::prediction_edits_for_single_file_diff,
10};
11use anyhow::{Context as _, Result};
12use cloud_llm_client::{
13    AcceptEditPredictionBody, EditPredictionRejectReason, predict_edits_v3::RawCompletionRequest,
14};
15use edit_prediction_types::PredictedCursorPosition;
16use gpui::{App, AppContext as _, Entity, Task, TaskExt, WeakEntity, prelude::*};
17use language::{
18    Buffer, BufferSnapshot, DiagnosticSeverity, EditPredictionPromptFormat, OffsetRangeExt as _,
19    ToOffset as _, ZetaVersion, language_settings::all_language_settings, text_diff,
20};
21use project::Project;
22use release_channel::AppVersion;
23use text::{Anchor, Bias, Point};
24use ui::SharedString;
25use workspace::notifications::simple_message_notification::MessageNotification;
26use workspace::notifications::{NotificationId, show_app_notification};
27use workspace::workspace_error::{ErrorAction, ErrorSeverity, WorkspaceError};
28
29use std::{ops::Range, path::Path, sync::Arc};
30use zeta_prompt::{
31    FilePosition, ParsedOutput, Zeta2PromptInput, Zeta3PromptInput, ZetaFormat,
32    excerpt_ranges_for_format, format_zeta_prompt, get_prefill, parse_zeta2_model_output,
33    stop_tokens_for_format,
34    zeta1::{self, EDITABLE_REGION_END_MARKER},
35};
36
37use crate::open_ai_compatible::{
38    load_open_ai_compatible_api_key_if_needed, send_custom_server_request,
39};
40
41pub(crate) fn request_prediction_with_zeta(
42    store: &mut EditPredictionStore,
43    EditPredictionModelInput {
44        buffer,
45        snapshot,
46        position,
47        related_files,
48        editable_context,
49        events,
50        debug_tx,
51        mode,
52        trigger,
53        project,
54        diagnostic_search_range,
55        can_collect_data,
56        is_open_source,
57        allow_jump,
58        ..
59    }: EditPredictionModelInput,
60    context_task: Option<Task<Result<CapturedPredictionContext>>>,
61    prompt_history_boundary: Option<PromptHistoryBoundary>,
62    repo_url: Option<String>,
63    cx: &mut Context<EditPredictionStore>,
64) -> Task<Result<Option<EditPredictionResult>>> {
65    let settings = &all_language_settings(None, cx).edit_predictions;
66    let provider = settings.provider;
67    let custom_server_settings = match provider {
68        settings::EditPredictionProvider::Ollama => settings.ollama.clone(),
69        settings::EditPredictionProvider::OpenAiCompatibleApi => {
70            settings.open_ai_compatible_api.clone()
71        }
72        _ => None,
73    };
74
75    let http_client = cx.http_client();
76    let request_start = cx.background_executor().now();
77    let raw_config = store.zeta2_raw_config().cloned();
78    let preferred_experiment = store.preferred_experiment().map(|s| s.to_owned());
79    let open_ai_compatible_api_key = load_open_ai_compatible_api_key_if_needed(provider, cx);
80
81    let excerpt_path = buffer_path_with_id_fallback(snapshot.file(), &snapshot.text, cx);
82
83    let repo_url = repo_url.filter(|_| can_collect_data);
84    let client = store.client.clone();
85    let llm_token = store.llm_token.clone();
86    let organization_id = store
87        .user_store
88        .read(cx)
89        .current_organization()
90        .map(|organization| organization.id.clone());
91    let app_version = AppVersion::global(cx);
92
93    enum PendingPrediction {
94        Resolved {
95            inputs: EditPredictionInputs,
96            buffer: Entity<Buffer>,
97            snapshot: BufferSnapshot,
98            edits: Vec<(Range<Anchor>, Arc<str>)>,
99            cursor_position: Option<PredictedCursorPosition>,
100            editable_range_in_buffer: Option<Range<usize>>,
101        },
102        Patch {
103            inputs: EditPredictionInputs,
104            project: Entity<Project>,
105            buffer: Entity<Buffer>,
106            snapshot: BufferSnapshot,
107            patch: String,
108        },
109    }
110
111    let project_for_request = project.clone();
112    let request_task = cx.background_spawn({
113        async move {
114            let local_zeta_version = custom_server_settings
115                .as_ref()
116                .and_then(|settings| match settings.prompt_format {
117                    EditPredictionPromptFormat::Zeta(version) => Some(version),
118                    EditPredictionPromptFormat::Infer => {
119                        match settings.model.to_ascii_lowercase().as_str() {
120                            "zeta" | "zeta1" => Some(ZetaVersion::Zeta1),
121                            "zeta2" => Some(ZetaVersion::Zeta2),
122                            "zeta2.1" => Some(ZetaVersion::Zeta2_1),
123                            _ => None,
124                        }
125                    }
126                    _ => None,
127                })
128                .unwrap_or_default();
129            let zeta_format = raw_config
130                .as_ref()
131                .map(|config| config.format)
132                .or(match local_zeta_version {
133                    ZetaVersion::Zeta1 => None,
134                    ZetaVersion::Zeta2 => Some(ZetaFormat::V0211SeedCoder),
135                    ZetaVersion::Zeta2_1 => Some(ZetaFormat::V0318SeedMultiRegions),
136                })
137                .unwrap_or_default();
138
139            enum RequestInput {
140                V3 {
141                    full_context_offset_range: Range<usize>,
142                    prompt_input: Zeta2PromptInput,
143                    formatted_prompt: Option<String>,
144                },
145                V4(Zeta3PromptInput),
146            }
147
148            let cursor_offset = position.to_offset(&snapshot);
149            let request_input =
150                if allow_jump && custom_server_settings.is_none() && raw_config.is_none() {
151                    let cursor_point = snapshot.offset_to_point(cursor_offset);
152                    let editable_context = editable_context
153                        .context("missing editable context task")?
154                        .await?;
155                    let active_buffer_diagnostics = active_buffer_diagnostics(
156                        &snapshot,
157                        diagnostic_search_range,
158                        cursor_point.row,
159                        ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT,
160                    );
161                    let syntax_ranges =
162                        compute_syntax_ranges(&snapshot, cursor_offset, &(0..snapshot.len()));
163
164                    RequestInput::V4(Zeta3PromptInput {
165                        cursor_path: excerpt_path,
166                        cursor_position: FilePosition {
167                            row: cursor_point.row,
168                            column: cursor_point.column,
169                        },
170                        events,
171                        editable_context,
172                        syntax_ranges,
173                        active_buffer_diagnostics,
174                        in_open_source_repo: is_open_source,
175                        can_collect_data,
176                        repo_url,
177                    })
178                } else {
179                    let (full_context_offset_range, prompt_input) = zeta2_prompt_input(
180                        &snapshot,
181                        related_files,
182                        events,
183                        diagnostic_search_range,
184                        excerpt_path,
185                        cursor_offset,
186                        is_open_source,
187                        can_collect_data,
188                        repo_url,
189                    );
190                    let formatted_prompt = format_zeta_prompt(&prompt_input, zeta_format);
191
192                    RequestInput::V3 {
193                        full_context_offset_range,
194                        prompt_input,
195                        formatted_prompt,
196                    }
197                };
198
199            if let Some(debug_tx) = &debug_tx {
200                let prompt = match &request_input {
201                    RequestInput::V3 {
202                        formatted_prompt, ..
203                    } => formatted_prompt.clone(),
204                    RequestInput::V4(_) => None,
205                };
206                debug_tx
207                    .unbounded_send(DebugEvent::EditPredictionStarted(
208                        EditPredictionStartedDebugEvent {
209                            buffer: buffer.downgrade(),
210                            prompt,
211                            position,
212                        },
213                    ))
214                    .ok();
215            }
216
217            log::trace!("Sending edit prediction request");
218
219            let (full_context_offset_range, prompt_input, formatted_prompt) = match request_input {
220                RequestInput::V3 {
221                    full_context_offset_range,
222                    prompt_input,
223                    formatted_prompt,
224                } => (full_context_offset_range, prompt_input, formatted_prompt),
225                RequestInput::V4(v4_prompt_input) => {
226                    let inputs = EditPredictionInputs::V3(v4_prompt_input.clone());
227                    let (response, usage) = EditPredictionStore::send_v4_request(
228                        v4_prompt_input,
229                        preferred_experiment.clone(),
230                        client,
231                        llm_token,
232                        organization_id,
233                        app_version,
234                        trigger,
235                        mode,
236                    )
237                    .await?;
238
239                    let request_id = EditPredictionId(response.request_id.into());
240                    let model_version = response.model_version;
241                    return anyhow::Ok((
242                        Some((
243                            request_id,
244                            Some(PendingPrediction::Patch {
245                                inputs,
246                                project: project_for_request,
247                                buffer,
248                                snapshot: snapshot.clone(),
249                                patch: response.patch,
250                            }),
251                            model_version,
252                        )),
253                        usage,
254                    ));
255                }
256            };
257
258            let Some((request_id, output, model_version, usage)) =
259                (if let Some(custom_settings) = &custom_server_settings {
260                    let max_tokens = custom_settings.max_output_tokens * 4;
261
262                    Some(match local_zeta_version {
263                        ZetaVersion::Zeta1 => {
264                            let ranges = &prompt_input.excerpt_ranges;
265                            let editable_range_in_excerpt = ranges.editable_350.clone();
266                            let prompt = zeta1::format_zeta1_from_input(
267                                &prompt_input,
268                                editable_range_in_excerpt.clone(),
269                                ranges.editable_350_context_150.clone(),
270                            );
271                            let stop_tokens = vec![
272                                EDITABLE_REGION_END_MARKER.to_string(),
273                                format!("{EDITABLE_REGION_END_MARKER}\n"),
274                                format!("{EDITABLE_REGION_END_MARKER}\n\n"),
275                                format!("{EDITABLE_REGION_END_MARKER}\n\n\n"),
276                            ];
277
278                            let (response_text, request_id) = send_custom_server_request(
279                                provider,
280                                custom_settings,
281                                prompt,
282                                max_tokens,
283                                stop_tokens,
284                                open_ai_compatible_api_key.clone(),
285                                &http_client,
286                            )
287                            .await?;
288
289                            let request_id = EditPredictionId(request_id.into());
290                            let output_text = zeta1::clean_zeta1_model_output(&response_text);
291                            let parsed_output = output_text.map(|text| ParsedOutput {
292                                new_editable_region: text,
293                                range_in_excerpt: editable_range_in_excerpt,
294                                cursor_offset_in_new_editable_region: None,
295                            });
296
297                            (request_id, parsed_output, None, None)
298                        }
299                        ZetaVersion::Zeta2 | ZetaVersion::Zeta2_1 => {
300                            let Some(prompt) = formatted_prompt.clone() else {
301                                return Ok((None, None));
302                            };
303                            let prefill = get_prefill(&prompt_input, zeta_format);
304                            let prompt = format!("{prompt}{prefill}");
305
306                            let (response_text, request_id) = send_custom_server_request(
307                                provider,
308                                custom_settings,
309                                prompt,
310                                max_tokens,
311                                stop_tokens_for_format(zeta_format)
312                                    .iter()
313                                    .map(|token| token.to_string())
314                                    .collect(),
315                                open_ai_compatible_api_key.clone(),
316                                &http_client,
317                            )
318                            .await?;
319
320                            let request_id = EditPredictionId(request_id.into());
321                            let output_text = if response_text.is_empty() {
322                                None
323                            } else {
324                                let output = format!("{prefill}{response_text}");
325                                Some(parse_zeta2_model_output(
326                                    &output,
327                                    zeta_format,
328                                    &prompt_input,
329                                )?)
330                            };
331
332                            (request_id, output_text, None, None)
333                        }
334                    })
335                } else if let Some(config) = &raw_config {
336                    let Some(prompt) = format_zeta_prompt(&prompt_input, config.format) else {
337                        return Ok((None, None));
338                    };
339                    let prefill = get_prefill(&prompt_input, config.format);
340                    let prompt = format!("{prompt}{prefill}");
341                    let environment = config
342                        .environment
343                        .clone()
344                        .or_else(|| Some(config.format.to_string().to_lowercase()));
345                    let request = RawCompletionRequest {
346                        model: config.model_id.clone().unwrap_or_default(),
347                        prompt,
348                        temperature: None,
349                        stop: stop_tokens_for_format(config.format)
350                            .iter()
351                            .map(|token| std::borrow::Cow::Borrowed(*token))
352                            .collect(),
353                        max_tokens: Some(2048),
354                        environment,
355                    };
356
357                    let (mut response, usage) = EditPredictionStore::send_raw_llm_request(
358                        request,
359                        client,
360                        None,
361                        llm_token,
362                        organization_id,
363                        app_version,
364                    )
365                    .await?;
366
367                    let request_id = EditPredictionId(response.id.clone().into());
368                    let output = if let Some(choice) = response.choices.pop() {
369                        let response = &choice.text;
370                        let output = format!("{prefill}{response}");
371                        Some(parse_zeta2_model_output(
372                            &output,
373                            config.format,
374                            &prompt_input,
375                        )?)
376                    } else {
377                        None
378                    };
379
380                    Some((request_id, output, None, usage))
381                } else {
382                    let (response, usage) = EditPredictionStore::send_v3_request(
383                        prompt_input.clone(),
384                        preferred_experiment.clone(),
385                        client,
386                        llm_token,
387                        organization_id,
388                        app_version,
389                        trigger,
390                        mode,
391                    )
392                    .await?;
393
394                    let request_id = EditPredictionId(response.request_id.into());
395                    let model_version = response.model_version;
396                    let parsed_output = ParsedOutput {
397                        new_editable_region: response.output,
398                        range_in_excerpt: response.editable_range,
399                        cursor_offset_in_new_editable_region: response.cursor_offset,
400                    };
401
402                    Some((request_id, Some(parsed_output), model_version, usage))
403                })
404            else {
405                return Ok((None, None));
406            };
407
408            log::trace!("Got edit prediction response");
409
410            let (new_editable_region, cursor_offset_in_output, editable_range_in_excerpt) =
411                if let Some(output) = output {
412                    let ParsedOutput {
413                        new_editable_region: output_text,
414                        range_in_excerpt: editable_range_in_excerpt,
415                        cursor_offset_in_new_editable_region: cursor_offset_in_output,
416                    } = output;
417                    (
418                        Some(output_text),
419                        cursor_offset_in_output,
420                        editable_range_in_excerpt,
421                    )
422                } else {
423                    let (editable_range, _) =
424                        excerpt_ranges_for_format(zeta_format, &prompt_input.excerpt_ranges);
425                    (None, None, editable_range)
426                };
427
428            let editable_range_in_buffer = editable_range_in_excerpt.start
429                + full_context_offset_range.start
430                ..editable_range_in_excerpt.end + full_context_offset_range.start;
431
432            if let Some(debug_tx) = &debug_tx {
433                debug_tx
434                    .unbounded_send(DebugEvent::EditPredictionFinished(
435                        EditPredictionFinishedDebugEvent {
436                            buffer: buffer.downgrade(),
437                            position,
438                            model_output: new_editable_region.clone(),
439                        },
440                    ))
441                    .ok();
442            }
443            let (edits, cursor_position) = new_editable_region
444                .map(|mut output_text| {
445                    let mut old_text = snapshot
446                        .text_for_range(editable_range_in_buffer.clone())
447                        .collect::<String>();
448
449                    if !output_text.is_empty() && !output_text.ends_with('\n') {
450                        output_text.push('\n');
451                    }
452                    if !old_text.is_empty() && !old_text.ends_with('\n') {
453                        old_text.push('\n');
454                    }
455
456                    compute_edits_and_cursor_position(
457                        old_text,
458                        &output_text,
459                        editable_range_in_buffer.start,
460                        cursor_offset_in_output,
461                        &snapshot,
462                    )
463                })
464                .unwrap_or_default();
465
466            let prediction = Some(PendingPrediction::Resolved {
467                inputs: EditPredictionInputs::V2(prompt_input),
468                buffer,
469                snapshot: snapshot.clone(),
470                edits,
471                cursor_position,
472                editable_range_in_buffer: Some(editable_range_in_buffer),
473            });
474
475            anyhow::Ok((Some((request_id, prediction, model_version)), usage))
476        }
477    });
478
479    cx.spawn(async move |this, cx| {
480        let Some((id, prediction, model_version)) =
481            handle_api_response(&this, request_task.await, cx)?
482        else {
483            return Ok(None);
484        };
485        let request_duration = cx.background_executor().now() - request_start;
486
487        let Some(prediction) = prediction else {
488            return Ok(None);
489        };
490
491        let (
492            inputs,
493            edited_buffer,
494            edited_buffer_snapshot,
495            edits,
496            cursor_position,
497            editable_range_in_buffer,
498        ) = match prediction {
499            PendingPrediction::Resolved {
500                inputs,
501                buffer,
502                snapshot,
503                edits,
504                cursor_position,
505                editable_range_in_buffer,
506            } => (
507                inputs,
508                buffer,
509                snapshot,
510                edits,
511                cursor_position,
512                editable_range_in_buffer,
513            ),
514            PendingPrediction::Patch {
515                inputs,
516                project,
517                buffer: fallback_buffer,
518                snapshot: fallback_snapshot,
519                patch,
520            } => {
521                let (buffer, snapshot, edits, cursor_position) =
522                    match prediction_edits_for_single_file_diff(&patch, &project, cx).await {
523                        Ok(Some(edits)) => edits,
524                        Ok(None) => {
525                            return Ok(Some(
526                                EditPredictionResult::new(
527                                    id,
528                                    &fallback_buffer,
529                                    &fallback_snapshot,
530                                    Arc::new([]),
531                                    None,
532                                    None,
533                                    inputs,
534                                    model_version,
535                                    trigger,
536                                    request_duration,
537                                    cx,
538                                )
539                                .await,
540                            ));
541                        }
542                        Err(error) => {
543                            log::error!("failed to apply edit prediction patch: {error:?}");
544                            return Ok(Some(EditPredictionResult::new_rejected(
545                                id,
546                                &fallback_buffer,
547                                &fallback_snapshot,
548                                inputs,
549                                model_version,
550                                trigger,
551                                request_duration,
552                                EditPredictionRejectReason::PatchApplyFailed,
553                            )));
554                        }
555                    };
556                let editable_range_in_buffer =
557                    edits
558                        .iter()
559                        .fold(None, |range: Option<Range<usize>>, (edit_range, _)| {
560                            let edit_range = edit_range.start.to_offset(&snapshot)
561                                ..edit_range.end.to_offset(&snapshot);
562                            Some(match range {
563                                Some(range) => {
564                                    range.start.min(edit_range.start)..range.end.max(edit_range.end)
565                                }
566                                None => edit_range,
567                            })
568                        });
569
570                (
571                    inputs,
572                    buffer,
573                    snapshot,
574                    edits,
575                    cursor_position,
576                    editable_range_in_buffer,
577                )
578            }
579        };
580
581        let result = EditPredictionResult::new(
582            id,
583            &edited_buffer,
584            &edited_buffer_snapshot,
585            edits.into(),
586            cursor_position,
587            editable_range_in_buffer
588                .clone()
589                .map(|range| edited_buffer_snapshot.anchor_range_inside(range)),
590            inputs,
591            model_version,
592            trigger,
593            request_duration,
594            cx,
595        )
596        .await;
597
598        if let Some(editable_range_in_buffer) = editable_range_in_buffer.clone() {
599            let prediction = &result.prediction;
600            let weak_this = this.clone();
601            let request_id = prediction.id.clone();
602            let edited_buffer = edited_buffer.clone();
603            let edited_buffer_snapshot = edited_buffer_snapshot.clone();
604            let edit_preview = prediction.edit_preview.clone();
605            let model_version = prediction.model_version.clone();
606            cx.spawn(async move |cx| {
607                weak_this
608                    .update(cx, |this, cx| {
609                        this.enqueue_settled_prediction(
610                            request_id.clone(),
611                            &project,
612                            &edited_buffer,
613                            &edited_buffer_snapshot,
614                            editable_range_in_buffer,
615                            &edit_preview,
616                            context_task,
617                            prompt_history_boundary,
618                            model_version,
619                            request_duration,
620                            cx,
621                        );
622                    })
623                    .ok();
624            })
625            .detach();
626        }
627
628        Ok(Some(result))
629    })
630}
631
632fn handle_api_response<T>(
633    this: &WeakEntity<EditPredictionStore>,
634    response: Result<(T, Option<client::EditPredictionUsage>)>,
635    cx: &mut gpui::AsyncApp,
636) -> Result<T> {
637    match response {
638        Ok((data, usage)) => {
639            if let Some(usage) = usage {
640                this.update(cx, |this, cx| {
641                    this.user_store.update(cx, |user_store, cx| {
642                        user_store.update_edit_prediction_usage(usage, cx);
643                    });
644                })
645                .ok();
646            }
647            Ok(data)
648        }
649        Err(err) => {
650            if err.is::<CloudRequestTimeoutError>() {
651                this.update(cx, |this, cx| this.back_off_requests_after_timeout(cx))
652                    .ok();
653            }
654
655            if err.is::<ZedUpdateRequiredError>() {
656                cx.update(|cx| {
657                    this.update(cx, |this, _cx| {
658                        this.update_required = true;
659                    })
660                    .ok();
661
662                    let message: SharedString = err.to_string().into();
663
664                    struct UpdateRequiredError {
665                        message: SharedString,
666                    }
667                    impl WorkspaceError for UpdateRequiredError {
668                        fn primary_message(&self) -> SharedString {
669                            self.message.clone()
670                        }
671                        fn severity(&self) -> ErrorSeverity {
672                            ErrorSeverity::Critical
673                        }
674                        fn primary_action(&self) -> ErrorAction {
675                            ErrorAction::link(
676                                "Update Omega",
677                                "https://github.com/OpenAgentsInc/omega/releases",
678                            )
679                        }
680                    }
681
682                    show_app_notification(
683                        NotificationId::unique::<ZedUpdateRequiredError>(),
684                        cx,
685                        move |cx| {
686                            cx.new({
687                                let message = message.clone();
688                                move |cx| {
689                                    let error = UpdateRequiredError { message };
690                                    MessageNotification::from_workspace_error(error, cx)
691                                }
692                            })
693                        },
694                    );
695                });
696            }
697            Err(err)
698        }
699    }
700}
701
702const ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT: usize = 100;
703const MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT: usize = 20;
704pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT: usize = 512;
705pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512;
706
707pub(crate) fn active_buffer_diagnostics(
708    snapshot: &language::BufferSnapshot,
709    diagnostic_search_range: Range<Point>,
710    cursor_row: u32,
711    additional_context_token_count: usize,
712) -> Vec<zeta_prompt::ActiveBufferDiagnostic> {
713    let mut diagnostics = snapshot
714        .diagnostics_in_range::<Point, Point>(diagnostic_search_range, false)
715        .collect::<Vec<_>>();
716    diagnostics.sort_by_key(|entry| {
717        cursor_row.abs_diff(entry.range.start.row) + cursor_row.abs_diff(entry.range.end.row)
718    });
719
720    diagnostics
721        .into_iter()
722        .map(|entry| {
723            let diagnostic_point_range = entry.range.clone();
724            let snippet_point_range = cursor_excerpt::expand_context_syntactically_then_linewise(
725                snapshot,
726                diagnostic_point_range.clone(),
727                additional_context_token_count,
728            );
729
730            let severity = match entry.diagnostic.severity {
731                DiagnosticSeverity::ERROR => Some(1),
732                DiagnosticSeverity::WARNING => Some(2),
733                DiagnosticSeverity::INFORMATION => Some(3),
734                DiagnosticSeverity::HINT => Some(4),
735                _ => None,
736            };
737            (
738                severity,
739                zeta_prompt::clamp_text_to_token_count(
740                    &entry.diagnostic.message,
741                    MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT,
742                )
743                .to_string(),
744                diagnostic_point_range,
745                snippet_point_range,
746            )
747        })
748        .take(MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT)
749        .map(
750            |(severity, message, diagnostic_point_range, snippet_point_range)| {
751                let (snippet, diagnostic_range_in_snippet) = if snippet_point_range.start
752                    == Point::new(0, 0)
753                    && snippet_point_range.end == snapshot.max_point()
754                {
755                    (String::new(), 0..0)
756                } else {
757                    let snippet = snapshot
758                        .text_for_range(snippet_point_range.clone())
759                        .collect::<String>();
760                    let snippet = zeta_prompt::clamp_text_to_token_count(
761                        &snippet,
762                        MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT,
763                    )
764                    .to_string();
765                    let snippet_start_offset = snippet_point_range.start.to_offset(snapshot);
766                    let diagnostic_offset_range = diagnostic_point_range.to_offset(snapshot);
767                    let diagnostic_range_start = diagnostic_offset_range
768                        .start
769                        .saturating_sub(snippet_start_offset)
770                        .min(snippet.len());
771                    let diagnostic_range_end = diagnostic_offset_range
772                        .end
773                        .saturating_sub(snippet_start_offset)
774                        .min(snippet.len());
775                    (snippet, diagnostic_range_start..diagnostic_range_end)
776                };
777                zeta_prompt::ActiveBufferDiagnostic {
778                    severity,
779                    message,
780                    snippet,
781                    snippet_buffer_row_range: diagnostic_point_range.start.row
782                        ..diagnostic_point_range.end.row,
783                    diagnostic_range_in_snippet,
784                }
785            },
786        )
787        .collect()
788}
789
790pub fn zeta2_prompt_input(
791    snapshot: &language::BufferSnapshot,
792    mut related_files: Vec<zeta_prompt::RelatedFile>,
793    events: Vec<Arc<zeta_prompt::Event>>,
794    diagnostic_search_range: Range<Point>,
795    excerpt_path: Arc<Path>,
796    cursor_offset: usize,
797    is_open_source: bool,
798    can_collect_data: bool,
799    repo_url: Option<String>,
800) -> (Range<usize>, zeta_prompt::Zeta2PromptInput) {
801    let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
802        compute_cursor_excerpt(snapshot, cursor_offset);
803
804    let cursor_excerpt: Arc<str> = snapshot
805        .text_for_range(excerpt_point_range.clone())
806        .collect::<String>()
807        .into();
808    let syntax_ranges = compute_syntax_ranges(snapshot, cursor_offset, &excerpt_offset_range);
809    let excerpt_ranges = zeta_prompt::compute_legacy_excerpt_ranges(
810        &cursor_excerpt,
811        cursor_offset_in_excerpt,
812        &syntax_ranges,
813    );
814
815    let active_buffer_diagnostics = active_buffer_diagnostics(
816        snapshot,
817        diagnostic_search_range,
818        snapshot.offset_to_point(cursor_offset).row,
819        ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT,
820    );
821    for file in &mut related_files {
822        for excerpt in &mut file.excerpts {
823            excerpt.context_source = zeta_prompt::ContextSource::Lsp;
824        }
825    }
826
827    let prompt_input = zeta_prompt::Zeta2PromptInput {
828        cursor_path: excerpt_path,
829        cursor_excerpt,
830        cursor_offset_in_excerpt,
831        excerpt_start_row: Some(excerpt_point_range.start.row),
832        events,
833        related_files: Some(related_files),
834        active_buffer_diagnostics,
835        excerpt_ranges,
836        syntax_ranges: Some(syntax_ranges),
837        in_open_source_repo: is_open_source,
838        can_collect_data,
839        repo_url,
840    };
841    (excerpt_offset_range, prompt_input)
842}
843
844pub(crate) fn edit_prediction_accepted(
845    store: &EditPredictionStore,
846    current_prediction: CurrentEditPrediction,
847    cx: &App,
848) {
849    if store.zeta2_raw_config().is_some() {
850        return;
851    }
852
853    let request_id = current_prediction.prediction.id.to_string();
854    let model_version = current_prediction.prediction.model_version;
855    let e2e_latency = current_prediction.e2e_latency;
856    let client = store.client.clone();
857    let llm_token = store.llm_token.clone();
858    let organization_id = store
859        .user_store
860        .read(cx)
861        .current_organization()
862        .map(|organization| organization.id.clone());
863    let app_version = AppVersion::global(cx);
864
865    cx.background_spawn(async move {
866        let body = serde_json::to_string(&AcceptEditPredictionBody {
867            request_id,
868            model_version,
869            e2e_latency_ms: Some(e2e_latency.as_millis()),
870        })?;
871
872        let url = client
873            .http_client()
874            .build_zed_llm_url("/predict_edits/accept", &[])?;
875        EditPredictionStore::send_api_request::<()>(
876            move |builder| Ok(builder.uri(url.as_ref()).body(body.clone().into())?),
877            client,
878            llm_token,
879            organization_id,
880            app_version,
881        )
882        .await?;
883        anyhow::Ok(())
884    })
885    .detach_and_log_err(cx);
886}
887
888pub fn compute_edits(
889    old_text: String,
890    new_text: &str,
891    offset: usize,
892    snapshot: &BufferSnapshot,
893) -> Vec<(Range<Anchor>, Arc<str>)> {
894    compute_edits_and_cursor_position(old_text, new_text, offset, None, snapshot).0
895}
896
897pub fn compute_edits_and_cursor_position(
898    old_text: String,
899    new_text: &str,
900    offset: usize,
901    cursor_offset_in_new_text: Option<usize>,
902    snapshot: &BufferSnapshot,
903) -> (
904    Vec<(Range<Anchor>, Arc<str>)>,
905    Option<PredictedCursorPosition>,
906) {
907    let diffs = text_diff(&old_text, new_text);
908
909    // Delta represents the cumulative change in byte count from all preceding edits.
910    // new_offset = old_offset + delta, so old_offset = new_offset - delta
911    let mut delta: isize = 0;
912    let mut cursor_position: Option<PredictedCursorPosition> = None;
913    let buffer_len = snapshot.len();
914
915    let edits = diffs
916        .iter()
917        .map(|(raw_old_range, new_text)| {
918            // Compute cursor position if it falls within or before this edit.
919            if let (Some(cursor_offset), None) = (cursor_offset_in_new_text, cursor_position) {
920                let edit_start_in_new = (raw_old_range.start as isize + delta) as usize;
921                let edit_end_in_new = edit_start_in_new + new_text.len();
922
923                if cursor_offset < edit_start_in_new {
924                    let cursor_in_old = (cursor_offset as isize - delta) as usize;
925                    let buffer_offset = (offset + cursor_in_old).min(buffer_len);
926                    cursor_position = Some(PredictedCursorPosition::at_anchor(
927                        snapshot.anchor_after(buffer_offset),
928                    ));
929                } else if cursor_offset < edit_end_in_new {
930                    let buffer_offset = (offset + raw_old_range.start).min(buffer_len);
931                    let offset_within_insertion = cursor_offset - edit_start_in_new;
932                    cursor_position = Some(PredictedCursorPosition::new(
933                        snapshot.anchor_before(buffer_offset),
934                        offset_within_insertion,
935                    ));
936                }
937
938                delta += new_text.len() as isize - raw_old_range.len() as isize;
939            }
940
941            // Compute the edit with prefix/suffix trimming.
942            let mut old_range = raw_old_range.clone();
943            let old_slice = &old_text[old_range.clone()];
944
945            let prefix_len = common_prefix(old_slice.chars(), new_text.chars());
946            let suffix_len = common_prefix(
947                old_slice[prefix_len..].chars().rev(),
948                new_text[prefix_len..].chars().rev(),
949            );
950
951            old_range.start += offset;
952            old_range.end += offset;
953            old_range.start += prefix_len;
954            old_range.end -= suffix_len;
955
956            old_range.start = old_range.start.min(buffer_len);
957            old_range.end = old_range.end.min(buffer_len);
958
959            let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
960            let range = if old_range.is_empty() {
961                let anchor = snapshot.anchor_after(old_range.start);
962                anchor..anchor
963            } else {
964                snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
965            };
966            (range, new_text)
967        })
968        .collect();
969
970    if let (Some(cursor_offset), None) = (cursor_offset_in_new_text, cursor_position) {
971        let cursor_in_old = (cursor_offset as isize - delta) as usize;
972        let buffer_offset = snapshot.clip_offset(offset + cursor_in_old, Bias::Right);
973        cursor_position = Some(PredictedCursorPosition::at_anchor(
974            snapshot.anchor_after(buffer_offset),
975        ));
976    }
977
978    (edits, cursor_position)
979}
980
981fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
982    a.zip(b)
983        .take_while(|(a, b)| a == b)
984        .map(|(a, _)| a.len_utf8())
985        .sum()
986}
987