use anyhow::{Context as _, Result, anyhow};use buffer_diff::BufferDiff;use client::{Client, EditPredictionUsage, UserStore, global_llm_token};use cloud_api_client::LlmApiToken;use cloud_api_types::{    EditPredictionRecentFile, EditPredictionSettledKeptChars,    MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST, OrganizationId, SettledEditPrediction,    SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody,    SubmitEditPredictionSettledBatchBody, SubmitEditPredictionSettledResponse,};use cloud_llm_client::predict_edits_v3::{    PREDICT_EDITS_MODE_HEADER_NAME, PREDICT_EDITS_REQUEST_ID_HEADER_NAME,    PREDICT_EDITS_TRIGGER_HEADER_NAME, PredictEditsMode, PredictEditsV3Request,    PredictEditsV3Response, RawCompletionRequest, RawCompletionResponse,};use cloud_llm_client::predict_edits_v4::{PredictEditsV4Request, PredictEditsV4Response};use cloud_llm_client::{    EditPredictionRejectReason, EditPredictionRejection,    MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST, MINIMUM_REQUIRED_VERSION_HEADER_NAME,    PREFERRED_EXPERIMENT_HEADER_NAME, PredictEditsRequestTrigger, RejectEditPredictionsBodyRef,    ZED_VERSION_HEADER_NAME,};use collections::{HashMap, HashSet};use copilot::{Copilot, Reinstall, SignIn, SignOut};use credentials_provider::CredentialsProvider;use db::kvp::{Dismissable, KeyValueStore};use edit_prediction_context::{RelatedExcerptStore, RelatedExcerptStoreEvent, RelatedFile};use edit_prediction_types::EditPredictionRequestTrigger;use feature_flags::{FeatureFlag, FeatureFlagAppExt as _, PresenceFlag, register_feature_flag};use futures::{    AsyncReadExt as _, FutureExt as _, StreamExt as _,    channel::mpsc::{self, UnboundedReceiver},    select_biased,};use git::repository::FileHistoryChangedFileSets;use gpui::BackgroundExecutor;use gpui::TaskExt;use gpui::http_client::Url;use gpui::{    App, AsyncApp, Context, Entity, EntityId, Global, SharedString, Task, WeakEntity, actions,    http_client::{self, AsyncBody, Method},    prelude::*,};use heapless::Vec as ArrayVec;use language::{    Anchor, Buffer, BufferEditSource, BufferSnapshot, EditPredictionPromptFormat,    EditPredictionsMode, EditPreview, File, OffsetRangeExt, Point, TextBufferSnapshot, ToOffset,    ToPoint, language_settings::all_language_settings,};use project::{DisableAiSettings, Project, ProjectPath, WorktreeId};use release_channel::AppVersion;use semver::Version;use serde::de::DeserializeOwned;use settings::{    EditPredictionDataCollectionChoice, EditPredictionProvider, Settings as _, update_settings_file,};use std::collections::{VecDeque, hash_map};use std::env;use std::rc::Rc;use text::{AnchorRangeExt, Edit};use workspace::{AppState, Workspace};use zeta_prompt::ContextSource;use zeta_prompt::{Zeta2PromptInput, Zeta3PromptInput, ZetaFormat}; use std::mem;use std::ops::Range;use std::path::Path;use std::str::FromStr as _;use std::sync::Arc;use std::time::{Duration, Instant}; use thiserror::Error;use util::{ResultExt as _, rel_path::RelPath}; pub mod cursor_excerpt;pub mod data_collection;pub mod example_spec;pub mod fim;mod license_detection;pub mod mercury;pub mod metrics;pub mod ollama;mod onboarding_modal;pub mod open_ai_response;mod prediction; pub mod udiff; pub mod open_ai_compatible;mod zed_edit_prediction_delegate;pub mod zeta; #[cfg(test)]mod edit_prediction_tests; use crate::cursor_excerpt::expand_context_syntactically_then_linewise;use crate::data_collection::{CapturedPredictionContext, capture_prediction_context};use crate::example_spec::RecentFile;use crate::license_detection::LicenseDetectionWatcher;use crate::mercury::Mercury;pub use crate::metrics::{KeptRateResult, compute_kept_rate};use crate::onboarding_modal::ZedPredictModal;use crate::prediction::EditPredictionResult;pub use crate::prediction::{EditPrediction, EditPredictionId, EditPredictionInputs};pub use language_model::ApiKeyState;pub use telemetry_events::EditPredictionRating;pub use zed_edit_prediction_delegate::ZedEditPredictionDelegate; actions!(    edit_prediction,    [        /// Resets the edit prediction onboarding state.        ResetOnboarding,        /// Clears the edit prediction history.        ClearHistory,    ]); /// Maximum number of events to track.const EVENT_COUNT_MAX: usize = 10;const RECENT_PATH_COUNT_MAX: usize = 20;const CHANGE_GROUPING_LINE_SPAN: u32 = 8;const EDIT_HISTORY_DIFF_SIZE_LIMIT: usize = 2048 * 3; // ~2048 tokens or ~50% of typical prompt budgetconst COLLABORATOR_EDIT_LOCALITY_CONTEXT_TOKENS: usize = 512;const GIT_CHANGED_FILE_SETS_COMMIT_LIMIT: usize = 100;const LAST_CHANGE_GROUPING_TIME: Duration = Duration::from_secs(1);const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice";const REJECT_REQUEST_DEBOUNCE: Duration = Duration::from_secs(15);const REQUEST_TIMEOUT_BACKOFF: Duration = Duration::from_secs(10); const EDIT_PREDICTION_SETTLED_TTL: Duration = Duration::from_secs(60 * 5);const EDIT_PREDICTION_SETTLED_QUIESCENCE: Duration = Duration::from_secs(10);const EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS: usize = 4;const EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES: usize = 4 * 1024; pub struct EditPredictionJumpsFeatureFlag; impl FeatureFlag for EditPredictionJumpsFeatureFlag {    const NAME: &'static str = "edit_prediction_jumps";    type Value = PresenceFlag;}register_feature_flag!(EditPredictionJumpsFeatureFlag); #[derive(Clone)]struct EditPredictionStoreGlobal(Entity<EditPredictionStore>); impl Global for EditPredictionStoreGlobal {} /// Configuration for using the raw Zeta2 endpoint./// When set, the client uses the raw endpoint and constructs the prompt itself./// The version is also used as the Baseten environment name (lowercased).#[derive(Clone)]pub struct Zeta2RawConfig {    pub model_id: Option<String>,    pub environment: Option<String>,    pub format: ZetaFormat,} pub struct EditPredictionStore {    client: Arc<Client>,    user_store: Entity<UserStore>,    llm_token: LlmApiToken,    _fetch_experiments_task: Task<()>,    projects: HashMap<EntityId, ProjectState>,    update_required: bool,    edit_prediction_model: EditPredictionModel,    zeta2_raw_config: Option<Zeta2RawConfig>,    request_backoff_until: Option<Instant>,    preferred_experiment: Option<String>,    available_experiments: Vec<String>,    pub mercury: Mercury,    legacy_data_collection_enabled: bool,    reject_predictions_tx: mpsc::UnboundedSender<EditPredictionRejectionPayload>,    settled_predictions_tx: mpsc::UnboundedSender<Instant>,    rateable_predictions: VecDeque<EditPrediction>,    rated_predictions: HashSet<EditPredictionId>,    #[cfg(test)]    settled_event_callback: Option<Box<dyn Fn(EditPredictionId, String)>>,    credentials_provider: Arc<dyn CredentialsProvider>,} pub(crate) struct EditPredictionRejectionPayload {    rejection: EditPredictionRejection,    organization_id: Option<OrganizationId>,} #[derive(Copy, Clone, PartialEq, Eq)]pub enum EditPredictionModel {    Zeta,    Fim { format: EditPredictionPromptFormat },    Mercury,} pub struct EditPredictionModelInput {    project: Entity<Project>,    buffer: Entity<Buffer>,    snapshot: BufferSnapshot,    position: Anchor,    events: Vec<Arc<zeta_prompt::Event>>,    related_files: Vec<RelatedFile>,    editable_context: Option<Task<anyhow::Result<Vec<RelatedFile>>>>,    mode: PredictEditsMode,    trigger: PredictEditsRequestTrigger,    diagnostic_search_range: Range<Point>,    debug_tx: Option<mpsc::UnboundedSender<DebugEvent>>,    can_collect_data: bool,    is_open_source: bool,    allow_jump: bool,} #[derive(Debug)]pub enum DebugEvent {    ContextRetrievalStarted(ContextRetrievalStartedDebugEvent),    ContextRetrievalFinished(ContextRetrievalFinishedDebugEvent),    EditPredictionStarted(EditPredictionStartedDebugEvent),    EditPredictionFinished(EditPredictionFinishedDebugEvent),} #[derive(Debug)]pub struct ContextRetrievalStartedDebugEvent {    pub project_entity_id: EntityId,    pub timestamp: Instant,    pub search_prompt: String,} #[derive(Debug)]pub struct ContextRetrievalFinishedDebugEvent {    pub project_entity_id: EntityId,    pub timestamp: Instant,    pub metadata: Vec<(&'static str, SharedString)>,} #[derive(Debug)]pub struct EditPredictionStartedDebugEvent {    pub buffer: WeakEntity<Buffer>,    pub position: Anchor,    pub prompt: Option<String>,} #[derive(Debug)]pub struct EditPredictionFinishedDebugEvent {    pub buffer: WeakEntity<Buffer>,    pub position: Anchor,    pub model_output: Option<String>,} /// An event with associated metadata for reconstructing buffer state.#[derive(Clone)]pub struct StoredEvent {    pub event: Arc<zeta_prompt::Event>,    pub old_snapshot: TextBufferSnapshot,    pub new_snapshot_version: clock::Global,    pub total_edit_range: Range<Anchor>,    pub(crate) file_context: Option<Entity<StoredFileContext>>,} pub(crate) struct StoredFileContext {    pub(crate) uncommitted_diff: Option<Entity<BufferDiff>>,    pub(crate) git_changed_file_sets: Option<Arc<FileHistoryChangedFileSets>>,    pub(crate) git_changed_file_sets_task: Option<Task<()>>,} impl StoredEvent {    fn can_merge(        &self,        next_old_event: &StoredEvent,        latest_snapshot: &TextBufferSnapshot,        latest_edit_range: &Range<Anchor>,    ) -> bool {        // Events must be for the same buffer and be contiguous across included snapshots to be mergeable.        if self.old_snapshot.remote_id() != next_old_event.old_snapshot.remote_id() {            return false;        }        if self.old_snapshot.remote_id() != latest_snapshot.remote_id() {            return false;        }        if self.new_snapshot_version != next_old_event.old_snapshot.version {            return false;        }        if !latest_snapshot            .version            .observed_all(&next_old_event.new_snapshot_version)        {            return false;        }         let a_is_predicted = matches!(            self.event.as_ref(),            zeta_prompt::Event::BufferChange {                predicted: true,                ..            }        );        let b_is_predicted = matches!(            next_old_event.event.as_ref(),            zeta_prompt::Event::BufferChange {                predicted: true,                ..            }        );         // If events come from the same source (both predicted or both manual) then        // we would have coalesced them already.        if a_is_predicted == b_is_predicted {            return false;        }         let left_range = self.total_edit_range.to_point(latest_snapshot);        let right_range = next_old_event.total_edit_range.to_point(latest_snapshot);        let latest_range = latest_edit_range.to_point(latest_snapshot);         // Events near to the latest edit are not merged if their sources differ.        if lines_between_ranges(&left_range, &latest_range)            .min(lines_between_ranges(&right_range, &latest_range))            <= CHANGE_GROUPING_LINE_SPAN        {            return false;        }         // Events that are distant from each other are not merged.        if lines_between_ranges(&left_range, &right_range) > CHANGE_GROUPING_LINE_SPAN {            return false;        }         true    }} fn lines_between_ranges(left: &Range<Point>, right: &Range<Point>) -> u32 {    if left.start > right.end {        return left.start.row - right.end.row;    }    if right.start > left.end {        return right.start.row - left.end.row;    }    0} fn push_recent_file(files: &mut VecDeque<RecentFile>, mut file: RecentFile) {    if let Some(ix) = files.iter().position(|probe| probe.path == file.path)        && let Some(previous) = files.remove(ix)        && file.cursor_position.is_none()    {        file.cursor_position = previous.cursor_position;    }    files.push_front(file);    files.truncate(RECENT_PATH_COUNT_MAX);} struct ProjectState {    events: VecDeque<StoredEvent>,    last_event: Option<LastEvent>,    next_last_event_seq: u64,    recently_viewed_files: VecDeque<RecentFile>,    recently_opened_files: VecDeque<RecentFile>,    registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,    file_contexts: HashMap<ProjectPath, WeakEntity<StoredFileContext>>,    current_prediction: Option<CurrentEditPrediction>,    last_edit_source: Option<BufferEditSource>,    next_pending_prediction_id: usize,    pending_predictions: ArrayVec<PendingPrediction, 2, u8>,    pending_prediction_captures: Vec<PendingPredictionCapture>,    debug_tx: Option<mpsc::UnboundedSender<DebugEvent>>,    last_edit_prediction_refresh: Option<(EntityId, Instant)>,    cancelled_predictions: HashSet<usize>,    context: Entity<RelatedExcerptStore>,    license_detection_watchers: HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,    _subscriptions: [gpui::Subscription; 2],    copilot: Option<Entity<Copilot>>,} impl ProjectState {    pub fn events(&self, cx: &App) -> Vec<StoredEvent> {        self.events            .iter()            .cloned()            .chain(self.last_event.as_ref().iter().flat_map(|event| {                let (one, two) = event.split_by_pause();                let one = one.finalize(&self.license_detection_watchers, cx);                let two = two.and_then(|two| two.finalize(&self.license_detection_watchers, cx));                one.into_iter().chain(two)            }))            .collect()    }     fn cancel_pending_prediction(        &mut self,        pending_prediction: PendingPrediction,        cx: &mut Context<EditPredictionStore>,    ) {        self.cancelled_predictions.insert(pending_prediction.id);         if pending_prediction.drop_on_cancel {            drop(pending_prediction.task);        } else {            cx.spawn(async move |this, cx| {                let Some((prediction_id, model_version)) = pending_prediction.task.await else {                    return;                };                 this.update(cx, |this, cx| {                    this.reject_prediction(                        prediction_id,                        EditPredictionRejectReason::Canceled,                        false,                        model_version,                        None,                        cx,                    );                })                .ok();            })            .detach()        }    }     fn active_buffer(        &self,        project: &Entity<Project>,        cx: &App,    ) -> Option<(Entity<Buffer>, Option<Anchor>)> {        let project = project.read(cx);        let active_path = project.path_for_entry(project.active_entry()?, cx)?;        let active_buffer = project.buffer_store().read(cx).get_by_path(&active_path)?;        let registered_buffer = self.registered_buffers.get(&active_buffer.entity_id())?;        Some((active_buffer, registered_buffer.last_position))    }     fn file_context_for_path(        &mut self,        path: ProjectPath,        cx: &mut Context<EditPredictionStore>,    ) -> Entity<StoredFileContext> {        if let Some(context) = self            .file_contexts            .get_mut(&path)            .and_then(|entry| entry.upgrade())        {            context        } else {            let context = cx.new(|_| StoredFileContext {                uncommitted_diff: None,                git_changed_file_sets: None,                git_changed_file_sets_task: None,            });            self.file_contexts.insert(path, context.downgrade());            context        }    }     fn update_recent_file_cursor(&mut self, path: &Path, cursor_position: usize) {        for file in &mut self.recently_opened_files {            if file.path.as_ref() == path && file.cursor_position.is_none() {                file.cursor_position = Some(cursor_position);            }        }        for file in &mut self.recently_viewed_files {            if file.path.as_ref() == path {                file.cursor_position = Some(cursor_position);            }        }    }     fn finalize_last_event(&mut self, cx: &mut Context<EditPredictionStore>) {        let Some(last_event) = self.last_event.take() else {            return;        };        let event = last_event.finalize(&self.license_detection_watchers, cx);         for capture in &mut self.pending_prediction_captures {            capture.try_record_future_event(                &last_event,                event.as_ref(),                &self.license_detection_watchers,                cx,            );        }         let Some(event) = event else {            return;        };        if self.events.len() + 1 >= EVENT_COUNT_MAX {            self.events.pop_front();        }        self.events.push_back(event);    }     fn clear_history(&mut self) {        self.events.clear();        self.last_event.take();        for capture in &mut self.pending_prediction_captures {            capture.sample_data = None;        }    }} #[derive(Debug, Clone)]struct CurrentEditPrediction {    pub requested_by: EntityId,    pub prediction: EditPrediction,

Lines 1–500 of 3556. Load more