3687 lines · 129.7 KB · Rust
RawDownload
use std::fmt::Write as _;use std::io::Write;use std::ops::Range;use std::sync::Arc;use std::time::Instant; use agent_settings::{AgentSettings, CompletionMode};use anyhow::{Result, anyhow};use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};use chrono::{DateTime, Utc};use collections::HashMap;use editor::display_map::CreaseMetadata;use feature_flags::{self, FeatureFlagAppExt};use futures::future::Shared;use futures::{FutureExt, StreamExt as _};use git::repository::DiffType;use gpui::{    AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task,    WeakEntity,};use language_model::{    ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,    LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest,    LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,    LanguageModelToolResultContent, LanguageModelToolUseId, MessageContent,    ModelRequestLimitReachedError, PaymentRequiredError, RequestUsage, Role, SelectedModel,    StopReason, TokenUsage,};use postage::stream::Stream as _;use project::Project;use project::git_store::{GitStore, GitStoreCheckpoint, RepositoryState};use prompt_store::{ModelContext, PromptBuilder};use proto::Plan;use schemars::JsonSchema;use serde::{Deserialize, Serialize};use settings::Settings;use thiserror::Error;use ui::Window;use util::{ResultExt as _, post_inc};use uuid::Uuid;use zed_llm_client::{CompletionIntent, CompletionRequestStatus}; use crate::ThreadStore;use crate::context::{AgentContext, AgentContextHandle, ContextLoadResult, LoadedContext};use crate::thread_store::{    SerializedCrease, SerializedLanguageModel, SerializedMessage, SerializedMessageSegment,    SerializedThread, SerializedToolResult, SerializedToolUse, SharedProjectContext,};use crate::tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState}; #[derive(    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,)]pub struct ThreadId(Arc<str>); impl ThreadId {    pub fn new() -> Self {        Self(Uuid::new_v4().to_string().into())    }} impl std::fmt::Display for ThreadId {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        write!(f, "{}", self.0)    }} impl From<&str> for ThreadId {    fn from(value: &str) -> Self {        Self(value.into())    }} /// The ID of the user prompt that initiated a request.////// This equates to the user physically submitting a message to the model (e.g., by pressing the Enter key).#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]pub struct PromptId(Arc<str>); impl PromptId {    pub fn new() -> Self {        Self(Uuid::new_v4().to_string().into())    }} impl std::fmt::Display for PromptId {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        write!(f, "{}", self.0)    }} #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]pub struct MessageId(pub(crate) usize); impl MessageId {    fn post_inc(&mut self) -> Self {        Self(post_inc(&mut self.0))    }} /// Stored information that can be used to resurrect a context crease when creating an editor for a past message.#[derive(Clone, Debug)]pub struct MessageCrease {    pub range: Range<usize>,    pub metadata: CreaseMetadata,    /// None for a deserialized message, Some otherwise.    pub context: Option<AgentContextHandle>,} /// A message in a [`Thread`].#[derive(Debug, Clone)]pub struct Message {    pub id: MessageId,    pub role: Role,    pub segments: Vec<MessageSegment>,    pub loaded_context: LoadedContext,    pub creases: Vec<MessageCrease>,    pub is_hidden: bool,} impl Message {    /// Returns whether the message contains any meaningful text that should be displayed    /// The model sometimes runs tool without producing any text or just a marker ([`USING_TOOL_MARKER`])    pub fn should_display_content(&self) -> bool {        self.segments.iter().all(|segment| segment.should_display())    }     pub fn push_thinking(&mut self, text: &str, signature: Option<String>) {        if let Some(MessageSegment::Thinking {            text: segment,            signature: current_signature,        }) = self.segments.last_mut()        {            if let Some(signature) = signature {                *current_signature = Some(signature);            }            segment.push_str(text);        } else {            self.segments.push(MessageSegment::Thinking {                text: text.to_string(),                signature,            });        }    }     pub fn push_text(&mut self, text: &str) {        if let Some(MessageSegment::Text(segment)) = self.segments.last_mut() {            segment.push_str(text);        } else {            self.segments.push(MessageSegment::Text(text.to_string()));        }    }     pub fn to_string(&self) -> String {        let mut result = String::new();         if !self.loaded_context.text.is_empty() {            result.push_str(&self.loaded_context.text);        }         for segment in &self.segments {            match segment {                MessageSegment::Text(text) => result.push_str(text),                MessageSegment::Thinking { text, .. } => {                    result.push_str("<think>\n");                    result.push_str(text);                    result.push_str("\n</think>");                }                MessageSegment::RedactedThinking(_) => {}            }        }         result    }} #[derive(Debug, Clone, PartialEq, Eq)]pub enum MessageSegment {    Text(String),    Thinking {        text: String,        signature: Option<String>,    },    RedactedThinking(Vec<u8>),} impl MessageSegment {    pub fn should_display(&self) -> bool {        match self {            Self::Text(text) => text.is_empty(),            Self::Thinking { text, .. } => text.is_empty(),            Self::RedactedThinking(_) => false,        }    }} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct ProjectSnapshot {    pub worktree_snapshots: Vec<WorktreeSnapshot>,    pub unsaved_buffer_paths: Vec<String>,    pub timestamp: DateTime<Utc>,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct WorktreeSnapshot {    pub worktree_path: String,    pub git_state: Option<GitState>,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct GitState {    pub remote_url: Option<String>,    pub head_sha: Option<String>,    pub current_branch: Option<String>,    pub diff: Option<String>,} #[derive(Clone, Debug)]pub struct ThreadCheckpoint {    message_id: MessageId,    git_checkpoint: GitStoreCheckpoint,} #[derive(Copy, Clone, Debug, PartialEq, Eq)]pub enum ThreadFeedback {    Positive,    Negative,} pub enum LastRestoreCheckpoint {    Pending {        message_id: MessageId,    },    Error {        message_id: MessageId,        error: String,    },} impl LastRestoreCheckpoint {    pub fn message_id(&self) -> MessageId {        match self {            LastRestoreCheckpoint::Pending { message_id } => *message_id,            LastRestoreCheckpoint::Error { message_id, .. } => *message_id,        }    }} #[derive(Clone, Debug, Default, Serialize, Deserialize)]pub enum DetailedSummaryState {    #[default]    NotGenerated,    Generating {        message_id: MessageId,    },    Generated {        text: SharedString,        message_id: MessageId,    },} impl DetailedSummaryState {    fn text(&self) -> Option<SharedString> {        if let Self::Generated { text, .. } = self {            Some(text.clone())        } else {            None        }    }} #[derive(Default, Debug)]pub struct TotalTokenUsage {    pub total: usize,    pub max: usize,} impl TotalTokenUsage {    pub fn ratio(&self) -> TokenUsageRatio {        #[cfg(debug_assertions)]        let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")            .unwrap_or("0.8".to_string())            .parse()            .unwrap();        #[cfg(not(debug_assertions))]        let warning_threshold: f32 = 0.8;         // When the maximum is unknown because there is no selected model,        // avoid showing the token limit warning.        if self.max == 0 {            TokenUsageRatio::Normal        } else if self.total >= self.max {            TokenUsageRatio::Exceeded        } else if self.total as f32 / self.max as f32 >= warning_threshold {            TokenUsageRatio::Warning        } else {            TokenUsageRatio::Normal        }    }     pub fn add(&self, tokens: usize) -> TotalTokenUsage {        TotalTokenUsage {            total: self.total + tokens,            max: self.max,        }    }} #[derive(Debug, Default, PartialEq, Eq)]pub enum TokenUsageRatio {    #[default]    Normal,    Warning,    Exceeded,} #[derive(Debug, Clone, Copy)]pub enum QueueState {    Sending,    Queued { position: usize },    Started,} /// A thread of conversation with the LLM.pub struct Thread {    id: ThreadId,    updated_at: DateTime<Utc>,    summary: ThreadSummary,    pending_summary: Task<Option<()>>,    detailed_summary_task: Task<Option<()>>,    detailed_summary_tx: postage::watch::Sender<DetailedSummaryState>,    detailed_summary_rx: postage::watch::Receiver<DetailedSummaryState>,    completion_mode: agent_settings::CompletionMode,    messages: Vec<Message>,    next_message_id: MessageId,    last_prompt_id: PromptId,    project_context: SharedProjectContext,    checkpoints_by_message: HashMap<MessageId, ThreadCheckpoint>,    completion_count: usize,    pending_completions: Vec<PendingCompletion>,    project: Entity<Project>,    prompt_builder: Arc<PromptBuilder>,    tools: Entity<ToolWorkingSet>,    tool_use: ToolUseState,    action_log: Entity<ActionLog>,    last_restore_checkpoint: Option<LastRestoreCheckpoint>,    pending_checkpoint: Option<ThreadCheckpoint>,    initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,    request_token_usage: Vec<TokenUsage>,    cumulative_token_usage: TokenUsage,    exceeded_window_error: Option<ExceededWindowError>,    last_usage: Option<RequestUsage>,    tool_use_limit_reached: bool,    feedback: Option<ThreadFeedback>,    message_feedback: HashMap<MessageId, ThreadFeedback>,    last_auto_capture_at: Option<Instant>,    last_received_chunk_at: Option<Instant>,    request_callback: Option<        Box<dyn FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>])>,    >,    remaining_turns: u32,    configured_model: Option<ConfiguredModel>,} #[derive(Clone, Debug, PartialEq, Eq)]pub enum ThreadSummary {    Pending,    Generating,    Ready(SharedString),    Error,} impl ThreadSummary {    pub const DEFAULT: SharedString = SharedString::new_static("New Thread");     pub fn or_default(&self) -> SharedString {        self.unwrap_or(Self::DEFAULT)    }     pub fn unwrap_or(&self, message: impl Into<SharedString>) -> SharedString {        self.ready().unwrap_or_else(|| message.into())    }     pub fn ready(&self) -> Option<SharedString> {        match self {            ThreadSummary::Ready(summary) => Some(summary.clone()),            ThreadSummary::Pending | ThreadSummary::Generating | ThreadSummary::Error => None,        }    }} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct ExceededWindowError {    /// Model used when last message exceeded context window    model_id: LanguageModelId,    /// Token count including last message    token_count: usize,} impl Thread {    pub fn new(        project: Entity<Project>,        tools: Entity<ToolWorkingSet>,        prompt_builder: Arc<PromptBuilder>,        system_prompt: SharedProjectContext,        cx: &mut Context<Self>,    ) -> Self {        let (detailed_summary_tx, detailed_summary_rx) = postage::watch::channel();        let configured_model = LanguageModelRegistry::read_global(cx).default_model();         Self {            id: ThreadId::new(),            updated_at: Utc::now(),            summary: ThreadSummary::Pending,            pending_summary: Task::ready(None),            detailed_summary_task: Task::ready(None),            detailed_summary_tx,            detailed_summary_rx,            completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,            messages: Vec::new(),            next_message_id: MessageId(0),            last_prompt_id: PromptId::new(),            project_context: system_prompt,            checkpoints_by_message: HashMap::default(),            completion_count: 0,            pending_completions: Vec::new(),            project: project.clone(),            prompt_builder,            tools: tools.clone(),            last_restore_checkpoint: None,            pending_checkpoint: None,            tool_use: ToolUseState::new(tools.clone()),            action_log: cx.new(|_| ActionLog::new(project.clone())),            initial_project_snapshot: {                let project_snapshot = Self::project_snapshot(project, cx);                cx.foreground_executor()                    .spawn(async move { Some(project_snapshot.await) })                    .shared()            },            request_token_usage: Vec::new(),            cumulative_token_usage: TokenUsage::default(),            exceeded_window_error: None,            last_usage: None,            tool_use_limit_reached: false,            feedback: None,            message_feedback: HashMap::default(),            last_auto_capture_at: None,            last_received_chunk_at: None,            request_callback: None,            remaining_turns: u32::MAX,            configured_model,        }    }     pub fn deserialize(        id: ThreadId,        serialized: SerializedThread,        project: Entity<Project>,        tools: Entity<ToolWorkingSet>,        prompt_builder: Arc<PromptBuilder>,        project_context: SharedProjectContext,        window: Option<&mut Window>, // None in headless mode        cx: &mut Context<Self>,    ) -> Self {        let next_message_id = MessageId(            serialized                .messages                .last()                .map(|message| message.id.0 + 1)                .unwrap_or(0),        );        let tool_use = ToolUseState::from_serialized_messages(            tools.clone(),            &serialized.messages,            project.clone(),            window,            cx,        );        let (detailed_summary_tx, detailed_summary_rx) =            postage::watch::channel_with(serialized.detailed_summary_state);         let configured_model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| {            serialized                .model                .and_then(|model| {                    let model = SelectedModel {                        provider: model.provider.clone().into(),                        model: model.model.clone().into(),                    };                    registry.select_model(&model, cx)                })                .or_else(|| registry.default_model())        });         let completion_mode = serialized            .completion_mode            .unwrap_or_else(|| AgentSettings::get_global(cx).preferred_completion_mode);         Self {            id,

Lines 1–500 of 3687.