1443 lines · 49.1 KB · Rust
RawDownload
use crate::{ContextServerRegistry, SystemPromptTemplate, Template, Templates};use acp_thread::{MentionUri, UserMessageId};use action_log::ActionLog;use agent_client_protocol as acp;use agent_settings::{AgentProfileId, AgentSettings};use anyhow::{Context as _, Result, anyhow};use assistant_tool::adapt_schema_to_format;use cloud_llm_client::{CompletionIntent, CompletionMode};use collections::IndexMap;use fs::Fs;use futures::{    channel::{mpsc, oneshot},    stream::FuturesUnordered,};use gpui::{App, Context, Entity, SharedString, Task};use language_model::{    LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelImage,    LanguageModelProviderId, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelRequestTool, LanguageModelToolResult, LanguageModelToolResultContent,    LanguageModelToolSchemaFormat, LanguageModelToolUse, LanguageModelToolUseId, Role, StopReason,};use project::Project;use prompt_store::ProjectContext;use schemars::{JsonSchema, Schema};use serde::{Deserialize, Serialize};use settings::{Settings, update_settings_file};use smol::stream::StreamExt;use std::fmt::Write;use std::{cell::RefCell, collections::BTreeMap, path::Path, rc::Rc, sync::Arc};use util::{ResultExt, markdown::MarkdownCodeBlock}; #[derive(Debug, Clone, PartialEq, Eq)]pub enum Message {    User(UserMessage),    Agent(AgentMessage),} impl Message {    pub fn as_agent_message(&self) -> Option<&AgentMessage> {        match self {            Message::Agent(agent_message) => Some(agent_message),            _ => None,        }    }     pub fn to_markdown(&self) -> String {        match self {            Message::User(message) => message.to_markdown(),            Message::Agent(message) => message.to_markdown(),        }    }} #[derive(Debug, Clone, PartialEq, Eq)]pub struct UserMessage {    pub id: UserMessageId,    pub content: Vec<UserMessageContent>,} #[derive(Debug, Clone, PartialEq, Eq)]pub enum UserMessageContent {    Text(String),    Mention { uri: MentionUri, content: String },    Image(LanguageModelImage),} impl UserMessage {    pub fn to_markdown(&self) -> String {        let mut markdown = String::from("## User\n\n");         for content in &self.content {            match content {                UserMessageContent::Text(text) => {                    markdown.push_str(text);                    markdown.push('\n');                }                UserMessageContent::Image(_) => {                    markdown.push_str("<image />\n");                }                UserMessageContent::Mention { uri, content } => {                    if !content.is_empty() {                        markdown.push_str(&format!("{}\n\n{}\n", uri.to_link(), content));                    } else {                        markdown.push_str(&format!("{}\n", uri.to_link()));                    }                }            }        }         markdown    }     fn to_request(&self) -> LanguageModelRequestMessage {        let mut message = LanguageModelRequestMessage {            role: Role::User,            content: Vec::with_capacity(self.content.len()),            cache: false,        };         const OPEN_CONTEXT: &str = "<context>\n\            The following items were attached by the user. \            They are up-to-date and don't need to be re-read.\n\n";         const OPEN_FILES_TAG: &str = "<files>";        const OPEN_SYMBOLS_TAG: &str = "<symbols>";        const OPEN_THREADS_TAG: &str = "<threads>";        const OPEN_RULES_TAG: &str =            "<rules>\nThe user has specified the following rules that should be applied:\n";         let mut file_context = OPEN_FILES_TAG.to_string();        let mut symbol_context = OPEN_SYMBOLS_TAG.to_string();        let mut thread_context = OPEN_THREADS_TAG.to_string();        let mut rules_context = OPEN_RULES_TAG.to_string();         for chunk in &self.content {            let chunk = match chunk {                UserMessageContent::Text(text) => {                    language_model::MessageContent::Text(text.clone())                }                UserMessageContent::Image(value) => {                    language_model::MessageContent::Image(value.clone())                }                UserMessageContent::Mention { uri, content } => {                    match uri {                        MentionUri::File(path) | MentionUri::Symbol(path, _) => {                            write!(                                &mut symbol_context,                                "\n{}",                                MarkdownCodeBlock {                                    tag: &codeblock_tag(&path),                                    text: &content.to_string(),                                }                            )                            .ok();                        }                        MentionUri::Thread(_session_id) => {                            write!(&mut thread_context, "\n{}\n", content).ok();                        }                        MentionUri::Rule(_user_prompt_id) => {                            write!(                                &mut rules_context,                                "\n{}",                                MarkdownCodeBlock {                                    tag: "",                                    text: &content                                }                            )                            .ok();                        }                    }                     language_model::MessageContent::Text(uri.to_link())                }            };             message.content.push(chunk);        }         let len_before_context = message.content.len();         if file_context.len() > OPEN_FILES_TAG.len() {            file_context.push_str("</files>\n");            message                .content                .push(language_model::MessageContent::Text(file_context));        }         if symbol_context.len() > OPEN_SYMBOLS_TAG.len() {            symbol_context.push_str("</symbols>\n");            message                .content                .push(language_model::MessageContent::Text(symbol_context));        }         if thread_context.len() > OPEN_THREADS_TAG.len() {            thread_context.push_str("</threads>\n");            message                .content                .push(language_model::MessageContent::Text(thread_context));        }         if rules_context.len() > OPEN_RULES_TAG.len() {            rules_context.push_str("</user_rules>\n");            message                .content                .push(language_model::MessageContent::Text(rules_context));        }         if message.content.len() > len_before_context {            message.content.insert(                len_before_context,                language_model::MessageContent::Text(OPEN_CONTEXT.into()),            );            message                .content                .push(language_model::MessageContent::Text("</context>".into()));        }         message    }} impl AgentMessage {    pub fn to_markdown(&self) -> String {        let mut markdown = String::from("## Assistant\n\n");         for content in &self.content {            match content {                AgentMessageContent::Text(text) => {                    markdown.push_str(text);                    markdown.push('\n');                }                AgentMessageContent::Thinking { text, .. } => {                    markdown.push_str("<think>");                    markdown.push_str(text);                    markdown.push_str("</think>\n");                }                AgentMessageContent::RedactedThinking(_) => {                    markdown.push_str("<redacted_thinking />\n")                }                AgentMessageContent::Image(_) => {                    markdown.push_str("<image />\n");                }                AgentMessageContent::ToolUse(tool_use) => {                    markdown.push_str(&format!(                        "**Tool Use**: {} (ID: {})\n",                        tool_use.name, tool_use.id                    ));                    markdown.push_str(&format!(                        "{}\n",                        MarkdownCodeBlock {                            tag: "json",                            text: &format!("{:#}", tool_use.input)                        }                    ));                }            }        }         for tool_result in self.tool_results.values() {            markdown.push_str(&format!(                "**Tool Result**: {} (ID: {})\n\n",                tool_result.tool_name, tool_result.tool_use_id            ));            if tool_result.is_error {                markdown.push_str("**ERROR:**\n");            }             match &tool_result.content {                LanguageModelToolResultContent::Text(text) => {                    writeln!(markdown, "{text}\n").ok();                }                LanguageModelToolResultContent::Image(_) => {                    writeln!(markdown, "<image />\n").ok();                }            }             if let Some(output) = tool_result.output.as_ref() {                writeln!(                    markdown,                    "**Debug Output**:\n\n```json\n{}\n```\n",                    serde_json::to_string_pretty(output).unwrap()                )                .unwrap();            }        }         markdown    }     pub fn to_request(&self) -> Vec<LanguageModelRequestMessage> {        let mut content = Vec::with_capacity(self.content.len());        for chunk in &self.content {            let chunk = match chunk {                AgentMessageContent::Text(text) => {                    language_model::MessageContent::Text(text.clone())                }                AgentMessageContent::Thinking { text, signature } => {                    language_model::MessageContent::Thinking {                        text: text.clone(),                        signature: signature.clone(),                    }                }                AgentMessageContent::RedactedThinking(value) => {                    language_model::MessageContent::RedactedThinking(value.clone())                }                AgentMessageContent::ToolUse(value) => {                    language_model::MessageContent::ToolUse(value.clone())                }                AgentMessageContent::Image(value) => {                    language_model::MessageContent::Image(value.clone())                }            };            content.push(chunk);        }         let mut messages = vec![LanguageModelRequestMessage {            role: Role::Assistant,            content,            cache: false,        }];         if !self.tool_results.is_empty() {            let mut tool_results = Vec::with_capacity(self.tool_results.len());            for tool_result in self.tool_results.values() {                tool_results.push(language_model::MessageContent::ToolResult(                    tool_result.clone(),                ));            }            messages.push(LanguageModelRequestMessage {                role: Role::User,                content: tool_results,                cache: false,            });        }         messages    }} #[derive(Default, Debug, Clone, PartialEq, Eq)]pub struct AgentMessage {    pub content: Vec<AgentMessageContent>,    pub tool_results: IndexMap<LanguageModelToolUseId, LanguageModelToolResult>,} #[derive(Debug, Clone, PartialEq, Eq)]pub enum AgentMessageContent {    Text(String),    Thinking {        text: String,        signature: Option<String>,    },    RedactedThinking(String),    Image(LanguageModelImage),    ToolUse(LanguageModelToolUse),} #[derive(Debug)]pub enum AgentResponseEvent {    Text(String),    Thinking(String),    ToolCall(acp::ToolCall),    ToolCallUpdate(acp_thread::ToolCallUpdate),    ToolCallAuthorization(ToolCallAuthorization),    Stop(acp::StopReason),} #[derive(Debug)]pub struct ToolCallAuthorization {    pub tool_call: acp::ToolCall,    pub options: Vec<acp::PermissionOption>,    pub response: oneshot::Sender<acp::PermissionOptionId>,} pub struct Thread {    messages: Vec<Message>,    completion_mode: CompletionMode,    /// Holds the task that handles agent interaction until the end of the turn.    /// Survives across multiple requests as the model performs tool calls and    /// we run tools, report their results.    running_turn: Option<Task<()>>,    pending_agent_message: Option<AgentMessage>,    tools: BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,    context_server_registry: Entity<ContextServerRegistry>,    profile_id: AgentProfileId,    project_context: Rc<RefCell<ProjectContext>>,    templates: Arc<Templates>,    pub selected_model: Arc<dyn LanguageModel>,    project: Entity<Project>,    action_log: Entity<ActionLog>,} impl Thread {    pub fn new(        project: Entity<Project>,        project_context: Rc<RefCell<ProjectContext>>,        context_server_registry: Entity<ContextServerRegistry>,        action_log: Entity<ActionLog>,        templates: Arc<Templates>,        default_model: Arc<dyn LanguageModel>,        cx: &mut Context<Self>,    ) -> Self {        let profile_id = AgentSettings::get_global(cx).default_profile.clone();        Self {            messages: Vec::new(),            completion_mode: CompletionMode::Normal,            running_turn: None,            pending_agent_message: None,            tools: BTreeMap::default(),            context_server_registry,            profile_id,            project_context,            templates,            selected_model: default_model,            project,            action_log,        }    }     pub fn project(&self) -> &Entity<Project> {        &self.project    }     pub fn action_log(&self) -> &Entity<ActionLog> {        &self.action_log    }     pub fn set_mode(&mut self, mode: CompletionMode) {        self.completion_mode = mode;    }     #[cfg(any(test, feature = "test-support"))]    pub fn last_message(&self) -> Option<Message> {        if let Some(message) = self.pending_agent_message.clone() {            Some(Message::Agent(message))        } else {            self.messages.last().cloned()        }    }     pub fn add_tool(&mut self, tool: impl AgentTool) {        self.tools.insert(tool.name(), tool.erase());    }     pub fn remove_tool(&mut self, name: &str) -> bool {        self.tools.remove(name).is_some()    }     pub fn set_profile(&mut self, profile_id: AgentProfileId) {        self.profile_id = profile_id;    }     pub fn cancel(&mut self) {        // TODO: do we need to emit a stop::cancel for ACP?        self.running_turn.take();        self.flush_pending_agent_message();    }     pub fn truncate(&mut self, message_id: UserMessageId) -> Result<()> {        self.cancel();        let Some(position) = self.messages.iter().position(            |msg| matches!(msg, Message::User(UserMessage { id, .. }) if id == &message_id),        ) else {            return Err(anyhow!("Message not found"));        };        self.messages.truncate(position);        Ok(())    }     /// Sending a message results in the model streaming a response, which could include tool calls.    /// After calling tools, the model will stops and waits for any outstanding tool calls to be completed and their results sent.    /// The returned channel will report all the occurrences in which the model stops before erroring or ending its turn.    pub fn send<T>(        &mut self,        message_id: UserMessageId,        content: impl IntoIterator<Item = T>,        cx: &mut Context<Self>,    ) -> mpsc::UnboundedReceiver<Result<AgentResponseEvent, LanguageModelCompletionError>>    where        T: Into<UserMessageContent>,    {        let model = self.selected_model.clone();        let content = content.into_iter().map(Into::into).collect::<Vec<_>>();        log::info!("Thread::send called with model: {:?}", model.name());        log::debug!("Thread::send content: {:?}", content);         cx.notify();        let (events_tx, events_rx) =            mpsc::unbounded::<Result<AgentResponseEvent, LanguageModelCompletionError>>();        let event_stream = AgentResponseEventStream(events_tx);         let user_message_ix = self.messages.len();        self.messages.push(Message::User(UserMessage {            id: message_id,            content,        }));        log::info!("Total messages in thread: {}", self.messages.len());        self.running_turn = Some(cx.spawn(async move |thread, cx| {            log::info!("Starting agent turn execution");            let turn_result = async {                // Perform one request, then keep looping if the model makes tool calls.                let mut completion_intent = CompletionIntent::UserPrompt;                'outer: loop {                    log::debug!(                        "Building completion request with intent: {:?}",                        completion_intent                    );                    let request = thread.update(cx, |thread, cx| {                        thread.build_completion_request(completion_intent, cx)                    })?;                     // Stream events, appending to messages and collecting up tool uses.                    log::info!("Calling model.stream_completion");                    let mut events = model.stream_completion(request, cx).await?;                    log::debug!("Stream completion started successfully");                     let mut tool_uses = FuturesUnordered::new();                    while let Some(event) = events.next().await {                        match event {

Lines 1–500 of 1443.