use agent::{RequestKind, ThreadEvent, ThreadStore};use anyhow::{Context as _, Result, anyhow};use assistant_tool::ToolWorkingSet;use client::proto::LspWorkProgress;use collections::HashMap;use dap::DapRegistry;use futures::channel::mpsc;use futures::{FutureExt, StreamExt as _, select_biased};use gpui::{App, AppContext as _, AsyncApp, Entity, Task};use handlebars::Handlebars;use language::{DiagnosticSeverity, OffsetRangeExt};use language_model::{    LanguageModel, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,    StopReason, TokenUsage,};use project::{LspStore, Project, ProjectPath};use serde::{Deserialize, Serialize};use std::fmt::Write as _;use std::fs::File;use std::io::Write as _;use std::sync::{Arc, Mutex};use std::time::Duration;use std::{    fs,    path::{Path, PathBuf},};use unindent::Unindent as _;use util::ResultExt as _;use util::command::new_smol_command;use util::serde::default_true; use crate::AgentAppState; pub const EXAMPLES_DIR: &str = "./crates/eval/examples";pub const REPOS_DIR: &str = "./crates/eval/repos";pub const WORKTREES_DIR: &str = "./crates/eval/worktrees"; const THREAD_EVENT_TIMEOUT: Duration = Duration::from_secs(60 * 2); #[derive(Clone, Debug, Deserialize)]pub struct ExampleBase {    pub url: String,    pub revision: String,    pub language_extension: Option<String>,    pub insert_id: Option<String>,    #[serde(default = "default_true")]    pub require_lsp: bool,} #[derive(Clone, Debug)]pub struct Example {    pub name: String,    /// Content of `base.toml`    pub base: ExampleBase,    /// Content of `prompt.md`    pub prompt: String,    /// Content of `criteria.md`    pub criteria: String,    /// Markdown output file to append to    pub output_file: Option<Arc<Mutex<File>>>,    /// Path to markdown output file    pub output_file_path: PathBuf,    /// Prefix used for logging that identifies this example    pub log_prefix: String,} #[derive(Debug, Serialize, Deserialize, Clone)]pub struct RunOutput {    pub repository_diff: String,    pub diagnostics: String,    pub response_count: usize,    pub token_usage: TokenUsage,    pub tool_use_counts: HashMap<Arc<str>, u32>,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct JudgeInput {    pub repository_diff: String,    pub criteria: String,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct JudgeOutput {    pub analysis: String,    pub score: u32,} impl Example {    /// Load an example from a directory containing base.toml, prompt.md, and criteria.md    pub fn load_from_directory(dir_path: &Path, run_dir: &Path) -> Result<Self> {        let name = Self::name_from_path(dir_path);        let base_path = dir_path.join("base.toml");        let prompt_path = dir_path.join("prompt.md");        let criteria_path = dir_path.join("criteria.md");         let output_file_path = run_dir.join(format!(            "{}.md",            dir_path.file_name().unwrap().to_str().unwrap()        ));         Ok(Example {            name: name.clone(),            base: toml::from_str(&fs::read_to_string(&base_path)?)?,            prompt: fs::read_to_string(prompt_path.clone())?,            criteria: fs::read_to_string(criteria_path.clone())?,            output_file: None,            output_file_path,            log_prefix: name,        })    }     pub fn set_log_prefix_style(&mut self, color: &str, name_width: usize) {        self.log_prefix = format!(            "{}{:<width$}\x1b[0m | ",            color,            self.name,            width = name_width        );    }     pub fn name_from_path(path: &Path) -> String {        path.file_name().unwrap().to_string_lossy().to_string()    }     pub fn worktree_path(&self) -> PathBuf {        Path::new(WORKTREES_DIR)            .canonicalize()            .context(format!("No such directory {WORKTREES_DIR}"))            .unwrap()            .join(&self.name)    }     /// Set up the example by checking out the specified Git revision    pub async fn setup(&mut self) -> Result<()> {        let repo_path = repo_path_for_url(&self.base.url);         println!("{}Fetching", self.log_prefix);         run_git(            &repo_path,            &["fetch", "--depth", "1", "origin", &self.base.revision],        )        .await?;         let worktree_path = self.worktree_path();         if worktree_path.is_dir() {            println!("{}Resetting existing worktree", self.log_prefix);             // TODO: consider including "-x" to remove ignored files. The downside of this is that            // it will also remove build artifacts, and so prevent incremental reuse there.            run_git(&worktree_path, &["clean", "--force", "-d"]).await?;            run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;            run_git(&worktree_path, &["checkout", &self.base.revision]).await?;        } else {            println!("{}Creating worktree", self.log_prefix);             let worktree_path_string = worktree_path.to_string_lossy().to_string();             run_git(                &repo_path,                &[                    "worktree",                    "add",                    "-f",                    &worktree_path_string,                    &self.base.revision,                ],            )            .await?;        }         // Create the output file        let output_file = Arc::new(Mutex::new(File::create(&self.output_file_path)?));        self.output_file = Some(output_file);         Ok(())    }     /// Returns the output file, panicking if it's not set    fn output_file(&self) -> Arc<Mutex<File>> {        self.output_file            .clone()            .expect("Output file not created. Call setup() first.")    }     pub fn run(        &self,        model: Arc<dyn LanguageModel>,        app_state: Arc<AgentAppState>,        cx: &mut App,    ) -> Task<Result<RunOutput>> {        let project = Project::local(            app_state.client.clone(),            app_state.node_runtime.clone(),            app_state.user_store.clone(),            app_state.languages.clone(),            Arc::new(DapRegistry::default()),            app_state.fs.clone(),            None,            cx,        );         let worktree_path = self.worktree_path();        let worktree = project.update(cx, |project, cx| {            project.create_worktree(&worktree_path, true, cx)        });         let tools = cx.new(|_| ToolWorkingSet::default());        let thread_store =            ThreadStore::load(project.clone(), tools, app_state.prompt_builder.clone(), cx);        let this = self.clone();         cx.spawn(async move |cx| {            let worktree = worktree.await?;             // Wait for worktree scan to finish before choosing a file to open.            worktree                .update(cx, |worktree, _cx| {                    worktree.as_local().unwrap().scan_complete()                })?                .await;             let lsp_open_handle_and_store = if this.base.require_lsp {                let language_extension = this.base.language_extension.as_deref().context(                    "language_extension field is required in base.toml when `require_lsp == true`",                )?;                 // Open a file that matches the language to cause LSP to start.                let language_file = worktree.read_with(cx, |worktree, _cx| {                    worktree                        .files(false, 0)                        .find_map(|e| {                            if e.path.clone().extension().and_then(|ext| ext.to_str())                                == Some(language_extension)                            {                                Some(ProjectPath {                                    worktree_id: worktree.id(),                                    path: e.path.clone(),                                })                            } else {                                None                            }                        })                        .context("Failed to find a file for example language")                })??;                 let open_language_file_buffer_task = project.update(cx, |project, cx| {                    project.open_buffer(language_file.clone(), cx)                })?;                 let language_file_buffer = open_language_file_buffer_task.await?;                 let (lsp_open_handle, lsp_store) = project.update(cx, |project, cx| {                    (                        project.register_buffer_with_language_servers(&language_file_buffer, cx),                        project.lsp_store().clone(),                    )                })?;                 // TODO: remove this once the diagnostics tool waits for new diagnostics                cx.background_executor().timer(Duration::new(5, 0)).await;                wait_for_lang_server(&lsp_store, this.log_prefix.clone(), cx).await?;                 lsp_store.update(cx, |lsp_store, cx| {                    lsp_open_handle.update(cx, |buffer, cx| {                        buffer.update(cx, |buffer, cx| {                            let has_language_server = lsp_store                                .language_servers_for_local_buffer(buffer, cx)                                .next()                                .is_some();                            if has_language_server {                                Ok(())                            } else {                                Err(anyhow!(                                    "`{:?}` was opened to cause the language server to start, \                                    but no language servers are registered for its buffer. \                                    Set `require_lsp = false` in `base.toml` to skip this.",                                    language_file                                ))                            }                        })                    })                })??;                 Some((lsp_open_handle, lsp_store))            } else {                None            };             if std::env::var("ZED_EVAL_SETUP_ONLY").is_ok() {                return Err(anyhow!("Setup only mode"));            }             let thread_store = thread_store.await;            let thread =                thread_store.update(cx, |thread_store, cx| thread_store.create_thread(cx))?;             {                let output_file_ref = this.output_file();                let mut output_file = output_file_ref.lock().unwrap();                writeln!(&mut output_file, "👤 USER:").log_err();                writeln!(&mut output_file, "{}", this.prompt).log_err();                writeln!(&mut output_file, "🤖 ASSISTANT:").log_err();                output_file.flush().log_err();            }             let tool_use_counts: Arc<Mutex<HashMap<Arc<str>, u32>>> =                Mutex::new(HashMap::default()).into();             let (thread_event_tx, mut thread_event_rx) = mpsc::unbounded();             let subscription = cx.subscribe(&thread, move |_thread, event: &ThreadEvent, _cx| {                thread_event_tx.unbounded_send(event.clone()).log_err();            });             let event_handler_task = cx.spawn({                // Need to clone the Arc here because the reference from output_file() won't live long enough                let output_file = this.output_file.clone().unwrap();                let log_prefix = this.log_prefix.clone();                let tool_use_counts = tool_use_counts.clone();                let thread = thread.downgrade();                async move |cx| {                    loop {                        let event = select_biased! {                            event = thread_event_rx.next() => event,                            _ = cx.background_executor().timer(THREAD_EVENT_TIMEOUT).fuse() => {                                return Err(anyhow!("Agentic loop stalled - waited {:?} without any events", THREAD_EVENT_TIMEOUT));                            }                        };                        let Some(event) = event else {                            return Err(anyhow!("ThreadEvent channel ended early"));                        };                         let mut output_file = output_file.lock().unwrap();                         match event {                            ThreadEvent::Stopped(reason) => match reason {                                Ok(StopReason::EndTurn) => {                                    return Ok(());                                }                                Ok(StopReason::MaxTokens) => {                                    return Err(anyhow!("Exceeded maximum tokens"));                                }                                Ok(StopReason::ToolUse) => {                                    if std::env::var("ZED_EVAL_DEBUG").is_ok() {                                        println!("{}StopReason: Tool use", log_prefix);                                    }                                }                                Err(error) => {                                    return Err(anyhow!(error.clone()));                                }                            },                            ThreadEvent::ShowError(thread_error) => {                                break Err(anyhow!(thread_error.clone()));                            }                            ThreadEvent::StreamedAssistantText(_, chunk) => {                                write!(&mut output_file, "{}", chunk).log_err();                            }                            ThreadEvent::StreamedAssistantThinking(_, chunk) => {                                write!(&mut output_file, "{}", chunk).log_err();                            }                            ThreadEvent::UsePendingTools { tool_uses } => {                                writeln!(&mut output_file, "\n\nUSING TOOLS:").log_err();                                for tool_use in tool_uses {                                    writeln!(&mut output_file, "{}: {}", tool_use.name, tool_use.input)                                        .log_err();                                }                            }                            ThreadEvent::ToolFinished {                                tool_use_id,                                pending_tool_use,                                ..                            } => {                                if let Some(tool_use) = pending_tool_use {                                    let message = format!("TOOL FINISHED: {}", tool_use.name);                                    println!("{}{message}", log_prefix);                                    writeln!(&mut output_file, "\n{}", message).log_err();                                }                                thread.update(cx, |thread, _cx| {                                    if let Some(tool_result) = thread.tool_result(&tool_use_id) {                                        writeln!(&mut output_file, "\n{}\n", tool_result.content).log_err();                                        let mut tool_use_counts = tool_use_counts.lock().unwrap();                                        *tool_use_counts                                            .entry(tool_result.tool_name.clone())                                            .or_insert(0) += 1;                                    }                                })?;                            }                            ThreadEvent::ToolConfirmationNeeded => {                                panic!("{}Bug: Tool confirmation should not be required in eval", log_prefix);                            },                            ThreadEvent::StreamedCompletion |                            ThreadEvent::MessageAdded(_) |                            ThreadEvent::MessageEdited(_) |                            ThreadEvent::MessageDeleted(_) |                            ThreadEvent::SummaryChanged |                            ThreadEvent::SummaryGenerated |                            ThreadEvent::CheckpointChanged => {                                if std::env::var("ZED_EVAL_DEBUG").is_ok() {                                    println!("{}Event: {:#?}", log_prefix, event);                                }                            }                        }                         output_file.flush().log_err();                    }                }            });             thread.update(cx, |thread, cx| {                let context = vec![];                thread.insert_user_message(this.prompt.clone(), context, None, cx);                thread.send_to_model(model, RequestKind::Chat, cx);            })?;             event_handler_task.await?;             println!("{}Stopped", this.log_prefix);             if let Some((_, lsp_store)) = lsp_open_handle_and_store.as_ref() {                wait_for_lang_server(lsp_store, this.log_prefix.clone(), cx).await?;            }             println!("{}Getting repository diff", this.log_prefix);            let repository_diff = this.repository_diff().await?;             println!("{}Getting diagnostics", this.log_prefix);            let diagnostics = cx                .update(move |cx| {                    cx.spawn(async move |cx| query_lsp_diagnostics(project, cx).await)                })?                .await?;            println!("{}Got diagnostics", this.log_prefix);             drop(subscription);            drop(lsp_open_handle_and_store);             thread.update(cx, |thread, _cx| {                let response_count = thread                    .messages()                    .filter(|message| message.role == language_model::Role::Assistant)                    .count();                RunOutput {                    repository_diff,                    diagnostics,                    response_count,                    token_usage: thread.cumulative_token_usage(),                    tool_use_counts: tool_use_counts.lock().unwrap().clone(),                }            })        })    }     pub async fn judge(        &self,        model: Arc<dyn LanguageModel>,        repository_diff: String,        cx: &AsyncApp,    ) -> Result<JudgeOutput> {        let judge_prompt = include_str!("judge_prompt.hbs");        let judge_prompt_name = "judge_prompt";        let mut handlebars = Handlebars::new();        handlebars.register_template_string(judge_prompt_name, judge_prompt)?;        let prompt = handlebars.render(            judge_prompt_name,            &JudgeInput {                repository_diff,                criteria: self.criteria.clone(),            },        )?;         let request = LanguageModelRequest {            messages: vec![LanguageModelRequestMessage {                role: Role::User,                content: vec![MessageContent::Text(prompt)],                cache: false,            }],            temperature: None,            tools: Vec::new(),            stop: Vec::new(),        };         let response = send_language_model_request(model, request, cx).await?;         let output_file_ref = self.output_file();        let mut output_file = output_file_ref.lock().unwrap();         writeln!(&mut output_file, "\n\n").log_err();        writeln!(&mut output_file, "========================================").log_err();        writeln!(&mut output_file, "              JUDGE OUTPUT              ").log_err();        writeln!(&mut output_file, "========================================").log_err();        writeln!(&mut output_file, "\n{}", &response).log_err();         parse_judge_output(&response)    }     pub async fn repository_diff(&self) -> Result<String> {        let worktree_path = self.worktree_path();        run_git(&worktree_path, &["add", "-N"]).await?;

Lines 1–500 of 730. Load more