1mod completion_diff_element;
   2mod init;
   3mod input_excerpt;
   4mod license_detection;
   5mod onboarding_modal;
   6mod onboarding_telemetry;
   7mod rate_completion_modal;
   8
   9pub(crate) use completion_diff_element::*;
  10use db::kvp::{Dismissable, KEY_VALUE_STORE};
  11pub use init::*;
  12use inline_completion::DataCollectionState;
  13use license_detection::LICENSE_FILES_TO_CHECK;
  14pub use license_detection::is_license_eligible_for_data_collection;
  15pub use rate_completion_modal::*;
  16
  17use anyhow::{Context as _, Result, anyhow};
  18use arrayvec::ArrayVec;
  19use client::{Client, CloudUserStore, EditPredictionUsage};
  20use cloud_llm_client::{
  21    AcceptEditPredictionBody, EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
  22    PredictEditsBody, PredictEditsResponse, ZED_VERSION_HEADER_NAME,
  23};
  24use collections::{HashMap, HashSet, VecDeque};
  25use futures::AsyncReadExt;
  26use gpui::{
  27    App, AppContext as _, AsyncApp, Context, Entity, EntityId, Global, SemanticVersion,
  28    Subscription, Task, WeakEntity, actions,
  29};
  30use http_client::{AsyncBody, HttpClient, Method, Request, Response};
  31use input_excerpt::excerpt_for_cursor_position;
  32use language::{
  33    Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, ToOffset, ToPoint, text_diff,
  34};
  35use language_model::{LlmApiToken, RefreshLlmTokenListener};
  36use postage::watch;
  37use project::Project;
  38use release_channel::AppVersion;
  39use settings::WorktreeId;
  40use std::str::FromStr;
  41use std::{
  42    borrow::Cow,
  43    cmp,
  44    fmt::Write,
  45    future::Future,
  46    mem,
  47    ops::Range,
  48    path::Path,
  49    rc::Rc,
  50    sync::Arc,
  51    time::{Duration, Instant},
  52};
  53use telemetry_events::InlineCompletionRating;
  54use thiserror::Error;
  55use util::ResultExt;
  56use uuid::Uuid;
  57use workspace::Workspace;
  58use workspace::notifications::{ErrorMessagePrompt, NotificationId};
  59use worktree::Worktree;
  60
  61const CURSOR_MARKER: &'static str = "<|user_cursor_is_here|>";
  62const START_OF_FILE_MARKER: &'static str = "<|start_of_file|>";
  63const EDITABLE_REGION_START_MARKER: &'static str = "<|editable_region_start|>";
  64const EDITABLE_REGION_END_MARKER: &'static str = "<|editable_region_end|>";
  65const BUFFER_CHANGE_GROUPING_INTERVAL: Duration = Duration::from_secs(1);
  66const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice";
  67
  68const MAX_CONTEXT_TOKENS: usize = 150;
  69const MAX_REWRITE_TOKENS: usize = 350;
  70const MAX_EVENT_TOKENS: usize = 500;
  71
  72/// Maximum number of events to track.
  73const MAX_EVENT_COUNT: usize = 16;
  74
  75actions!(
  76    edit_prediction,
  77    [
  78        /// Clears the edit prediction history.
  79        ClearHistory
  80    ]
  81);
  82
  83#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
  84pub struct InlineCompletionId(Uuid);
  85
  86impl From<InlineCompletionId> for gpui::ElementId {
  87    fn from(value: InlineCompletionId) -> Self {
  88        gpui::ElementId::Uuid(value.0)
  89    }
  90}
  91
  92impl std::fmt::Display for InlineCompletionId {
  93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
  94        write!(f, "{}", self.0)
  95    }
  96}
  97
  98struct ZedPredictUpsell;
  99
 100impl Dismissable for ZedPredictUpsell {
 101    const KEY: &'static str = "dismissed-edit-predict-upsell";
 102
 103    fn dismissed() -> bool {
 104        // To make this backwards compatible with older versions of Zed, we
 105        // check if the user has seen the previous Edit Prediction Onboarding
 106        // before, by checking the data collection choice which was written to
 107        // the database once the user clicked on "Accept and Enable"
 108        if KEY_VALUE_STORE
 109            .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
 110            .log_err()
 111            .map_or(false, |s| s.is_some())
 112        {
 113            return true;
 114        }
 115
 116        KEY_VALUE_STORE
 117            .read_kvp(Self::KEY)
 118            .log_err()
 119            .map_or(false, |s| s.is_some())
 120    }
 121}
 122
 123pub fn should_show_upsell_modal(cloud_user_store: &Entity<CloudUserStore>, cx: &App) -> bool {
 124    if cloud_user_store.read(cx).has_accepted_tos() {
 125        !ZedPredictUpsell::dismissed()
 126    } else {
 127        true
 128    }
 129}
 130
 131#[derive(Clone)]
 132struct ZetaGlobal(Entity<Zeta>);
 133
 134impl Global for ZetaGlobal {}
 135
 136#[derive(Clone)]
 137pub struct InlineCompletion {
 138    id: InlineCompletionId,
 139    path: Arc<Path>,
 140    excerpt_range: Range<usize>,
 141    cursor_offset: usize,
 142    edits: Arc<[(Range<Anchor>, String)]>,
 143    snapshot: BufferSnapshot,
 144    edit_preview: EditPreview,
 145    input_outline: Arc<str>,
 146    input_events: Arc<str>,
 147    input_excerpt: Arc<str>,
 148    output_excerpt: Arc<str>,
 149    request_sent_at: Instant,
 150    response_received_at: Instant,
 151}
 152
 153impl InlineCompletion {
 154    fn latency(&self) -> Duration {
 155        self.response_received_at
 156            .duration_since(self.request_sent_at)
 157    }
 158
 159    fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, String)>> {
 160        interpolate(&self.snapshot, new_snapshot, self.edits.clone())
 161    }
 162}
 163
 164fn interpolate(
 165    old_snapshot: &BufferSnapshot,
 166    new_snapshot: &BufferSnapshot,
 167    current_edits: Arc<[(Range<Anchor>, String)]>,
 168) -> Option<Vec<(Range<Anchor>, String)>> {
 169    let mut edits = Vec::new();
 170
 171    let mut model_edits = current_edits.into_iter().peekable();
 172    for user_edit in new_snapshot.edits_since::<usize>(&old_snapshot.version) {
 173        while let Some((model_old_range, _)) = model_edits.peek() {
 174            let model_old_range = model_old_range.to_offset(old_snapshot);
 175            if model_old_range.end < user_edit.old.start {
 176                let (model_old_range, model_new_text) = model_edits.next().unwrap();
 177                edits.push((model_old_range.clone(), model_new_text.clone()));
 178            } else {
 179                break;
 180            }
 181        }
 182
 183        if let Some((model_old_range, model_new_text)) = model_edits.peek() {
 184            let model_old_offset_range = model_old_range.to_offset(old_snapshot);
 185            if user_edit.old == model_old_offset_range {
 186                let user_new_text = new_snapshot
 187                    .text_for_range(user_edit.new.clone())
 188                    .collect::<String>();
 189
 190                if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
 191                    if !model_suffix.is_empty() {
 192                        let anchor = old_snapshot.anchor_after(user_edit.old.end);
 193                        edits.push((anchor..anchor, model_suffix.to_string()));
 194                    }
 195
 196                    model_edits.next();
 197                    continue;
 198                }
 199            }
 200        }
 201
 202        return None;
 203    }
 204
 205    edits.extend(model_edits.cloned());
 206
 207    if edits.is_empty() { None } else { Some(edits) }
 208}
 209
 210impl std::fmt::Debug for InlineCompletion {
 211    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 212        f.debug_struct("InlineCompletion")
 213            .field("id", &self.id)
 214            .field("path", &self.path)
 215            .field("edits", &self.edits)
 216            .finish_non_exhaustive()
 217    }
 218}
 219
 220pub struct Zeta {
 221    workspace: Option<WeakEntity<Workspace>>,
 222    client: Arc<Client>,
 223    events: VecDeque<Event>,
 224    registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,
 225    shown_completions: VecDeque<InlineCompletion>,
 226    rated_completions: HashSet<InlineCompletionId>,
 227    data_collection_choice: Entity<DataCollectionChoice>,
 228    llm_token: LlmApiToken,
 229    _llm_token_subscription: Subscription,
 230    /// Whether an update to a newer version of Zed is required to continue using Zeta.
 231    update_required: bool,
 232    cloud_user_store: Entity<CloudUserStore>,
 233    license_detection_watchers: HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
 234}
 235
 236impl Zeta {
 237    pub fn global(cx: &mut App) -> Option<Entity<Self>> {
 238        cx.try_global::<ZetaGlobal>().map(|global| global.0.clone())
 239    }
 240
 241    pub fn register(
 242        workspace: Option<WeakEntity<Workspace>>,
 243        worktree: Option<Entity<Worktree>>,
 244        client: Arc<Client>,
 245        cloud_user_store: Entity<CloudUserStore>,
 246        cx: &mut App,
 247    ) -> Entity<Self> {
 248        let this = Self::global(cx).unwrap_or_else(|| {
 249            let entity = cx.new(|cx| Self::new(workspace, client, cloud_user_store, cx));
 250            cx.set_global(ZetaGlobal(entity.clone()));
 251            entity
 252        });
 253
 254        this.update(cx, move |this, cx| {
 255            if let Some(worktree) = worktree {
 256                worktree.update(cx, |worktree, cx| {
 257                    this.license_detection_watchers
 258                        .entry(worktree.id())
 259                        .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(worktree, cx)));
 260                });
 261            }
 262        });
 263
 264        this
 265    }
 266
 267    pub fn clear_history(&mut self) {
 268        self.events.clear();
 269    }
 270
 271    pub fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
 272        self.cloud_user_store.read(cx).edit_prediction_usage()
 273    }
 274
 275    fn new(
 276        workspace: Option<WeakEntity<Workspace>>,
 277        client: Arc<Client>,
 278        cloud_user_store: Entity<CloudUserStore>,
 279        cx: &mut Context<Self>,
 280    ) -> Self {
 281        let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
 282
 283        let data_collection_choice = Self::load_data_collection_choices();
 284        let data_collection_choice = cx.new(|_| data_collection_choice);
 285
 286        Self {
 287            workspace,
 288            client,
 289            events: VecDeque::new(),
 290            shown_completions: VecDeque::new(),
 291            rated_completions: HashSet::default(),
 292            registered_buffers: HashMap::default(),
 293            data_collection_choice,
 294            llm_token: LlmApiToken::default(),
 295            _llm_token_subscription: cx.subscribe(
 296                &refresh_llm_token_listener,
 297                |this, _listener, _event, cx| {
 298                    let client = this.client.clone();
 299                    let llm_token = this.llm_token.clone();
 300                    cx.spawn(async move |_this, _cx| {
 301                        llm_token.refresh(&client).await?;
 302                        anyhow::Ok(())
 303                    })
 304                    .detach_and_log_err(cx);
 305                },
 306            ),
 307            update_required: false,
 308            license_detection_watchers: HashMap::default(),
 309            cloud_user_store,
 310        }
 311    }
 312
 313    fn push_event(&mut self, event: Event) {
 314        if let Some(Event::BufferChange {
 315            new_snapshot: last_new_snapshot,
 316            timestamp: last_timestamp,
 317            ..
 318        }) = self.events.back_mut()
 319        {
 320            // Coalesce edits for the same buffer when they happen one after the other.
 321            let Event::BufferChange {
 322                old_snapshot,
 323                new_snapshot,
 324                timestamp,
 325            } = &event;
 326
 327            if timestamp.duration_since(*last_timestamp) <= BUFFER_CHANGE_GROUPING_INTERVAL
 328                && old_snapshot.remote_id() == last_new_snapshot.remote_id()
 329                && old_snapshot.version == last_new_snapshot.version
 330            {
 331                *last_new_snapshot = new_snapshot.clone();
 332                *last_timestamp = *timestamp;
 333                return;
 334            }
 335        }
 336
 337        self.events.push_back(event);
 338        if self.events.len() >= MAX_EVENT_COUNT {
 339            // These are halved instead of popping to improve prompt caching.
 340            self.events.drain(..MAX_EVENT_COUNT / 2);
 341        }
 342    }
 343
 344    pub fn register_buffer(&mut self, buffer: &Entity<Buffer>, cx: &mut Context<Self>) {
 345        let buffer_id = buffer.entity_id();
 346        let weak_buffer = buffer.downgrade();
 347
 348        if let std::collections::hash_map::Entry::Vacant(entry) =
 349            self.registered_buffers.entry(buffer_id)
 350        {
 351            let snapshot = buffer.read(cx).snapshot();
 352
 353            entry.insert(RegisteredBuffer {
 354                snapshot,
 355                _subscriptions: [
 356                    cx.subscribe(buffer, move |this, buffer, event, cx| {
 357                        this.handle_buffer_event(buffer, event, cx);
 358                    }),
 359                    cx.observe_release(buffer, move |this, _buffer, _cx| {
 360                        this.registered_buffers.remove(&weak_buffer.entity_id());
 361                    }),
 362                ],
 363            });
 364        };
 365    }
 366
 367    fn handle_buffer_event(
 368        &mut self,
 369        buffer: Entity<Buffer>,
 370        event: &language::BufferEvent,
 371        cx: &mut Context<Self>,
 372    ) {
 373        if let language::BufferEvent::Edited = event {
 374            self.report_changes_for_buffer(&buffer, cx);
 375        }
 376    }
 377
 378    fn request_completion_impl<F, R>(
 379        &mut self,
 380        workspace: Option<Entity<Workspace>>,
 381        project: Option<&Entity<Project>>,
 382        buffer: &Entity<Buffer>,
 383        cursor: language::Anchor,
 384        can_collect_data: bool,
 385        cx: &mut Context<Self>,
 386        perform_predict_edits: F,
 387    ) -> Task<Result<Option<InlineCompletion>>>
 388    where
 389        F: FnOnce(PerformPredictEditsParams) -> R + 'static,
 390        R: Future<Output = Result<(PredictEditsResponse, Option<EditPredictionUsage>)>>
 391            + Send
 392            + 'static,
 393    {
 394        let snapshot = self.report_changes_for_buffer(&buffer, cx);
 395        let diagnostic_groups = snapshot.diagnostic_groups(None);
 396        let cursor_point = cursor.to_point(&snapshot);
 397        let cursor_offset = cursor_point.to_offset(&snapshot);
 398        let events = self.events.clone();
 399        let path: Arc<Path> = snapshot
 400            .file()
 401            .map(|f| Arc::from(f.full_path(cx).as_path()))
 402            .unwrap_or_else(|| Arc::from(Path::new("untitled")));
 403
 404        let zeta = cx.entity();
 405        let client = self.client.clone();
 406        let llm_token = self.llm_token.clone();
 407        let app_version = AppVersion::global(cx);
 408
 409        let buffer = buffer.clone();
 410
 411        let local_lsp_store =
 412            project.and_then(|project| project.read(cx).lsp_store().read(cx).as_local());
 413        let diagnostic_groups = if let Some(local_lsp_store) = local_lsp_store {
 414            Some(
 415                diagnostic_groups
 416                    .into_iter()
 417                    .filter_map(|(language_server_id, diagnostic_group)| {
 418                        let language_server =
 419                            local_lsp_store.running_language_server_for_id(language_server_id)?;
 420
 421                        Some((
 422                            language_server.name(),
 423                            diagnostic_group.resolve::<usize>(&snapshot),
 424                        ))
 425                    })
 426                    .collect::<Vec<_>>(),
 427            )
 428        } else {
 429            None
 430        };
 431
 432        cx.spawn(async move |this, cx| {
 433            let request_sent_at = Instant::now();
 434
 435            struct BackgroundValues {
 436                input_events: String,
 437                input_excerpt: String,
 438                speculated_output: String,
 439                editable_range: Range<usize>,
 440                input_outline: String,
 441            }
 442
 443            let values = cx
 444                .background_spawn({
 445                    let snapshot = snapshot.clone();
 446                    let path = path.clone();
 447                    async move {
 448                        let path = path.to_string_lossy();
 449                        let input_excerpt = excerpt_for_cursor_position(
 450                            cursor_point,
 451                            &path,
 452                            &snapshot,
 453                            MAX_REWRITE_TOKENS,
 454                            MAX_CONTEXT_TOKENS,
 455                        );
 456                        let input_events = prompt_for_events(&events, MAX_EVENT_TOKENS);
 457                        let input_outline = prompt_for_outline(&snapshot);
 458
 459                        anyhow::Ok(BackgroundValues {
 460                            input_events,
 461                            input_excerpt: input_excerpt.prompt,
 462                            speculated_output: input_excerpt.speculated_output,
 463                            editable_range: input_excerpt.editable_range.to_offset(&snapshot),
 464                            input_outline,
 465                        })
 466                    }
 467                })
 468                .await?;
 469
 470            log::debug!(
 471                "Events:\n{}\nExcerpt:\n{:?}",
 472                values.input_events,
 473                values.input_excerpt
 474            );
 475
 476            let body = PredictEditsBody {
 477                input_events: values.input_events.clone(),
 478                input_excerpt: values.input_excerpt.clone(),
 479                speculated_output: Some(values.speculated_output),
 480                outline: Some(values.input_outline.clone()),
 481                can_collect_data,
 482                diagnostic_groups: diagnostic_groups.and_then(|diagnostic_groups| {
 483                    diagnostic_groups
 484                        .into_iter()
 485                        .map(|(name, diagnostic_group)| {
 486                            Ok((name.to_string(), serde_json::to_value(diagnostic_group)?))
 487                        })
 488                        .collect::<Result<Vec<_>>>()
 489                        .log_err()
 490                }),
 491            };
 492
 493            let response = perform_predict_edits(PerformPredictEditsParams {
 494                client,
 495                llm_token,
 496                app_version,
 497                body,
 498            })
 499            .await;
 500            let (response, usage) = match response {
 501                Ok(response) => response,
 502                Err(err) => {
 503                    if err.is::<ZedUpdateRequiredError>() {
 504                        cx.update(|cx| {
 505                            zeta.update(cx, |zeta, _cx| {
 506                                zeta.update_required = true;
 507                            });
 508
 509                            if let Some(workspace) = workspace {
 510                                workspace.update(cx, |workspace, cx| {
 511                                    workspace.show_notification(
 512                                        NotificationId::unique::<ZedUpdateRequiredError>(),
 513                                        cx,
 514                                        |cx| {
 515                                            cx.new(|cx| {
 516                                                ErrorMessagePrompt::new(err.to_string(), cx)
 517                                                    .with_link_button(
 518                                                        "Update Zed",
 519                                                        "https://zed.dev/releases",
 520                                                    )
 521                                            })
 522                                        },
 523                                    );
 524                                });
 525                            }
 526                        })
 527                        .ok();
 528                    }
 529
 530                    return Err(err);
 531                }
 532            };
 533
 534            log::debug!("completion response: {}", &response.output_excerpt);
 535
 536            if let Some(usage) = usage {
 537                this.update(cx, |this, cx| {
 538                    this.cloud_user_store.update(cx, |cloud_user_store, cx| {
 539                        cloud_user_store.update_edit_prediction_usage(usage, cx);
 540                    });
 541                })
 542                .ok();
 543            }
 544
 545            Self::process_completion_response(
 546                response,
 547                buffer,
 548                &snapshot,
 549                values.editable_range,
 550                cursor_offset,
 551                path,
 552                values.input_outline,
 553                values.input_events,
 554                values.input_excerpt,
 555                request_sent_at,
 556                &cx,
 557            )
 558            .await
 559        })
 560    }
 561
 562    // Generates several example completions of various states to fill the Zeta completion modal
 563    #[cfg(any(test, feature = "test-support"))]
 564    pub fn fill_with_fake_completions(&mut self, cx: &mut Context<Self>) -> Task<()> {
 565        use language::Point;
 566
 567        let test_buffer_text = indoc::indoc! {r#"a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 568            And maybe a short line
 569
 570            Then a few lines
 571
 572            and then another
 573            "#};
 574
 575        let project = None;
 576        let buffer = cx.new(|cx| Buffer::local(test_buffer_text, cx));
 577        let position = buffer.read(cx).anchor_before(Point::new(1, 0));
 578
 579        let completion_tasks = vec![
 580            self.fake_completion(
 581                project,
 582                &buffer,
 583                position,
 584                PredictEditsResponse {
 585                    request_id: Uuid::parse_str("e7861db5-0cea-4761-b1c5-ad083ac53a80").unwrap(),
 586                    output_excerpt: format!("{EDITABLE_REGION_START_MARKER}
 587a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 588[here's an edit]
 589And maybe a short line
 590Then a few lines
 591and then another
 592{EDITABLE_REGION_END_MARKER}
 593                        ", ),
 594                },
 595                cx,
 596            ),
 597            self.fake_completion(
 598                project,
 599                &buffer,
 600                position,
 601                PredictEditsResponse {
 602                    request_id: Uuid::parse_str("077c556a-2c49-44e2-bbc6-dafc09032a5e").unwrap(),
 603                    output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER}
 604a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 605And maybe a short line
 606[and another edit]
 607Then a few lines
 608and then another
 609{EDITABLE_REGION_END_MARKER}
 610                        "#),
 611                },
 612                cx,
 613            ),
 614            self.fake_completion(
 615                project,
 616                &buffer,
 617                position,
 618                PredictEditsResponse {
 619                    request_id: Uuid::parse_str("df8c7b23-3d1d-4f99-a306-1f6264a41277").unwrap(),
 620                    output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER}
 621a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 622And maybe a short line
 623
 624Then a few lines
 625
 626and then another
 627{EDITABLE_REGION_END_MARKER}
 628                        "#),
 629                },
 630                cx,
 631            ),
 632            self.fake_completion(
 633                project,
 634                &buffer,
 635                position,
 636                PredictEditsResponse {
 637                    request_id: Uuid::parse_str("c743958d-e4d8-44a8-aa5b-eb1e305c5f5c").unwrap(),
 638                    output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER}
 639a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 640And maybe a short line
 641
 642Then a few lines
 643
 644and then another
 645{EDITABLE_REGION_END_MARKER}
 646                        "#),
 647                },
 648                cx,
 649            ),
 650            self.fake_completion(
 651                project,
 652                &buffer,
 653                position,
 654                PredictEditsResponse {
 655                    request_id: Uuid::parse_str("ff5cd7ab-ad06-4808-986e-d3391e7b8355").unwrap(),
 656                    output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER}
 657a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 658And maybe a short line
 659Then a few lines
 660[a third completion]
 661and then another
 662{EDITABLE_REGION_END_MARKER}
 663                        "#),
 664                },
 665                cx,
 666            ),
 667            self.fake_completion(
 668                project,
 669                &buffer,
 670                position,
 671                PredictEditsResponse {
 672                    request_id: Uuid::parse_str("83cafa55-cdba-4b27-8474-1865ea06be94").unwrap(),
 673                    output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER}
 674a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 675And maybe a short line
 676and then another
 677[fourth completion example]
 678{EDITABLE_REGION_END_MARKER}
 679                        "#),
 680                },
 681                cx,
 682            ),
 683            self.fake_completion(
 684                project,
 685                &buffer,
 686                position,
 687                PredictEditsResponse {
 688                    request_id: Uuid::parse_str("d5bd3afd-8723-47c7-bd77-15a3a926867b").unwrap(),
 689                    output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER}
 690a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line
 691And maybe a short line
 692Then a few lines
 693and then another
 694[fifth and final completion]
 695{EDITABLE_REGION_END_MARKER}
 696                        "#),
 697                },
 698                cx,
 699            ),
 700        ];
 701
 702        cx.spawn(async move |zeta, cx| {
 703            for task in completion_tasks {
 704                task.await.unwrap();
 705            }
 706
 707            zeta.update(cx, |zeta, _cx| {
 708                zeta.shown_completions.get_mut(2).unwrap().edits = Arc::new([]);
 709                zeta.shown_completions.get_mut(3).unwrap().edits = Arc::new([]);
 710            })
 711            .ok();
 712        })
 713    }
 714
 715    #[cfg(any(test, feature = "test-support"))]
 716    pub fn fake_completion(
 717        &mut self,
 718        project: Option<&Entity<Project>>,
 719        buffer: &Entity<Buffer>,
 720        position: language::Anchor,
 721        response: PredictEditsResponse,
 722        cx: &mut Context<Self>,
 723    ) -> Task<Result<Option<InlineCompletion>>> {
 724        use std::future::ready;
 725
 726        self.request_completion_impl(None, project, buffer, position, false, cx, |_params| {
 727            ready(Ok((response, None)))
 728        })
 729    }
 730
 731    pub fn request_completion(
 732        &mut self,
 733        project: Option<&Entity<Project>>,
 734        buffer: &Entity<Buffer>,
 735        position: language::Anchor,
 736        can_collect_data: bool,
 737        cx: &mut Context<Self>,
 738    ) -> Task<Result<Option<InlineCompletion>>> {
 739        let workspace = self
 740            .workspace
 741            .as_ref()
 742            .and_then(|workspace| workspace.upgrade());
 743        self.request_completion_impl(
 744            workspace,
 745            project,
 746            buffer,
 747            position,
 748            can_collect_data,
 749            cx,
 750            Self::perform_predict_edits,
 751        )
 752    }
 753
 754    fn perform_predict_edits(
 755        params: PerformPredictEditsParams,
 756    ) -> impl Future<Output = Result<(PredictEditsResponse, Option<EditPredictionUsage>)>> {
 757        async move {
 758            let PerformPredictEditsParams {
 759                client,
 760                llm_token,
 761                app_version,
 762                body,
 763                ..
 764            } = params;
 765
 766            let http_client = client.http_client();
 767            let mut token = llm_token.acquire(&client).await?;
 768            let mut did_retry = false;
 769
 770            loop {
 771                let request_builder = http_client::Request::builder().method(Method::POST);
 772                let request_builder =
 773                    if let Ok(predict_edits_url) = std::env::var("ZED_PREDICT_EDITS_URL") {
 774                        request_builder.uri(predict_edits_url)
 775                    } else {
 776                        request_builder.uri(
 777                            http_client
 778                                .build_zed_llm_url("/predict_edits/v2", &[])?
 779                                .as_ref(),
 780                        )
 781                    };
 782                let request = request_builder
 783                    .header("Content-Type", "application/json")
 784                    .header("Authorization", format!("Bearer {}", token))
 785                    .header(ZED_VERSION_HEADER_NAME, app_version.to_string())
 786                    .body(serde_json::to_string(&body)?.into())?;
 787
 788                let mut response = http_client.send(request).await?;
 789
 790                if let Some(minimum_required_version) = response
 791                    .headers()
 792                    .get(MINIMUM_REQUIRED_VERSION_HEADER_NAME)
 793                    .and_then(|version| SemanticVersion::from_str(version.to_str().ok()?).ok())
 794                {
 795                    anyhow::ensure!(
 796                        app_version >= minimum_required_version,
 797                        ZedUpdateRequiredError {
 798                            minimum_version: minimum_required_version
 799                        }
 800                    );
 801                }
 802
 803                if response.status().is_success() {
 804                    let usage = EditPredictionUsage::from_headers(response.headers()).ok();
 805
 806                    let mut body = String::new();
 807                    response.body_mut().read_to_string(&mut body).await?;
 808                    return Ok((serde_json::from_str(&body)?, usage));
 809                } else if !did_retry
 810                    && response
 811                        .headers()
 812                        .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
 813                        .is_some()
 814                {
 815                    did_retry = true;
 816                    token = llm_token.refresh(&client).await?;
 817                } else {
 818                    let mut body = String::new();
 819                    response.body_mut().read_to_string(&mut body).await?;
 820                    anyhow::bail!(
 821                        "error predicting edits.\nStatus: {:?}\nBody: {}",
 822                        response.status(),
 823                        body
 824                    );
 825                }
 826            }
 827        }
 828    }
 829
 830    fn accept_edit_prediction(
 831        &mut self,
 832        request_id: InlineCompletionId,
 833        cx: &mut Context<Self>,
 834    ) -> Task<Result<()>> {
 835        let client = self.client.clone();
 836        let llm_token = self.llm_token.clone();
 837        let app_version = AppVersion::global(cx);
 838        cx.spawn(async move |this, cx| {
 839            let http_client = client.http_client();
 840            let mut response = llm_token_retry(&llm_token, &client, |token| {
 841                let request_builder = http_client::Request::builder().method(Method::POST);
 842                let request_builder =
 843                    if let Ok(accept_prediction_url) = std::env::var("ZED_ACCEPT_PREDICTION_URL") {
 844                        request_builder.uri(accept_prediction_url)
 845                    } else {
 846                        request_builder.uri(
 847                            http_client
 848                                .build_zed_llm_url("/predict_edits/accept", &[])?
 849                                .as_ref(),
 850                        )
 851                    };
 852                Ok(request_builder
 853                    .header("Content-Type", "application/json")
 854                    .header("Authorization", format!("Bearer {}", token))
 855                    .header(ZED_VERSION_HEADER_NAME, app_version.to_string())
 856                    .body(
 857                        serde_json::to_string(&AcceptEditPredictionBody {
 858                            request_id: request_id.0,
 859                        })?
 860                        .into(),
 861                    )?)
 862            })
 863            .await?;
 864
 865            if let Some(minimum_required_version) = response
 866                .headers()
 867                .get(MINIMUM_REQUIRED_VERSION_HEADER_NAME)
 868                .and_then(|version| SemanticVersion::from_str(version.to_str().ok()?).ok())
 869            {
 870                if app_version < minimum_required_version {
 871                    return Err(anyhow!(ZedUpdateRequiredError {
 872                        minimum_version: minimum_required_version
 873                    }));
 874                }
 875            }
 876
 877            if response.status().is_success() {
 878                if let Some(usage) = EditPredictionUsage::from_headers(response.headers()).ok() {
 879                    this.update(cx, |this, cx| {
 880                        this.cloud_user_store.update(cx, |cloud_user_store, cx| {
 881                            cloud_user_store.update_edit_prediction_usage(usage, cx);
 882                        });
 883                    })?;
 884                }
 885
 886                Ok(())
 887            } else {
 888                let mut body = String::new();
 889                response.body_mut().read_to_string(&mut body).await?;
 890                Err(anyhow!(
 891                    "error accepting edit prediction.\nStatus: {:?}\nBody: {}",
 892                    response.status(),
 893                    body
 894                ))
 895            }
 896        })
 897    }
 898
 899    fn process_completion_response(
 900        prediction_response: PredictEditsResponse,
 901        buffer: Entity<Buffer>,
 902        snapshot: &BufferSnapshot,
 903        editable_range: Range<usize>,
 904        cursor_offset: usize,
 905        path: Arc<Path>,
 906        input_outline: String,
 907        input_events: String,
 908        input_excerpt: String,
 909        request_sent_at: Instant,
 910        cx: &AsyncApp,
 911    ) -> Task<Result<Option<InlineCompletion>>> {
 912        let snapshot = snapshot.clone();
 913        let request_id = prediction_response.request_id;
 914        let output_excerpt = prediction_response.output_excerpt;
 915        cx.spawn(async move |cx| {
 916            let output_excerpt: Arc<str> = output_excerpt.into();
 917
 918            let edits: Arc<[(Range<Anchor>, String)]> = cx
 919                .background_spawn({
 920                    let output_excerpt = output_excerpt.clone();
 921                    let editable_range = editable_range.clone();
 922                    let snapshot = snapshot.clone();
 923                    async move { Self::parse_edits(output_excerpt, editable_range, &snapshot) }
 924                })
 925                .await?
 926                .into();
 927
 928            let Some((edits, snapshot, edit_preview)) = buffer.read_with(cx, {
 929                let edits = edits.clone();
 930                |buffer, cx| {
 931                    let new_snapshot = buffer.snapshot();
 932                    let edits: Arc<[(Range<Anchor>, String)]> =
 933                        interpolate(&snapshot, &new_snapshot, edits)?.into();
 934                    Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
 935                }
 936            })?
 937            else {
 938                return anyhow::Ok(None);
 939            };
 940
 941            let edit_preview = edit_preview.await;
 942
 943            Ok(Some(InlineCompletion {
 944                id: InlineCompletionId(request_id),
 945                path,
 946                excerpt_range: editable_range,
 947                cursor_offset,
 948                edits,
 949                edit_preview,
 950                snapshot,
 951                input_outline: input_outline.into(),
 952                input_events: input_events.into(),
 953                input_excerpt: input_excerpt.into(),
 954                output_excerpt,
 955                request_sent_at,
 956                response_received_at: Instant::now(),
 957            }))
 958        })
 959    }
 960
 961    fn parse_edits(
 962        output_excerpt: Arc<str>,
 963        editable_range: Range<usize>,
 964        snapshot: &BufferSnapshot,
 965    ) -> Result<Vec<(Range<Anchor>, String)>> {
 966        let content = output_excerpt.replace(CURSOR_MARKER, "");
 967
 968        let start_markers = content
 969            .match_indices(EDITABLE_REGION_START_MARKER)
 970            .collect::<Vec<_>>();
 971        anyhow::ensure!(
 972            start_markers.len() == 1,
 973            "expected exactly one start marker, found {}",
 974            start_markers.len()
 975        );
 976
 977        let end_markers = content
 978            .match_indices(EDITABLE_REGION_END_MARKER)
 979            .collect::<Vec<_>>();
 980        anyhow::ensure!(
 981            end_markers.len() == 1,
 982            "expected exactly one end marker, found {}",
 983            end_markers.len()
 984        );
 985
 986        let sof_markers = content
 987            .match_indices(START_OF_FILE_MARKER)
 988            .collect::<Vec<_>>();
 989        anyhow::ensure!(
 990            sof_markers.len() <= 1,
 991            "expected at most one start-of-file marker, found {}",
 992            sof_markers.len()
 993        );
 994
 995        let codefence_start = start_markers[0].0;
 996        let content = &content[codefence_start..];
 997
 998        let newline_ix = content.find('\n').context("could not find newline")?;
 999        let content = &content[newline_ix + 1..];
1000
1001        let codefence_end = content
1002            .rfind(&format!("\n{EDITABLE_REGION_END_MARKER}"))
1003            .context("could not find end marker")?;
1004        let new_text = &content[..codefence_end];
1005
1006        let old_text = snapshot
1007            .text_for_range(editable_range.clone())
1008            .collect::<String>();
1009
1010        Ok(Self::compute_edits(
1011            old_text,
1012            new_text,
1013            editable_range.start,
1014            &snapshot,
1015        ))
1016    }
1017
1018    pub fn compute_edits(
1019        old_text: String,
1020        new_text: &str,
1021        offset: usize,
1022        snapshot: &BufferSnapshot,
1023    ) -> Vec<(Range<Anchor>, String)> {
1024        text_diff(&old_text, &new_text)
1025            .into_iter()
1026            .map(|(mut old_range, new_text)| {
1027                old_range.start += offset;
1028                old_range.end += offset;
1029
1030                let prefix_len = common_prefix(
1031                    snapshot.chars_for_range(old_range.clone()),
1032                    new_text.chars(),
1033                );
1034                old_range.start += prefix_len;
1035
1036                let suffix_len = common_prefix(
1037                    snapshot.reversed_chars_for_range(old_range.clone()),
1038                    new_text[prefix_len..].chars().rev(),
1039                );
1040                old_range.end = old_range.end.saturating_sub(suffix_len);
1041
1042                let new_text = new_text[prefix_len..new_text.len() - suffix_len].to_string();
1043                let range = if old_range.is_empty() {
1044                    let anchor = snapshot.anchor_after(old_range.start);
1045                    anchor..anchor
1046                } else {
1047                    snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
1048                };
1049                (range, new_text)
1050            })
1051            .collect()
1052    }
1053
1054    pub fn is_completion_rated(&self, completion_id: InlineCompletionId) -> bool {
1055        self.rated_completions.contains(&completion_id)
1056    }
1057
1058    pub fn completion_shown(&mut self, completion: &InlineCompletion, cx: &mut Context<Self>) {
1059        self.shown_completions.push_front(completion.clone());
1060        if self.shown_completions.len() > 50 {
1061            let completion = self.shown_completions.pop_back().unwrap();
1062            self.rated_completions.remove(&completion.id);
1063        }
1064        cx.notify();
1065    }
1066
1067    pub fn rate_completion(
1068        &mut self,
1069        completion: &InlineCompletion,
1070        rating: InlineCompletionRating,
1071        feedback: String,
1072        cx: &mut Context<Self>,
1073    ) {
1074        self.rated_completions.insert(completion.id);
1075        telemetry::event!(
1076            "Edit Prediction Rated",
1077            rating,
1078            input_events = completion.input_events,
1079            input_excerpt = completion.input_excerpt,
1080            input_outline = completion.input_outline,
1081            output_excerpt = completion.output_excerpt,
1082            feedback
1083        );
1084        self.client.telemetry().flush_events().detach();
1085        cx.notify();
1086    }
1087
1088    pub fn shown_completions(&self) -> impl DoubleEndedIterator<Item = &InlineCompletion> {
1089        self.shown_completions.iter()
1090    }
1091
1092    pub fn shown_completions_len(&self) -> usize {
1093        self.shown_completions.len()
1094    }
1095
1096    fn report_changes_for_buffer(
1097        &mut self,
1098        buffer: &Entity<Buffer>,
1099        cx: &mut Context<Self>,
1100    ) -> BufferSnapshot {
1101        self.register_buffer(buffer, cx);
1102
1103        let registered_buffer = self
1104            .registered_buffers
1105            .get_mut(&buffer.entity_id())
1106            .unwrap();
1107        let new_snapshot = buffer.read(cx).snapshot();
1108
1109        if new_snapshot.version != registered_buffer.snapshot.version {
1110            let old_snapshot = mem::replace(&mut registered_buffer.snapshot, new_snapshot.clone());
1111            self.push_event(Event::BufferChange {
1112                old_snapshot,
1113                new_snapshot: new_snapshot.clone(),
1114                timestamp: Instant::now(),
1115            });
1116        }
1117
1118        new_snapshot
1119    }
1120
1121    fn load_data_collection_choices() -> DataCollectionChoice {
1122        let choice = KEY_VALUE_STORE
1123            .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
1124            .log_err()
1125            .flatten();
1126
1127        match choice.as_deref() {
1128            Some("true") => DataCollectionChoice::Enabled,
1129            Some("false") => DataCollectionChoice::Disabled,
1130            Some(_) => {
1131                log::error!("unknown value in '{ZED_PREDICT_DATA_COLLECTION_CHOICE}'");
1132                DataCollectionChoice::NotAnswered
1133            }
1134            None => DataCollectionChoice::NotAnswered,
1135        }
1136    }
1137}
1138
1139struct PerformPredictEditsParams {
1140    pub client: Arc<Client>,
1141    pub llm_token: LlmApiToken,
1142    pub app_version: SemanticVersion,
1143    pub body: PredictEditsBody,
1144}
1145
1146#[derive(Error, Debug)]
1147#[error(
1148    "You must update to Zed version {minimum_version} or higher to continue using edit predictions."
1149)]
1150pub struct ZedUpdateRequiredError {
1151    minimum_version: SemanticVersion,
1152}
1153
1154struct LicenseDetectionWatcher {
1155    is_open_source_rx: watch::Receiver<bool>,
1156    _is_open_source_task: Task<()>,
1157}
1158
1159impl LicenseDetectionWatcher {
1160    pub fn new(worktree: &Worktree, cx: &mut Context<Worktree>) -> Self {
1161        let (mut is_open_source_tx, is_open_source_rx) = watch::channel_with::<bool>(false);
1162
1163        // Check if worktree is a single file, if so we do not need to check for a LICENSE file
1164        let task = if worktree.abs_path().is_file() {
1165            Task::ready(())
1166        } else {
1167            let loaded_files = LICENSE_FILES_TO_CHECK
1168                .iter()
1169                .map(Path::new)
1170                .map(|file| worktree.load_file(file, cx))
1171                .collect::<ArrayVec<_, { LICENSE_FILES_TO_CHECK.len() }>>();
1172
1173            cx.background_spawn(async move {
1174                for loaded_file in loaded_files.into_iter() {
1175                    let Ok(loaded_file) = loaded_file.await else {
1176                        continue;
1177                    };
1178
1179                    let path = &loaded_file.file.path;
1180                    if is_license_eligible_for_data_collection(&loaded_file.text) {
1181                        log::info!("detected '{path:?}' as open source license");
1182                        *is_open_source_tx.borrow_mut() = true;
1183                    } else {
1184                        log::info!("didn't detect '{path:?}' as open source license");
1185                    }
1186
1187                    // stop on the first license that successfully read
1188                    return;
1189                }
1190
1191                log::debug!("didn't find a license file to check, assuming closed source");
1192            })
1193        };
1194
1195        Self {
1196            is_open_source_rx,
1197            _is_open_source_task: task,
1198        }
1199    }
1200
1201    /// Answers false until we find out it's open source
1202    pub fn is_project_open_source(&self) -> bool {
1203        *self.is_open_source_rx.borrow()
1204    }
1205}
1206
1207fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
1208    a.zip(b)
1209        .take_while(|(a, b)| a == b)
1210        .map(|(a, _)| a.len_utf8())
1211        .sum()
1212}
1213
1214fn prompt_for_outline(snapshot: &BufferSnapshot) -> String {
1215    let mut input_outline = String::new();
1216
1217    writeln!(
1218        input_outline,
1219        "```{}",
1220        snapshot
1221            .file()
1222            .map_or(Cow::Borrowed("untitled"), |file| file
1223                .path()
1224                .to_string_lossy())
1225    )
1226    .unwrap();
1227
1228    if let Some(outline) = snapshot.outline(None) {
1229        for item in &outline.items {
1230            let spacing = " ".repeat(item.depth);
1231            writeln!(input_outline, "{}{}", spacing, item.text).unwrap();
1232        }
1233    }
1234
1235    writeln!(input_outline, "```").unwrap();
1236
1237    input_outline
1238}
1239
1240fn prompt_for_events(events: &VecDeque<Event>, mut remaining_tokens: usize) -> String {
1241    let mut result = String::new();
1242    for event in events.iter().rev() {
1243        let event_string = event.to_prompt();
1244        let event_tokens = tokens_for_bytes(event_string.len());
1245        if event_tokens > remaining_tokens {
1246            break;
1247        }
1248
1249        if !result.is_empty() {
1250            result.insert_str(0, "\n\n");
1251        }
1252        result.insert_str(0, &event_string);
1253        remaining_tokens -= event_tokens;
1254    }
1255    result
1256}
1257
1258struct RegisteredBuffer {
1259    snapshot: BufferSnapshot,
1260    _subscriptions: [gpui::Subscription; 2],
1261}
1262
1263#[derive(Clone)]
1264enum Event {
1265    BufferChange {
1266        old_snapshot: BufferSnapshot,
1267        new_snapshot: BufferSnapshot,
1268        timestamp: Instant,
1269    },
1270}
1271
1272impl Event {
1273    fn to_prompt(&self) -> String {
1274        match self {
1275            Event::BufferChange {
1276                old_snapshot,
1277                new_snapshot,
1278                ..
1279            } => {
1280                let mut prompt = String::new();
1281
1282                let old_path = old_snapshot
1283                    .file()
1284                    .map(|f| f.path().as_ref())
1285                    .unwrap_or(Path::new("untitled"));
1286                let new_path = new_snapshot
1287                    .file()
1288                    .map(|f| f.path().as_ref())
1289                    .unwrap_or(Path::new("untitled"));
1290                if old_path != new_path {
1291                    writeln!(prompt, "User renamed {:?} to {:?}\n", old_path, new_path).unwrap();
1292                }
1293
1294                let diff = language::unified_diff(&old_snapshot.text(), &new_snapshot.text());
1295                if !diff.is_empty() {
1296                    write!(
1297                        prompt,
1298                        "User edited {:?}:\n```diff\n{}\n```",
1299                        new_path, diff
1300                    )
1301                    .unwrap();
1302                }
1303
1304                prompt
1305            }
1306        }
1307    }
1308}
1309
1310#[derive(Debug, Clone)]
1311struct CurrentInlineCompletion {
1312    buffer_id: EntityId,
1313    completion: InlineCompletion,
1314}
1315
1316impl CurrentInlineCompletion {
1317    fn should_replace_completion(&self, old_completion: &Self, snapshot: &BufferSnapshot) -> bool {
1318        if self.buffer_id != old_completion.buffer_id {
1319            return true;
1320        }
1321
1322        let Some(old_edits) = old_completion.completion.interpolate(&snapshot) else {
1323            return true;
1324        };
1325        let Some(new_edits) = self.completion.interpolate(&snapshot) else {
1326            return false;
1327        };
1328
1329        if old_edits.len() == 1 && new_edits.len() == 1 {
1330            let (old_range, old_text) = &old_edits[0];
1331            let (new_range, new_text) = &new_edits[0];
1332            new_range == old_range && new_text.starts_with(old_text)
1333        } else {
1334            true
1335        }
1336    }
1337}
1338
1339struct PendingCompletion {
1340    id: usize,
1341    _task: Task<()>,
1342}
1343
1344#[derive(Debug, Clone, Copy)]
1345pub enum DataCollectionChoice {
1346    NotAnswered,
1347    Enabled,
1348    Disabled,
1349}
1350
1351impl DataCollectionChoice {
1352    pub fn is_enabled(self) -> bool {
1353        match self {
1354            Self::Enabled => true,
1355            Self::NotAnswered | Self::Disabled => false,
1356        }
1357    }
1358
1359    pub fn is_answered(self) -> bool {
1360        match self {
1361            Self::Enabled | Self::Disabled => true,
1362            Self::NotAnswered => false,
1363        }
1364    }
1365
1366    pub fn toggle(&self) -> DataCollectionChoice {
1367        match self {
1368            Self::Enabled => Self::Disabled,
1369            Self::Disabled => Self::Enabled,
1370            Self::NotAnswered => Self::Enabled,
1371        }
1372    }
1373}
1374
1375impl From<bool> for DataCollectionChoice {
1376    fn from(value: bool) -> Self {
1377        match value {
1378            true => DataCollectionChoice::Enabled,
1379            false => DataCollectionChoice::Disabled,
1380        }
1381    }
1382}
1383
1384pub struct ProviderDataCollection {
1385    /// When set to None, data collection is not possible in the provider buffer
1386    choice: Option<Entity<DataCollectionChoice>>,
1387    license_detection_watcher: Option<Rc<LicenseDetectionWatcher>>,
1388}
1389
1390impl ProviderDataCollection {
1391    pub fn new(zeta: Entity<Zeta>, buffer: Option<Entity<Buffer>>, cx: &mut App) -> Self {
1392        let choice_and_watcher = buffer.and_then(|buffer| {
1393            let file = buffer.read(cx).file()?;
1394
1395            if !file.is_local() || file.is_private() {
1396                return None;
1397            }
1398
1399            let zeta = zeta.read(cx);
1400            let choice = zeta.data_collection_choice.clone();
1401
1402            let license_detection_watcher = zeta
1403                .license_detection_watchers
1404                .get(&file.worktree_id(cx))
1405                .cloned()?;
1406
1407            Some((choice, license_detection_watcher))
1408        });
1409
1410        if let Some((choice, watcher)) = choice_and_watcher {
1411            ProviderDataCollection {
1412                choice: Some(choice),
1413                license_detection_watcher: Some(watcher),
1414            }
1415        } else {
1416            ProviderDataCollection {
1417                choice: None,
1418                license_detection_watcher: None,
1419            }
1420        }
1421    }
1422
1423    pub fn can_collect_data(&self, cx: &App) -> bool {
1424        self.is_data_collection_enabled(cx) && self.is_project_open_source()
1425    }
1426
1427    pub fn is_data_collection_enabled(&self, cx: &App) -> bool {
1428        self.choice
1429            .as_ref()
1430            .is_some_and(|choice| choice.read(cx).is_enabled())
1431    }
1432
1433    fn is_project_open_source(&self) -> bool {
1434        self.license_detection_watcher
1435            .as_ref()
1436            .is_some_and(|watcher| watcher.is_project_open_source())
1437    }
1438
1439    pub fn toggle(&mut self, cx: &mut App) {
1440        if let Some(choice) = self.choice.as_mut() {
1441            let new_choice = choice.update(cx, |choice, _cx| {
1442                let new_choice = choice.toggle();
1443                *choice = new_choice;
1444                new_choice
1445            });
1446
1447            db::write_and_log(cx, move || {
1448                KEY_VALUE_STORE.write_kvp(
1449                    ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
1450                    new_choice.is_enabled().to_string(),
1451                )
1452            });
1453        }
1454    }
1455}
1456
1457async fn llm_token_retry(
1458    llm_token: &LlmApiToken,
1459    client: &Arc<Client>,
1460    build_request: impl Fn(String) -> Result<Request<AsyncBody>>,
1461) -> Result<Response<AsyncBody>> {
1462    let mut did_retry = false;
1463    let http_client = client.http_client();
1464    let mut token = llm_token.acquire(client).await?;
1465    loop {
1466        let request = build_request(token.clone())?;
1467        let response = http_client.send(request).await?;
1468
1469        if !did_retry
1470            && !response.status().is_success()
1471            && response
1472                .headers()
1473                .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
1474                .is_some()
1475        {
1476            did_retry = true;
1477            token = llm_token.refresh(client).await?;
1478            continue;
1479        }
1480
1481        return Ok(response);
1482    }
1483}
1484
1485pub struct ZetaInlineCompletionProvider {
1486    zeta: Entity<Zeta>,
1487    pending_completions: ArrayVec<PendingCompletion, 2>,
1488    next_pending_completion_id: usize,
1489    current_completion: Option<CurrentInlineCompletion>,
1490    /// None if this is entirely disabled for this provider
1491    provider_data_collection: ProviderDataCollection,
1492    last_request_timestamp: Instant,
1493}
1494
1495impl ZetaInlineCompletionProvider {
1496    pub const THROTTLE_TIMEOUT: Duration = Duration::from_millis(300);
1497
1498    pub fn new(zeta: Entity<Zeta>, provider_data_collection: ProviderDataCollection) -> Self {
1499        Self {
1500            zeta,
1501            pending_completions: ArrayVec::new(),
1502            next_pending_completion_id: 0,
1503            current_completion: None,
1504            provider_data_collection,
1505            last_request_timestamp: Instant::now(),
1506        }
1507    }
1508}
1509
1510impl inline_completion::EditPredictionProvider for ZetaInlineCompletionProvider {
1511    fn name() -> &'static str {
1512        "zed-predict"
1513    }
1514
1515    fn display_name() -> &'static str {
1516        "Zed's Edit Predictions"
1517    }
1518
1519    fn show_completions_in_menu() -> bool {
1520        true
1521    }
1522
1523    fn show_tab_accept_marker() -> bool {
1524        true
1525    }
1526
1527    fn data_collection_state(&self, cx: &App) -> DataCollectionState {
1528        let is_project_open_source = self.provider_data_collection.is_project_open_source();
1529
1530        if self.provider_data_collection.is_data_collection_enabled(cx) {
1531            DataCollectionState::Enabled {
1532                is_project_open_source,
1533            }
1534        } else {
1535            DataCollectionState::Disabled {
1536                is_project_open_source,
1537            }
1538        }
1539    }
1540
1541    fn toggle_data_collection(&mut self, cx: &mut App) {
1542        self.provider_data_collection.toggle(cx);
1543    }
1544
1545    fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
1546        self.zeta.read(cx).usage(cx)
1547    }
1548
1549    fn is_enabled(
1550        &self,
1551        _buffer: &Entity<Buffer>,
1552        _cursor_position: language::Anchor,
1553        _cx: &App,
1554    ) -> bool {
1555        true
1556    }
1557
1558    fn needs_terms_acceptance(&self, cx: &App) -> bool {
1559        !self
1560            .zeta
1561            .read(cx)
1562            .cloud_user_store
1563            .read(cx)
1564            .has_accepted_tos()
1565    }
1566
1567    fn is_refreshing(&self) -> bool {
1568        !self.pending_completions.is_empty()
1569    }
1570
1571    fn refresh(
1572        &mut self,
1573        project: Option<Entity<Project>>,
1574        buffer: Entity<Buffer>,
1575        position: language::Anchor,
1576        _debounce: bool,
1577        cx: &mut Context<Self>,
1578    ) {
1579        if self.needs_terms_acceptance(cx) {
1580            return;
1581        }
1582
1583        if self.zeta.read(cx).update_required {
1584            return;
1585        }
1586
1587        if self
1588            .zeta
1589            .read(cx)
1590            .cloud_user_store
1591            .read_with(cx, |cloud_user_store, _cx| {
1592                cloud_user_store.account_too_young() || cloud_user_store.has_overdue_invoices()
1593            })
1594        {
1595            return;
1596        }
1597
1598        if let Some(current_completion) = self.current_completion.as_ref() {
1599            let snapshot = buffer.read(cx).snapshot();
1600            if current_completion
1601                .completion
1602                .interpolate(&snapshot)
1603                .is_some()
1604            {
1605                return;
1606            }
1607        }
1608
1609        let pending_completion_id = self.next_pending_completion_id;
1610        self.next_pending_completion_id += 1;
1611        let can_collect_data = self.provider_data_collection.can_collect_data(cx);
1612        let last_request_timestamp = self.last_request_timestamp;
1613
1614        let task = cx.spawn(async move |this, cx| {
1615            if let Some(timeout) = (last_request_timestamp + Self::THROTTLE_TIMEOUT)
1616                .checked_duration_since(Instant::now())
1617            {
1618                cx.background_executor().timer(timeout).await;
1619            }
1620
1621            let completion_request = this.update(cx, |this, cx| {
1622                this.last_request_timestamp = Instant::now();
1623                this.zeta.update(cx, |zeta, cx| {
1624                    zeta.request_completion(
1625                        project.as_ref(),
1626                        &buffer,
1627                        position,
1628                        can_collect_data,
1629                        cx,
1630                    )
1631                })
1632            });
1633
1634            let completion = match completion_request {
1635                Ok(completion_request) => {
1636                    let completion_request = completion_request.await;
1637                    completion_request.map(|c| {
1638                        c.map(|completion| CurrentInlineCompletion {
1639                            buffer_id: buffer.entity_id(),
1640                            completion,
1641                        })
1642                    })
1643                }
1644                Err(error) => Err(error),
1645            };
1646            let Some(new_completion) = completion
1647                .context("edit prediction failed")
1648                .log_err()
1649                .flatten()
1650            else {
1651                this.update(cx, |this, cx| {
1652                    if this.pending_completions[0].id == pending_completion_id {
1653                        this.pending_completions.remove(0);
1654                    } else {
1655                        this.pending_completions.clear();
1656                    }
1657
1658                    cx.notify();
1659                })
1660                .ok();
1661                return;
1662            };
1663
1664            this.update(cx, |this, cx| {
1665                if this.pending_completions[0].id == pending_completion_id {
1666                    this.pending_completions.remove(0);
1667                } else {
1668                    this.pending_completions.clear();
1669                }
1670
1671                if let Some(old_completion) = this.current_completion.as_ref() {
1672                    let snapshot = buffer.read(cx).snapshot();
1673                    if new_completion.should_replace_completion(&old_completion, &snapshot) {
1674                        this.zeta.update(cx, |zeta, cx| {
1675                            zeta.completion_shown(&new_completion.completion, cx);
1676                        });
1677                        this.current_completion = Some(new_completion);
1678                    }
1679                } else {
1680                    this.zeta.update(cx, |zeta, cx| {
1681                        zeta.completion_shown(&new_completion.completion, cx);
1682                    });
1683                    this.current_completion = Some(new_completion);
1684                }
1685
1686                cx.notify();
1687            })
1688            .ok();
1689        });
1690
1691        // We always maintain at most two pending completions. When we already
1692        // have two, we replace the newest one.
1693        if self.pending_completions.len() <= 1 {
1694            self.pending_completions.push(PendingCompletion {
1695                id: pending_completion_id,
1696                _task: task,
1697            });
1698        } else if self.pending_completions.len() == 2 {
1699            self.pending_completions.pop();
1700            self.pending_completions.push(PendingCompletion {
1701                id: pending_completion_id,
1702                _task: task,
1703            });
1704        }
1705    }
1706
1707    fn cycle(
1708        &mut self,
1709        _buffer: Entity<Buffer>,
1710        _cursor_position: language::Anchor,
1711        _direction: inline_completion::Direction,
1712        _cx: &mut Context<Self>,
1713    ) {
1714        // Right now we don't support cycling.
1715    }
1716
1717    fn accept(&mut self, cx: &mut Context<Self>) {
1718        let completion_id = self
1719            .current_completion
1720            .as_ref()
1721            .map(|completion| completion.completion.id);
1722        if let Some(completion_id) = completion_id {
1723            self.zeta
1724                .update(cx, |zeta, cx| {
1725                    zeta.accept_edit_prediction(completion_id, cx)
1726                })
1727                .detach();
1728        }
1729        self.pending_completions.clear();
1730    }
1731
1732    fn discard(&mut self, _cx: &mut Context<Self>) {
1733        self.pending_completions.clear();
1734        self.current_completion.take();
1735    }
1736
1737    fn suggest(
1738        &mut self,
1739        buffer: &Entity<Buffer>,
1740        cursor_position: language::Anchor,
1741        cx: &mut Context<Self>,
1742    ) -> Option<inline_completion::InlineCompletion> {
1743        let CurrentInlineCompletion {
1744            buffer_id,
1745            completion,
1746            ..
1747        } = self.current_completion.as_mut()?;
1748
1749        // Invalidate previous completion if it was generated for a different buffer.
1750        if *buffer_id != buffer.entity_id() {
1751            self.current_completion.take();
1752            return None;
1753        }
1754
1755        let buffer = buffer.read(cx);
1756        let Some(edits) = completion.interpolate(&buffer.snapshot()) else {
1757            self.current_completion.take();
1758            return None;
1759        };
1760
1761        let cursor_row = cursor_position.to_point(buffer).row;
1762        let (closest_edit_ix, (closest_edit_range, _)) =
1763            edits.iter().enumerate().min_by_key(|(_, (range, _))| {
1764                let distance_from_start = cursor_row.abs_diff(range.start.to_point(buffer).row);
1765                let distance_from_end = cursor_row.abs_diff(range.end.to_point(buffer).row);
1766                cmp::min(distance_from_start, distance_from_end)
1767            })?;
1768
1769        let mut edit_start_ix = closest_edit_ix;
1770        for (range, _) in edits[..edit_start_ix].iter().rev() {
1771            let distance_from_closest_edit =
1772                closest_edit_range.start.to_point(buffer).row - range.end.to_point(buffer).row;
1773            if distance_from_closest_edit <= 1 {
1774                edit_start_ix -= 1;
1775            } else {
1776                break;
1777            }
1778        }
1779
1780        let mut edit_end_ix = closest_edit_ix + 1;
1781        for (range, _) in &edits[edit_end_ix..] {
1782            let distance_from_closest_edit =
1783                range.start.to_point(buffer).row - closest_edit_range.end.to_point(buffer).row;
1784            if distance_from_closest_edit <= 1 {
1785                edit_end_ix += 1;
1786            } else {
1787                break;
1788            }
1789        }
1790
1791        Some(inline_completion::InlineCompletion {
1792            id: Some(completion.id.to_string().into()),
1793            edits: edits[edit_start_ix..edit_end_ix].to_vec(),
1794            edit_preview: Some(completion.edit_preview.clone()),
1795        })
1796    }
1797}
1798
1799fn tokens_for_bytes(bytes: usize) -> usize {
1800    /// Typical number of string bytes per token for the purposes of limiting model input. This is
1801    /// intentionally low to err on the side of underestimating limits.
1802    const BYTES_PER_TOKEN_GUESS: usize = 3;
1803    bytes / BYTES_PER_TOKEN_GUESS
1804}
1805
1806#[cfg(test)]
1807mod tests {
1808    use client::UserStore;
1809    use client::test::FakeServer;
1810    use clock::FakeSystemClock;
1811    use cloud_api_types::{
1812        AuthenticatedUser, CreateLlmTokenResponse, GetAuthenticatedUserResponse, LlmToken, PlanInfo,
1813    };
1814    use cloud_llm_client::{CurrentUsage, Plan, UsageData, UsageLimit};
1815    use gpui::TestAppContext;
1816    use http_client::FakeHttpClient;
1817    use indoc::indoc;
1818    use language::Point;
1819    use settings::SettingsStore;
1820
1821    use super::*;
1822
1823    fn make_get_authenticated_user_response() -> GetAuthenticatedUserResponse {
1824        GetAuthenticatedUserResponse {
1825            user: AuthenticatedUser {
1826                id: 1,
1827                metrics_id: "metrics-id-1".to_string(),
1828                avatar_url: "".to_string(),
1829                github_login: "".to_string(),
1830                name: None,
1831                is_staff: false,
1832                accepted_tos_at: None,
1833            },
1834            feature_flags: vec![],
1835            plan: PlanInfo {
1836                plan: Plan::ZedPro,
1837                subscription_period: None,
1838                usage: CurrentUsage {
1839                    model_requests: UsageData {
1840                        used: 0,
1841                        limit: UsageLimit::Limited(500),
1842                    },
1843                    edit_predictions: UsageData {
1844                        used: 250,
1845                        limit: UsageLimit::Unlimited,
1846                    },
1847                },
1848                trial_started_at: None,
1849                is_usage_based_billing_enabled: false,
1850                is_account_too_young: false,
1851                has_overdue_invoices: false,
1852            },
1853        }
1854    }
1855
1856    #[gpui::test]
1857    async fn test_inline_completion_basic_interpolation(cx: &mut TestAppContext) {
1858        let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
1859        let edits: Arc<[(Range<Anchor>, String)]> = cx.update(|cx| {
1860            to_completion_edits(
1861                [(2..5, "REM".to_string()), (9..11, "".to_string())],
1862                &buffer,
1863                cx,
1864            )
1865            .into()
1866        });
1867
1868        let edit_preview = cx
1869            .read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
1870            .await;
1871
1872        let completion = InlineCompletion {
1873            edits,
1874            edit_preview,
1875            path: Path::new("").into(),
1876            snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
1877            id: InlineCompletionId(Uuid::new_v4()),
1878            excerpt_range: 0..0,
1879            cursor_offset: 0,
1880            input_outline: "".into(),
1881            input_events: "".into(),
1882            input_excerpt: "".into(),
1883            output_excerpt: "".into(),
1884            request_sent_at: Instant::now(),
1885            response_received_at: Instant::now(),
1886        };
1887
1888        cx.update(|cx| {
1889            assert_eq!(
1890                from_completion_edits(
1891                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1892                    &buffer,
1893                    cx
1894                ),
1895                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
1896            );
1897
1898            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
1899            assert_eq!(
1900                from_completion_edits(
1901                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1902                    &buffer,
1903                    cx
1904                ),
1905                vec![(2..2, "REM".to_string()), (6..8, "".to_string())]
1906            );
1907
1908            buffer.update(cx, |buffer, cx| buffer.undo(cx));
1909            assert_eq!(
1910                from_completion_edits(
1911                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1912                    &buffer,
1913                    cx
1914                ),
1915                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
1916            );
1917
1918            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
1919            assert_eq!(
1920                from_completion_edits(
1921                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1922                    &buffer,
1923                    cx
1924                ),
1925                vec![(3..3, "EM".to_string()), (7..9, "".to_string())]
1926            );
1927
1928            buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
1929            assert_eq!(
1930                from_completion_edits(
1931                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1932                    &buffer,
1933                    cx
1934                ),
1935                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
1936            );
1937
1938            buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
1939            assert_eq!(
1940                from_completion_edits(
1941                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1942                    &buffer,
1943                    cx
1944                ),
1945                vec![(9..11, "".to_string())]
1946            );
1947
1948            buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
1949            assert_eq!(
1950                from_completion_edits(
1951                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1952                    &buffer,
1953                    cx
1954                ),
1955                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
1956            );
1957
1958            buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
1959            assert_eq!(
1960                from_completion_edits(
1961                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1962                    &buffer,
1963                    cx
1964                ),
1965                vec![(4..4, "M".to_string())]
1966            );
1967
1968            buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
1969            assert_eq!(completion.interpolate(&buffer.read(cx).snapshot()), None);
1970        })
1971    }
1972
1973    #[gpui::test]
1974    async fn test_clean_up_diff(cx: &mut TestAppContext) {
1975        cx.update(|cx| {
1976            let settings_store = SettingsStore::test(cx);
1977            cx.set_global(settings_store);
1978            client::init_settings(cx);
1979        });
1980
1981        let edits = edits_for_prediction(
1982            indoc! {"
1983                fn main() {
1984                    let word_1 = \"lorem\";
1985                    let range = word.len()..word.len();
1986                }
1987            "},
1988            indoc! {"
1989                <|editable_region_start|>
1990                fn main() {
1991                    let word_1 = \"lorem\";
1992                    let range = word_1.len()..word_1.len();
1993                }
1994
1995                <|editable_region_end|>
1996            "},
1997            cx,
1998        )
1999        .await;
2000        assert_eq!(
2001            edits,
2002            [
2003                (Point::new(2, 20)..Point::new(2, 20), "_1".to_string()),
2004                (Point::new(2, 32)..Point::new(2, 32), "_1".to_string()),
2005            ]
2006        );
2007
2008        let edits = edits_for_prediction(
2009            indoc! {"
2010                fn main() {
2011                    let story = \"the quick\"
2012                }
2013            "},
2014            indoc! {"
2015                <|editable_region_start|>
2016                fn main() {
2017                    let story = \"the quick brown fox jumps over the lazy dog\";
2018                }
2019
2020                <|editable_region_end|>
2021            "},
2022            cx,
2023        )
2024        .await;
2025        assert_eq!(
2026            edits,
2027            [
2028                (
2029                    Point::new(1, 26)..Point::new(1, 26),
2030                    " brown fox jumps over the lazy dog".to_string()
2031                ),
2032                (Point::new(1, 27)..Point::new(1, 27), ";".to_string()),
2033            ]
2034        );
2035    }
2036
2037    #[gpui::test]
2038    async fn test_inline_completion_end_of_buffer(cx: &mut TestAppContext) {
2039        cx.update(|cx| {
2040            let settings_store = SettingsStore::test(cx);
2041            cx.set_global(settings_store);
2042            client::init_settings(cx);
2043        });
2044
2045        let buffer_content = "lorem\n";
2046        let completion_response = indoc! {"
2047            ```animals.js
2048            <|start_of_file|>
2049            <|editable_region_start|>
2050            lorem
2051            ipsum
2052            <|editable_region_end|>
2053            ```"};
2054
2055        let http_client = FakeHttpClient::create(move |req| async move {
2056            match (req.method(), req.uri().path()) {
2057                (&Method::GET, "/client/users/me") => Ok(http_client::Response::builder()
2058                    .status(200)
2059                    .body(
2060                        serde_json::to_string(&make_get_authenticated_user_response())
2061                            .unwrap()
2062                            .into(),
2063                    )
2064                    .unwrap()),
2065                (&Method::POST, "/client/llm_tokens") => Ok(http_client::Response::builder()
2066                    .status(200)
2067                    .body(
2068                        serde_json::to_string(&CreateLlmTokenResponse {
2069                            token: LlmToken("the-llm-token".to_string()),
2070                        })
2071                        .unwrap()
2072                        .into(),
2073                    )
2074                    .unwrap()),
2075                (&Method::POST, "/predict_edits/v2") => Ok(http_client::Response::builder()
2076                    .status(200)
2077                    .body(
2078                        serde_json::to_string(&PredictEditsResponse {
2079                            request_id: Uuid::parse_str("7e86480f-3536-4d2c-9334-8213e3445d45")
2080                                .unwrap(),
2081                            output_excerpt: completion_response.to_string(),
2082                        })
2083                        .unwrap()
2084                        .into(),
2085                    )
2086                    .unwrap()),
2087                _ => Ok(http_client::Response::builder()
2088                    .status(404)
2089                    .body("Not Found".into())
2090                    .unwrap()),
2091            }
2092        });
2093
2094        let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
2095        cx.update(|cx| {
2096            RefreshLlmTokenListener::register(client.clone(), cx);
2097        });
2098        // Construct the fake server to authenticate.
2099        let _server = FakeServer::for_client(42, &client, cx).await;
2100        let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
2101        let cloud_user_store =
2102            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
2103        let zeta = cx.new(|cx| Zeta::new(None, client, cloud_user_store, cx));
2104
2105        let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
2106        let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
2107        let completion_task = zeta.update(cx, |zeta, cx| {
2108            zeta.request_completion(None, &buffer, cursor, false, cx)
2109        });
2110
2111        let completion = completion_task.await.unwrap().unwrap();
2112        buffer.update(cx, |buffer, cx| {
2113            buffer.edit(completion.edits.iter().cloned(), None, cx)
2114        });
2115        assert_eq!(
2116            buffer.read_with(cx, |buffer, _| buffer.text()),
2117            "lorem\nipsum"
2118        );
2119    }
2120
2121    async fn edits_for_prediction(
2122        buffer_content: &str,
2123        completion_response: &str,
2124        cx: &mut TestAppContext,
2125    ) -> Vec<(Range<Point>, String)> {
2126        let completion_response = completion_response.to_string();
2127        let http_client = FakeHttpClient::create(move |req| {
2128            let completion = completion_response.clone();
2129            async move {
2130                match (req.method(), req.uri().path()) {
2131                    (&Method::GET, "/client/users/me") => Ok(http_client::Response::builder()
2132                        .status(200)
2133                        .body(
2134                            serde_json::to_string(&make_get_authenticated_user_response())
2135                                .unwrap()
2136                                .into(),
2137                        )
2138                        .unwrap()),
2139                    (&Method::POST, "/client/llm_tokens") => Ok(http_client::Response::builder()
2140                        .status(200)
2141                        .body(
2142                            serde_json::to_string(&CreateLlmTokenResponse {
2143                                token: LlmToken("the-llm-token".to_string()),
2144                            })
2145                            .unwrap()
2146                            .into(),
2147                        )
2148                        .unwrap()),
2149                    (&Method::POST, "/predict_edits/v2") => Ok(http_client::Response::builder()
2150                        .status(200)
2151                        .body(
2152                            serde_json::to_string(&PredictEditsResponse {
2153                                request_id: Uuid::new_v4(),
2154                                output_excerpt: completion,
2155                            })
2156                            .unwrap()
2157                            .into(),
2158                        )
2159                        .unwrap()),
2160                    _ => Ok(http_client::Response::builder()
2161                        .status(404)
2162                        .body("Not Found".into())
2163                        .unwrap()),
2164                }
2165            }
2166        });
2167
2168        let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
2169        cx.update(|cx| {
2170            RefreshLlmTokenListener::register(client.clone(), cx);
2171        });
2172        // Construct the fake server to authenticate.
2173        let _server = FakeServer::for_client(42, &client, cx).await;
2174        let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
2175        let cloud_user_store =
2176            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
2177        let zeta = cx.new(|cx| Zeta::new(None, client, cloud_user_store, cx));
2178
2179        let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
2180        let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
2181        let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
2182        let completion_task = zeta.update(cx, |zeta, cx| {
2183            zeta.request_completion(None, &buffer, cursor, false, cx)
2184        });
2185
2186        let completion = completion_task.await.unwrap().unwrap();
2187        completion
2188            .edits
2189            .into_iter()
2190            .map(|(old_range, new_text)| (old_range.to_point(&snapshot), new_text.clone()))
2191            .collect::<Vec<_>>()
2192    }
2193
2194    fn to_completion_edits(
2195        iterator: impl IntoIterator<Item = (Range<usize>, String)>,
2196        buffer: &Entity<Buffer>,
2197        cx: &App,
2198    ) -> Vec<(Range<Anchor>, String)> {
2199        let buffer = buffer.read(cx);
2200        iterator
2201            .into_iter()
2202            .map(|(range, text)| {
2203                (
2204                    buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
2205                    text,
2206                )
2207            })
2208            .collect()
2209    }
2210
2211    fn from_completion_edits(
2212        editor_edits: &[(Range<Anchor>, String)],
2213        buffer: &Entity<Buffer>,
2214        cx: &App,
2215    ) -> Vec<(Range<usize>, String)> {
2216        let buffer = buffer.read(cx);
2217        editor_edits
2218            .iter()
2219            .map(|(range, text)| {
2220                (
2221                    range.start.to_offset(buffer)..range.end.to_offset(buffer),
2222                    text.clone(),
2223                )
2224            })
2225            .collect()
2226    }
2227
2228    #[ctor::ctor]
2229    fn init_logger() {
2230        zlog::init_test();
2231    }
2232}