example.rs

   1use crate::{AgentAppState, ToolMetrics};
   2use agent::{ThreadEvent, ThreadStore};
   3use anyhow::{Context as _, Result, anyhow};
   4use assistant_tool::ToolWorkingSet;
   5use client::proto::LspWorkProgress;
   6use futures::channel::mpsc;
   7use futures::{FutureExt, StreamExt as _, select_biased};
   8use gpui::{App, AppContext as _, AsyncApp, Entity, Task};
   9use handlebars::Handlebars;
  10use language::{Buffer, DiagnosticSeverity, OffsetRangeExt};
  11use language_model::{
  12    LanguageModel, LanguageModelCompletionEvent, LanguageModelRequest, LanguageModelRequestMessage,
  13    MessageContent, Role, StopReason, TokenUsage,
  14};
  15use project::{Project, ProjectPath};
  16use serde::{Deserialize, Serialize};
  17use std::cell::RefCell;
  18use std::fmt::Write as _;
  19use std::fs::File;
  20use std::io::Write as _;
  21use std::rc::Rc;
  22use std::sync::{Arc, Mutex};
  23use std::time::Duration;
  24use std::{
  25    fs,
  26    path::{Path, PathBuf},
  27};
  28use unindent::Unindent as _;
  29use util::ResultExt as _;
  30use util::command::new_smol_command;
  31use util::markdown::MarkdownString;
  32use util::serde::default_true;
  33
  34pub const EXAMPLES_DIR: &str = "./crates/eval/examples";
  35pub const REPOS_DIR: &str = "./crates/eval/repos";
  36pub const WORKTREES_DIR: &str = "./crates/eval/worktrees";
  37
  38const THREAD_EVENT_TIMEOUT: Duration = Duration::from_secs(60 * 2);
  39
  40const ZED_REPO_URL: &str = "https://github.com/zed-industries/zed.git";
  41
  42#[derive(Clone, Debug, Deserialize)]
  43pub struct ExampleBase {
  44    pub url: String,
  45    pub revision: String,
  46    pub language_extension: Option<String>,
  47    pub insert_id: Option<String>,
  48    #[serde(default = "default_true")]
  49    pub require_lsp: bool,
  50    #[serde(default)]
  51    pub allow_preexisting_diagnostics: bool,
  52}
  53
  54impl ExampleBase {
  55    pub fn repo_name(&self) -> String {
  56        self.url
  57            .split('/')
  58            .next_back()
  59            .unwrap_or(&"")
  60            .trim_end_matches(".git")
  61            .into()
  62    }
  63}
  64
  65#[derive(Clone, Debug)]
  66pub struct Example {
  67    pub name: String,
  68    /// Content of `base.toml`
  69    pub base: ExampleBase,
  70    /// Content of `prompt.md`
  71    pub prompt: String,
  72    /// Content of `diff_criteria.md`
  73    pub diff_criteria: String,
  74    /// Content of `thread_criteria.md`, if that file exists (it's optional)
  75    pub thread_criteria: Option<String>,
  76    /// Path to the directory containing the requests and responses for the agentic loop
  77    pub run_directory_path: PathBuf,
  78    /// Prefix used for logging that identifies this example
  79    pub log_prefix: String,
  80}
  81
  82#[derive(Debug, Serialize, Deserialize, Clone)]
  83pub struct RunOutput {
  84    pub repository_diff: String,
  85    pub ran_diagnostics_check: bool,
  86    pub diagnostics_before: Option<String>,
  87    pub diagnostics_after: Option<String>,
  88    pub response_count: usize,
  89    pub token_usage: TokenUsage,
  90    pub tool_metrics: ToolMetrics,
  91    pub last_request: LanguageModelRequest,
  92}
  93
  94#[derive(Debug, Clone, Serialize, Deserialize)]
  95pub struct JudgeDiffInput {
  96    pub repository_diff: String,
  97    pub ran_diagnostics_check: bool,
  98    #[serde(skip_serializing_if = "Option::is_none")]
  99    pub diagnostics_before: Option<String>,
 100    #[serde(skip_serializing_if = "Option::is_none")]
 101    pub diagnostics_after: Option<String>,
 102    pub criteria: String,
 103}
 104
 105#[derive(Debug, Clone, Serialize, Deserialize)]
 106pub struct JudgeThreadInput {
 107    pub messages: String,
 108    pub criteria: String,
 109}
 110
 111#[derive(Debug, Clone, Serialize, Deserialize)]
 112pub struct JudgeResponse {
 113    pub analysis: String,
 114    pub score: u32,
 115}
 116
 117#[derive(Debug, Clone, Serialize, Deserialize)]
 118pub struct JudgeOutput {
 119    pub thread: Option<JudgeResponse>,
 120    pub diff: JudgeResponse,
 121}
 122
 123impl Example {
 124    /// Load an example from a directory containing base.toml, prompt.md, and criteria.md
 125    pub fn load_from_directory(dir_path: &Path, run_dir: &Path) -> Result<Self> {
 126        let name = Self::name_from_path(dir_path);
 127        let base_path = dir_path.join("base.toml");
 128        let prompt_path = dir_path.join("prompt.md");
 129        let diff_criteria_path = dir_path.join("diff_criteria.md");
 130        let thread_criteria_path = dir_path.join("thread_criteria.md");
 131        let thread_criteria = if thread_criteria_path.exists() {
 132            Some(fs::read_to_string(thread_criteria_path.clone())?)
 133        } else {
 134            None
 135        };
 136
 137        Ok(Example {
 138            name: name.clone(),
 139            base: toml::from_str(&fs::read_to_string(&base_path)?)?,
 140            prompt: fs::read_to_string(prompt_path.clone())?,
 141            thread_criteria,
 142            diff_criteria: fs::read_to_string(diff_criteria_path.clone())?,
 143            run_directory_path: run_dir.to_path_buf(),
 144            log_prefix: name,
 145        })
 146    }
 147
 148    pub fn set_repetition_number(&mut self, repetition_number: u32) {
 149        if repetition_number > 0 {
 150            self.name = format!("{}-{}", self.name, repetition_number);
 151        }
 152    }
 153
 154    pub fn example_output_directory(&self) -> PathBuf {
 155        self.run_directory_path.join(&self.name)
 156    }
 157
 158    pub fn set_log_prefix_style(&mut self, color: &str, name_width: usize) {
 159        self.log_prefix = format!(
 160            "{}{:<width$}\x1b[0m | ",
 161            color,
 162            self.name,
 163            width = name_width
 164        );
 165    }
 166
 167    pub fn name_from_path(path: &Path) -> String {
 168        path.file_name().unwrap().to_string_lossy().to_string()
 169    }
 170
 171    pub fn worktree_path(&self) -> PathBuf {
 172        Path::new(WORKTREES_DIR)
 173            .canonicalize()
 174            .context(format!("No such directory {WORKTREES_DIR}"))
 175            .unwrap()
 176            .join(&self.name)
 177            .join(self.base.repo_name())
 178    }
 179
 180    /// Set up the example by checking out the specified Git revision
 181    pub async fn setup(&mut self) -> Result<()> {
 182        let repo_path = repo_path_for_url(&self.base.url);
 183
 184        let revision_exists = run_git(
 185            &repo_path,
 186            &["rev-parse", &format!("{}^{{commit}}", self.base.revision)],
 187        )
 188        .await
 189        .is_ok();
 190
 191        if !revision_exists {
 192            println!(
 193                "{}Fetching revision {}",
 194                self.log_prefix, &self.base.revision
 195            );
 196            run_git(
 197                &repo_path,
 198                &["fetch", "--depth", "1", "origin", &self.base.revision],
 199            )
 200            .await?;
 201        }
 202
 203        let worktree_path = self.worktree_path();
 204
 205        if worktree_path.is_dir() {
 206            println!("{}Resetting existing worktree", self.log_prefix);
 207
 208            // TODO: consider including "-x" to remove ignored files. The downside of this is that
 209            // it will also remove build artifacts, and so prevent incremental reuse there.
 210            run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
 211            run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
 212            run_git(&worktree_path, &["checkout", &self.base.revision]).await?;
 213        } else {
 214            println!("{}Creating worktree", self.log_prefix);
 215
 216            let worktree_path_string = worktree_path.to_string_lossy().to_string();
 217
 218            run_git(
 219                &repo_path,
 220                &[
 221                    "worktree",
 222                    "add",
 223                    "-f",
 224                    &worktree_path_string,
 225                    &self.base.revision,
 226                ],
 227            )
 228            .await?;
 229        }
 230
 231        if self.base.url == ZED_REPO_URL {
 232            std::fs::write(worktree_path.join(".rules"), std::fs::read(".rules")?)?;
 233        }
 234
 235        std::fs::create_dir_all(self.example_output_directory())?;
 236
 237        Ok(())
 238    }
 239
 240    pub fn run(
 241        &self,
 242        model: Arc<dyn LanguageModel>,
 243        app_state: Arc<AgentAppState>,
 244        cx: &mut App,
 245    ) -> Task<Result<RunOutput>> {
 246        let project = Project::local(
 247            app_state.client.clone(),
 248            app_state.node_runtime.clone(),
 249            app_state.user_store.clone(),
 250            app_state.languages.clone(),
 251            app_state.fs.clone(),
 252            None,
 253            cx,
 254        );
 255
 256        let worktree_path = self.worktree_path();
 257        let worktree = project.update(cx, |project, cx| {
 258            project.create_worktree(&worktree_path, true, cx)
 259        });
 260
 261        let tools = cx.new(|_| ToolWorkingSet::default());
 262        let thread_store =
 263            ThreadStore::load(project.clone(), tools, app_state.prompt_builder.clone(), cx);
 264        let this = self.clone();
 265
 266        cx.spawn(async move |cx| {
 267            let worktree = worktree.await?;
 268
 269            // Wait for worktree scan to finish before choosing a file to open.
 270            worktree
 271                .update(cx, |worktree, _cx| {
 272                    worktree.as_local().unwrap().scan_complete()
 273                })?
 274                .await;
 275
 276            let lsp = if this.base.require_lsp {
 277                let language_extension = this.base.language_extension.as_deref().context(
 278                    "language_extension field is required in base.toml when `require_lsp == true`",
 279                )?;
 280
 281                // Open a file that matches the language to cause LSP to start.
 282                let language_file = worktree.read_with(cx, |worktree, _cx| {
 283                    worktree
 284                        .files(false, 0)
 285                        .find_map(|e| {
 286                            if e.path.clone().extension().and_then(|ext| ext.to_str())
 287                                == Some(language_extension)
 288                            {
 289                                Some(ProjectPath {
 290                                    worktree_id: worktree.id(),
 291                                    path: e.path.clone(),
 292                                })
 293                            } else {
 294                                None
 295                            }
 296                        })
 297                        .context("Failed to find a file for example language")
 298                })??;
 299
 300                let open_language_file_buffer_task = project.update(cx, |project, cx| {
 301                    project.open_buffer(language_file.clone(), cx)
 302                })?;
 303
 304                let language_file_buffer = open_language_file_buffer_task.await?;
 305
 306                let lsp_open_handle = project.update(cx, |project, cx| {
 307                    project.register_buffer_with_language_servers(&language_file_buffer, cx)
 308                })?;
 309
 310                wait_for_lang_server(&project, &language_file_buffer, this.log_prefix.clone(), cx).await?;
 311
 312                Some((lsp_open_handle, language_file_buffer))
 313            } else {
 314                None
 315            };
 316
 317            let diagnostics_before = query_lsp_diagnostics(project.clone(), cx).await?;
 318            if diagnostics_before.is_some() && !this.base.allow_preexisting_diagnostics {
 319                return Err(anyhow!("Example has pre-existing diagnostics. If you want to run this example regardless, set `allow_preexisting_diagnostics` to `true` in `base.toml`"));
 320            }
 321
 322            if std::env::var("ZED_EVAL_SETUP_ONLY").is_ok() {
 323                return Err(anyhow!("Setup only mode"));
 324            }
 325
 326            let thread_store = thread_store.await?;
 327            let thread =
 328                thread_store.update(cx, |thread_store, cx| thread_store.create_thread(cx))?;
 329            let last_request = Rc::new(RefCell::new(None));
 330
 331            thread.update(cx, |thread, _cx| {
 332                let mut request_count = 0;
 333                let example_dir_path = this.example_output_directory();
 334
 335                let last_request = Rc::clone(&last_request);
 336                thread.set_request_callback(move |request, response_events| {
 337                    *last_request.borrow_mut() = Some(request.clone());
 338
 339                    request_count += 1;
 340                    let messages_file_path = example_dir_path.join(format!("{request_count}.messages.md"));
 341                    let last_messages_file_path = example_dir_path.join("last.messages.md");
 342                    let request_markdown = RequestMarkdown::new(request);
 343                    let response_events_markdown = response_events_to_markdown(response_events);
 344
 345                    let messages = format!("{}\n\n{}", request_markdown.messages, response_events_markdown);
 346                    fs::write(messages_file_path, messages.clone()).expect("failed to write messages file");
 347                    fs::write(last_messages_file_path, messages).expect("failed to write last messages file");
 348
 349                    if request_count == 1 {
 350                        let tools_file_path = example_dir_path.join("tools.md");
 351                        fs::write(tools_file_path, request_markdown.tools).expect("failed to write tools file");
 352                    }
 353                });
 354            })?;
 355
 356            let tool_metrics = Arc::new(Mutex::new(ToolMetrics::default()));
 357
 358            let (thread_event_tx, mut thread_event_rx) = mpsc::unbounded();
 359
 360            let subscription = cx.subscribe(&thread, move |_thread, event: &ThreadEvent, _cx| {
 361                thread_event_tx.unbounded_send(event.clone()).log_err();
 362            });
 363
 364            let event_handler_task = cx.spawn({
 365                let log_prefix = this.log_prefix.clone();
 366                let tool_metrics = tool_metrics.clone();
 367                let thread = thread.downgrade();
 368                async move |cx| {
 369                    loop {
 370                        let event = select_biased! {
 371                            event = thread_event_rx.next() => event,
 372                            _ = cx.background_executor().timer(THREAD_EVENT_TIMEOUT).fuse() => {
 373                                return Err(anyhow!("Agentic loop stalled - waited {:?} without any events", THREAD_EVENT_TIMEOUT));
 374                            }
 375                        };
 376                        let Some(event) = event else {
 377                            return Err(anyhow!("ThreadEvent channel ended early"));
 378                        };
 379
 380                        match event {
 381                            ThreadEvent::Stopped(reason) => match reason {
 382                                Ok(StopReason::EndTurn) => {
 383                                    return Ok(());
 384                                }
 385                                Ok(StopReason::MaxTokens) => {
 386                                    return Err(anyhow!("Exceeded maximum tokens"));
 387                                }
 388                                Ok(StopReason::ToolUse) => {
 389                                    if std::env::var("ZED_EVAL_DEBUG").is_ok() {
 390                                        println!("{}StopReason: Tool use", log_prefix);
 391                                    }
 392                                }
 393                                Err(error) => {
 394                                    return Err(anyhow!(error.clone()));
 395                                }
 396                            },
 397                            ThreadEvent::ShowError(thread_error) => {
 398                                break Err(anyhow!(thread_error.clone()));
 399                            }
 400                            ThreadEvent::StreamedAssistantText(_, _)| ThreadEvent::StreamedAssistantThinking(_, _) | ThreadEvent::UsePendingTools { .. } => {
 401                            }
 402                            ThreadEvent::ToolFinished {
 403                                tool_use_id,
 404                                pending_tool_use,
 405                                ..
 406                            } => {
 407                                thread.update(cx, |thread, _cx| {
 408                                    if let Some(tool_use) = pending_tool_use {
 409                                        let mut tool_metrics = tool_metrics.lock().unwrap();
 410                                        if let Some(tool_result) = thread.tool_result(&tool_use_id) {
 411                                            let message = if tool_result.is_error {
 412                                                format!("TOOL FAILED: {}", tool_use.name)
 413                                            } else {
 414                                                format!("TOOL FINISHED: {}", tool_use.name)
 415                                            };
 416                                            println!("{log_prefix}{message}");
 417                                            tool_metrics.insert(tool_result.tool_name.clone(), !tool_result.is_error);
 418                                        } else {
 419                                            let message = format!("TOOL FINISHED WITHOUT RESULT: {}", tool_use.name);
 420                                            println!("{log_prefix}{message}");
 421                                            tool_metrics.insert(tool_use.name.clone(), true);
 422                                        }
 423                                    }
 424                                })?;
 425                            }
 426                            ThreadEvent::ToolConfirmationNeeded => {
 427                                panic!("{}Bug: Tool confirmation should not be required in eval", log_prefix);
 428                            },
 429                            ThreadEvent::StreamedToolUse { .. } |
 430                            ThreadEvent::StreamedCompletion |
 431                            ThreadEvent::MessageAdded(_) |
 432                            ThreadEvent::MessageEdited(_) |
 433                            ThreadEvent::MessageDeleted(_) |
 434                            ThreadEvent::SummaryChanged |
 435                            ThreadEvent::SummaryGenerated |
 436                            ThreadEvent::CheckpointChanged |
 437                            ThreadEvent::UsageUpdated(_) => {
 438                                if std::env::var("ZED_EVAL_DEBUG").is_ok() {
 439                                    println!("{}Event: {:#?}", log_prefix, event);
 440                                }
 441                            }
 442                        }
 443                    }
 444                }
 445            });
 446
 447            thread.update(cx, |thread, cx| {
 448                let context = vec![];
 449                thread.insert_user_message(this.prompt.clone(), context, None, cx);
 450                thread.send_to_model(model, cx);
 451            })?;
 452
 453            event_handler_task.await?;
 454
 455            println!("{}Stopped", this.log_prefix);
 456
 457            if let Some((_, language_file_buffer)) = lsp.as_ref() {
 458                wait_for_lang_server(&project, &language_file_buffer, this.log_prefix.clone(), cx).await?;
 459            }
 460
 461            println!("{}Getting repository diff", this.log_prefix);
 462            let repository_diff = this.repository_diff().await?;
 463
 464            let example_output_dir = this.example_output_directory();
 465            let repository_diff_path = example_output_dir.join("patch.diff");
 466            let mut repository_diff_output_file = File::create(&repository_diff_path)?;
 467            writeln!(&mut repository_diff_output_file, "{}", &repository_diff).log_err();
 468
 469            println!("{}Getting diagnostics", this.log_prefix);
 470            let diagnostics_after = cx
 471                .update(move |cx| {
 472                    cx.spawn(async move |cx| query_lsp_diagnostics(project, cx).await)
 473                })?
 474                .await?;
 475            println!("{}Got diagnostics", this.log_prefix);
 476
 477            let Some(last_request) = last_request.borrow_mut().take() else {
 478                return Err(anyhow!("No requests ran."));
 479            };
 480
 481            drop(subscription);
 482            drop(lsp);
 483
 484            if let Some(diagnostics_before) = &diagnostics_before {
 485                fs::write(example_output_dir.join("diagnostics_before.txt"), diagnostics_before)?;
 486            }
 487
 488            if let Some(diagnostics_after) = &diagnostics_after {
 489                fs::write(example_output_dir.join("diagnostics_after.txt"), diagnostics_after)?;
 490            }
 491
 492
 493            thread.update(cx, |thread, _cx| {
 494                let response_count = thread
 495                    .messages()
 496                    .filter(|message| message.role == language_model::Role::Assistant)
 497                    .count();
 498                RunOutput {
 499                    repository_diff,
 500                    ran_diagnostics_check: this.base.require_lsp,
 501                    diagnostics_before,
 502                    diagnostics_after,
 503                    response_count,
 504                    token_usage: thread.cumulative_token_usage(),
 505                    tool_metrics: tool_metrics.lock().unwrap().clone(),
 506                    last_request,
 507                }
 508            })
 509        })
 510    }
 511
 512    async fn judge_diff(
 513        &self,
 514        model: Arc<dyn LanguageModel>,
 515        run_output: &RunOutput,
 516        judge_number: u32,
 517        cx: &AsyncApp,
 518    ) -> Result<(String, JudgeResponse)> {
 519        let judge_diff_prompt = include_str!("judge_diff_prompt.hbs");
 520        let judge_diff_prompt_name = "judge_diff_prompt";
 521        let mut hbs = Handlebars::new();
 522        hbs.register_template_string(judge_diff_prompt_name, judge_diff_prompt)?;
 523
 524        let diff_prompt = hbs.render(
 525            judge_diff_prompt_name,
 526            &JudgeDiffInput {
 527                repository_diff: run_output.repository_diff.clone(),
 528                ran_diagnostics_check: run_output.ran_diagnostics_check,
 529                diagnostics_before: run_output.diagnostics_before.clone(),
 530                diagnostics_after: run_output.diagnostics_after.clone(),
 531                criteria: self.diff_criteria.clone(),
 532            },
 533        )?;
 534
 535        let request = LanguageModelRequest {
 536            thread_id: None,
 537            prompt_id: None,
 538            messages: vec![LanguageModelRequestMessage {
 539                role: Role::User,
 540                content: vec![MessageContent::Text(diff_prompt)],
 541                cache: false,
 542            }],
 543            temperature: None,
 544            tools: Vec::new(),
 545            stop: Vec::new(),
 546        };
 547
 548        let diff_response = send_language_model_request(model, request, cx).await?;
 549        let diff_output = JudgeResponse::parse(&diff_response)?;
 550
 551        println!(
 552            "{}Judge #{judge_number} - Diff score: {}",
 553            self.log_prefix, diff_output.score
 554        );
 555
 556        Ok((diff_response, diff_output))
 557    }
 558
 559    async fn judge_thread(
 560        &self,
 561        model: Arc<dyn LanguageModel>,
 562        run_output: &RunOutput,
 563        judge_number: u32,
 564        cx: &AsyncApp,
 565    ) -> Result<(String, Option<JudgeResponse>)> {
 566        if let Some(criteria) = self.thread_criteria.clone() {
 567            let judge_thread_prompt = include_str!("judge_thread_prompt.hbs");
 568            let judge_thread_prompt_name = "judge_thread_prompt";
 569            let mut hbs = Handlebars::new();
 570            hbs.register_template_string(judge_thread_prompt_name, judge_thread_prompt)?;
 571
 572            let request_markdown = RequestMarkdown::new(&run_output.last_request);
 573            let thread_prompt = hbs.render(
 574                judge_thread_prompt_name,
 575                &JudgeThreadInput {
 576                    messages: request_markdown.messages,
 577                    criteria,
 578                },
 579            )?;
 580
 581            let request = LanguageModelRequest {
 582                thread_id: None,
 583                prompt_id: None,
 584                messages: vec![LanguageModelRequestMessage {
 585                    role: Role::User,
 586                    content: vec![MessageContent::Text(thread_prompt)],
 587                    cache: false,
 588                }],
 589                temperature: None,
 590                tools: Vec::new(),
 591                stop: Vec::new(),
 592            };
 593
 594            let thread_response = send_language_model_request(model, request, cx).await?;
 595            let thread_output = JudgeResponse::parse(&thread_response)?;
 596
 597            println!(
 598                "{}Judge #{judge_number} - Thread score: {}",
 599                self.log_prefix, thread_output.score
 600            );
 601
 602            Ok((thread_response, Some(thread_output)))
 603        } else {
 604            let msg = "There were no criteria specified for this thread, so this example was not judged on its thread.".to_string();
 605            Ok((msg, None))
 606        }
 607    }
 608
 609    pub async fn judge(
 610        &self,
 611        model: Arc<dyn LanguageModel>,
 612        run_output: &RunOutput,
 613        judge_number: u32,
 614        cx: &AsyncApp,
 615    ) -> Result<JudgeOutput> {
 616        let mut output_file = File::create(
 617            self.example_output_directory()
 618                .join(format!("judge_{}.md", judge_number)),
 619        )
 620        .expect("failed to create judge.md");
 621
 622        println!("{}Running judge #{judge_number}", self.log_prefix);
 623
 624        let diff_task = self.judge_diff(model.clone(), &run_output, judge_number, cx);
 625        let thread_task = self.judge_thread(model.clone(), &run_output, judge_number, cx);
 626
 627        let (diff_result, thread_result) = futures::join!(diff_task, thread_task);
 628
 629        let (diff_response, diff_output) = diff_result?;
 630        let (thread_response, thread_output) = thread_result?;
 631
 632        writeln!(
 633            &mut output_file,
 634            "# Judgment\n\n## Thread\n\n{thread_response}\n\n## Diff\n\n{diff_response}",
 635        )
 636        .log_err();
 637
 638        Ok(JudgeOutput {
 639            thread: thread_output,
 640            diff: diff_output,
 641        })
 642    }
 643
 644    async fn repository_diff(&self) -> Result<String> {
 645        let worktree_path = self.worktree_path();
 646        run_git(&worktree_path, &["add", "."]).await?;
 647        let mut diff_args = vec!["diff", "--staged"];
 648        if self.base.url == ZED_REPO_URL {
 649            diff_args.push(":(exclude).rules");
 650        }
 651        run_git(&worktree_path, &diff_args).await
 652    }
 653}
 654
 655fn wait_for_lang_server(
 656    project: &Entity<Project>,
 657    buffer: &Entity<Buffer>,
 658    log_prefix: String,
 659    cx: &mut AsyncApp,
 660) -> Task<Result<()>> {
 661    println!("{}⏵ Waiting for language server", log_prefix);
 662
 663    let (mut tx, mut rx) = mpsc::channel(1);
 664
 665    let lsp_store = project
 666        .update(cx, |project, _| project.lsp_store())
 667        .unwrap();
 668
 669    let has_lang_server = buffer
 670        .update(cx, |buffer, cx| {
 671            lsp_store.update(cx, |lsp_store, cx| {
 672                lsp_store
 673                    .language_servers_for_local_buffer(&buffer, cx)
 674                    .next()
 675                    .is_some()
 676            })
 677        })
 678        .unwrap_or(false);
 679
 680    if has_lang_server {
 681        project
 682            .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))
 683            .unwrap()
 684            .detach();
 685    }
 686
 687    let subscriptions =
 688        [
 689            cx.subscribe(&lsp_store, {
 690                let log_prefix = log_prefix.clone();
 691                move |_, event, _| match event {
 692                    project::LspStoreEvent::LanguageServerUpdate {
 693                        message:
 694                            client::proto::update_language_server::Variant::WorkProgress(
 695                                LspWorkProgress {
 696                                    message: Some(message),
 697                                    ..
 698                                },
 699                            ),
 700                        ..
 701                    } => println!("{}{message}", log_prefix),
 702                    _ => {}
 703                }
 704            }),
 705            cx.subscribe(&project, {
 706                let buffer = buffer.clone();
 707                move |project, event, cx| match event {
 708                    project::Event::LanguageServerAdded(_, _, _) => {
 709                        let buffer = buffer.clone();
 710                        project
 711                            .update(cx, |project, cx| project.save_buffer(buffer, cx))
 712                            .detach();
 713                    }
 714                    project::Event::DiskBasedDiagnosticsFinished { .. } => {
 715                        tx.try_send(()).ok();
 716                    }
 717                    _ => {}
 718                }
 719            }),
 720        ];
 721
 722    cx.spawn(async move |cx| {
 723        let timeout = cx.background_executor().timer(Duration::new(60 * 5, 0));
 724        let result = futures::select! {
 725            _ = rx.next() => {
 726                println!("{}⚑ Language server idle", log_prefix);
 727                anyhow::Ok(())
 728            },
 729            _ = timeout.fuse() => {
 730                Err(anyhow!("LSP wait timed out after 5 minutes"))
 731            }
 732        };
 733        drop(subscriptions);
 734        result
 735    })
 736}
 737
 738async fn query_lsp_diagnostics(
 739    project: Entity<Project>,
 740    cx: &mut AsyncApp,
 741) -> Result<Option<String>> {
 742    let paths_with_diagnostics = project.update(cx, |project, cx| {
 743        project
 744            .diagnostic_summaries(true, cx)
 745            .filter(|(_, _, summary)| summary.error_count > 0 || summary.warning_count > 0)
 746            .map(|(project_path, _, _)| project_path)
 747            .collect::<Vec<_>>()
 748    })?;
 749
 750    if paths_with_diagnostics.is_empty() {
 751        return Ok(None);
 752    }
 753
 754    let mut output = String::new();
 755    for project_path in paths_with_diagnostics {
 756        let buffer = project
 757            .update(cx, |project, cx| project.open_buffer(project_path, cx))?
 758            .await?;
 759        let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
 760
 761        for (_, group) in snapshot.diagnostic_groups(None) {
 762            let entry = &group.entries[group.primary_ix];
 763            let range = entry.range.to_point(&snapshot);
 764            let severity = match entry.diagnostic.severity {
 765                DiagnosticSeverity::ERROR => "error",
 766                DiagnosticSeverity::WARNING => "warning",
 767                _ => continue,
 768            };
 769
 770            writeln!(
 771                output,
 772                "{} at line {}: {}",
 773                severity,
 774                range.start.row + 1,
 775                entry.diagnostic.message
 776            )?;
 777        }
 778    }
 779    anyhow::Ok(Some(output))
 780}
 781
 782impl JudgeResponse {
 783    fn parse(response: &str) -> Result<Self> {
 784        let analysis = get_tag("analysis", response)?.to_string();
 785        let score = get_tag("score", response)?
 786            .parse()
 787            .context("error parsing score")?;
 788
 789        Ok(Self { analysis, score })
 790    }
 791}
 792
 793fn get_tag(name: &'static str, response: &str) -> Result<String> {
 794    let start_tag = format!("<{}>", name);
 795    let end_tag = format!("</{}>", name);
 796
 797    let start_ix = response
 798        .find(&start_tag)
 799        .context(format!("{} start tag not found", name))?;
 800    let content_start_ix = start_ix + start_tag.len();
 801
 802    let end_ix = content_start_ix
 803        + response[content_start_ix..]
 804            .find(&end_tag)
 805            .context(format!("{} end tag not found", name))?;
 806
 807    let content = response[content_start_ix..end_ix].trim().unindent();
 808
 809    anyhow::Ok(content)
 810}
 811
 812pub fn repo_path_for_url(repo_url: &str) -> PathBuf {
 813    let repo_name = repo_url
 814        .trim_start_matches("https://")
 815        .replace(|c: char| !c.is_alphanumeric(), "-");
 816    Path::new(REPOS_DIR)
 817        .canonicalize()
 818        .context(format!("No such directory {REPOS_DIR}"))
 819        .unwrap()
 820        .join(repo_name)
 821}
 822
 823pub async fn run_git(repo_path: &Path, args: &[&str]) -> Result<String> {
 824    let output = new_smol_command("git")
 825        .current_dir(repo_path)
 826        .args(args)
 827        .output()
 828        .await?;
 829
 830    if output.status.success() {
 831        Ok(String::from_utf8(output.stdout)?.trim().to_string())
 832    } else {
 833        Err(anyhow!(
 834            "`git {}` within `{}` failed with status: {}\nstderr:\n{}\nstdout:\n{}",
 835            args.join(" "),
 836            repo_path.display(),
 837            output.status,
 838            String::from_utf8_lossy(&output.stderr),
 839            String::from_utf8_lossy(&output.stdout),
 840        ))
 841    }
 842}
 843
 844pub async fn send_language_model_request(
 845    model: Arc<dyn LanguageModel>,
 846    request: LanguageModelRequest,
 847    cx: &AsyncApp,
 848) -> anyhow::Result<String> {
 849    match model.stream_completion_text(request, &cx).await {
 850        Ok(mut stream) => {
 851            let mut full_response = String::new();
 852            while let Some(chunk_result) = stream.stream.next().await {
 853                match chunk_result {
 854                    Ok(chunk_str) => {
 855                        full_response.push_str(&chunk_str);
 856                    }
 857                    Err(err) => {
 858                        return Err(anyhow!(
 859                            "Error receiving response from language model: {err}"
 860                        ));
 861                    }
 862                }
 863            }
 864            Ok(full_response)
 865        }
 866        Err(err) => Err(anyhow!(
 867            "Failed to get response from language model. Error was: {err}"
 868        )),
 869    }
 870}
 871
 872struct RequestMarkdown {
 873    tools: String,
 874    messages: String,
 875}
 876
 877impl RequestMarkdown {
 878    fn new(request: &LanguageModelRequest) -> Self {
 879        let mut tools = String::new();
 880        let mut messages = String::new();
 881        let mut assistant_message_number: u32 = 1;
 882
 883        // Print the tools
 884        if !request.tools.is_empty() {
 885            for tool in &request.tools {
 886                write!(&mut tools, "# {}\n\n", tool.name).unwrap();
 887                write!(&mut tools, "{}\n\n", tool.description).unwrap();
 888                write!(
 889                    &mut tools,
 890                    "{}\n",
 891                    MarkdownString::code_block("json", &format!("{:#}", tool.input_schema))
 892                )
 893                .unwrap();
 894            }
 895        }
 896
 897        // Print the messages
 898        for message in &request.messages {
 899            match message.role {
 900                Role::System => messages.push_str("# ⚙️ SYSTEM\n\n"),
 901                Role::User => messages.push_str("# 👤 USER\n\n"),
 902                Role::Assistant => {
 903                    messages.push_str(&format!("# 🤖 ASSISTANT {assistant_message_number}\n\n"));
 904                    assistant_message_number += 1;
 905                }
 906            };
 907
 908            for content in &message.content {
 909                match content {
 910                    MessageContent::Text(text) => {
 911                        messages.push_str(text);
 912                        messages.push_str("\n\n");
 913                    }
 914                    MessageContent::Image(_) => {
 915                        messages.push_str("[IMAGE DATA]\n\n");
 916                    }
 917                    MessageContent::Thinking { text, signature } => {
 918                        messages.push_str("**Thinking**:\n\n");
 919                        if let Some(sig) = signature {
 920                            messages.push_str(&format!("Signature: {}\n\n", sig));
 921                        }
 922                        messages.push_str(text);
 923                        messages.push_str("\n");
 924                    }
 925                    MessageContent::RedactedThinking(items) => {
 926                        messages.push_str(&format!(
 927                            "**Redacted Thinking**: {} item(s)\n\n",
 928                            items.len()
 929                        ));
 930                    }
 931                    MessageContent::ToolUse(tool_use) => {
 932                        messages.push_str(&format!(
 933                            "**Tool Use**: {} (ID: {})\n",
 934                            tool_use.name, tool_use.id
 935                        ));
 936                        messages.push_str(&format!(
 937                            "{}\n",
 938                            MarkdownString::code_block("json", &format!("{:#}", tool_use.input))
 939                        ));
 940                    }
 941                    MessageContent::ToolResult(tool_result) => {
 942                        messages.push_str(&format!(
 943                            "**Tool Result**: {} (ID: {})\n\n",
 944                            tool_result.tool_name, tool_result.tool_use_id
 945                        ));
 946                        if tool_result.is_error {
 947                            messages.push_str("**ERROR:**\n");
 948                        }
 949                        messages.push_str(&format!("{}\n\n", tool_result.content));
 950                    }
 951                }
 952            }
 953        }
 954
 955        Self { tools, messages }
 956    }
 957}
 958
 959fn response_events_to_markdown(
 960    response_events: &[std::result::Result<LanguageModelCompletionEvent, String>],
 961) -> String {
 962    let mut response = String::new();
 963    // Print the response events if any
 964    response.push_str("# Response\n\n");
 965    let mut text_buffer = String::new();
 966    let mut thinking_buffer = String::new();
 967
 968    let flush_buffers =
 969        |output: &mut String, text_buffer: &mut String, thinking_buffer: &mut String| {
 970            if !text_buffer.is_empty() {
 971                output.push_str(&format!("**Text**:\n{}\n\n", text_buffer));
 972                text_buffer.clear();
 973            }
 974            if !thinking_buffer.is_empty() {
 975                output.push_str(&format!("**Thinking**:\n{}\n\n", thinking_buffer));
 976                thinking_buffer.clear();
 977            }
 978        };
 979
 980    for event in response_events {
 981        match event {
 982            Ok(LanguageModelCompletionEvent::Text(text)) => {
 983                text_buffer.push_str(text);
 984            }
 985            Ok(LanguageModelCompletionEvent::Thinking { text, .. }) => {
 986                thinking_buffer.push_str(text);
 987            }
 988            Ok(LanguageModelCompletionEvent::Stop(reason)) => {
 989                flush_buffers(&mut response, &mut text_buffer, &mut thinking_buffer);
 990                response.push_str(&format!("**Stop**: {:?}\n\n", reason));
 991            }
 992            Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
 993                flush_buffers(&mut response, &mut text_buffer, &mut thinking_buffer);
 994                response.push_str(&format!(
 995                    "**Tool Use**: {} (ID: {})\n",
 996                    tool_use.name, tool_use.id
 997                ));
 998                response.push_str(&format!(
 999                    "{}\n",
1000                    MarkdownString::code_block("json", &format!("{:#}", tool_use.input))
1001                ));
1002            }
1003            Ok(
1004                LanguageModelCompletionEvent::UsageUpdate(_)
1005                | LanguageModelCompletionEvent::StartMessage { .. },
1006            ) => {}
1007            Err(error) => {
1008                flush_buffers(&mut response, &mut text_buffer, &mut thinking_buffer);
1009                response.push_str(&format!("**Error**: {}\n\n", error));
1010            }
1011        }
1012    }
1013
1014    flush_buffers(&mut response, &mut text_buffer, &mut thinking_buffer);
1015
1016    response
1017}
1018
1019#[cfg(test)]
1020mod test {
1021    use super::*;
1022    use handlebars::Handlebars;
1023
1024    #[test]
1025    fn test_parse_judge_output() {
1026        let response = r#"
1027            <analysis>The model did a good job but there were still compilations errors.</analysis>
1028            <score>3</score>
1029        "#
1030        .unindent();
1031
1032        let output = JudgeResponse::parse(&response).unwrap();
1033        assert_eq!(
1034            output.analysis,
1035            "The model did a good job but there were still compilations errors."
1036        );
1037        assert_eq!(output.score, 3);
1038
1039        let response = r#"
1040            Text around ignored
1041
1042            <analysis>
1043                Failed to compile:
1044                - Error 1
1045                - Error 2
1046            </analysis>
1047
1048            <score>1</score>
1049        "#
1050        .unindent();
1051
1052        let output = JudgeResponse::parse(&response).unwrap();
1053        assert_eq!(output.analysis, "Failed to compile:\n- Error 1\n- Error 2");
1054        assert_eq!(output.score, 1);
1055    }
1056
1057    #[test]
1058    fn test_judge_prompt_with_diagnostics() {
1059        // Case 1: Both diagnostics before and after are present
1060        let input = JudgeDiffInput {
1061            repository_diff: "diff content goes here".to_string(),
1062            ran_diagnostics_check: true,
1063            diagnostics_before: Some("Error at line 10: variable not found".to_string()),
1064            diagnostics_after: Some("Error at line 15: missing semicolon".to_string()),
1065            criteria: "Fix all bugs".to_string(),
1066        };
1067
1068        let rendered = templates().render(JUDGE_PROMPT_NAME, &input).unwrap();
1069
1070        let expected_diagnostics_section = r#"
1071            Take into account the diagnostics before and after applying the change:
1072
1073            <diagnostics_before>
1074            Error at line 10: variable not found
1075            </diagnostics_before>
1076
1077            <diagnostics_after>
1078            Error at line 15: missing semicolon
1079            </diagnostics_after>
1080            "#
1081        .unindent();
1082
1083        assert!(rendered.contains(&expected_diagnostics_section));
1084    }
1085
1086    #[test]
1087    fn test_judge_prompt_with_empty_diagnostics() {
1088        // Case 2: Diagnostics check run but no diagnostics found
1089        let input = JudgeDiffInput {
1090            repository_diff: "diff content goes here".to_string(),
1091            ran_diagnostics_check: true,
1092            diagnostics_before: None,
1093            diagnostics_after: None,
1094            criteria: "Fix all bugs".to_string(),
1095        };
1096
1097        let rendered = templates().render(JUDGE_PROMPT_NAME, &input).unwrap();
1098
1099        let expected_diagnostics_section = r#"
1100            Take into account the diagnostics before and after applying the change:
1101
1102            <diagnostics_before>
1103            No diagnostics before applying the edits.
1104            </diagnostics_before>
1105
1106            <diagnostics_after>
1107            No diagnostics after applying the edits.
1108            </diagnostics_after>
1109            "#
1110        .unindent();
1111
1112        assert!(rendered.contains(&expected_diagnostics_section));
1113    }
1114
1115    #[test]
1116    fn test_judge_prompt_with_mixed_diagnostics() {
1117        let templates = templates();
1118
1119        // Case 3: Before diagnostics present, after diagnostics absent
1120        let input = JudgeDiffInput {
1121            repository_diff: "diff content goes here".to_string(),
1122            ran_diagnostics_check: true,
1123            diagnostics_before: Some("Error at line 10: variable not found".to_string()),
1124            diagnostics_after: None,
1125            criteria: "Fix all bugs".to_string(),
1126        };
1127
1128        let rendered = templates.render(JUDGE_PROMPT_NAME, &input).unwrap();
1129
1130        let expected_diagnostics_section = r#"
1131            Take into account the diagnostics before and after applying the change:
1132
1133            <diagnostics_before>
1134            Error at line 10: variable not found
1135            </diagnostics_before>
1136
1137            <diagnostics_after>
1138            No diagnostics after applying the edits.
1139            </diagnostics_after>
1140            "#
1141        .unindent();
1142
1143        assert!(rendered.contains(&expected_diagnostics_section));
1144
1145        // Case 4: Before diagnostics absent, after diagnostics present
1146        let input = JudgeDiffInput {
1147            repository_diff: "diff content goes here".to_string(),
1148            ran_diagnostics_check: true,
1149            diagnostics_before: None,
1150            diagnostics_after: Some("Error at line 15: missing semicolon".to_string()),
1151            criteria: "Fix all bugs".to_string(),
1152        };
1153
1154        let rendered = templates.render(JUDGE_PROMPT_NAME, &input).unwrap();
1155
1156        let expected_diagnostics_section = r#"
1157            Take into account the diagnostics before and after applying the change:
1158
1159            <diagnostics_before>
1160            No diagnostics before applying the edits.
1161            </diagnostics_before>
1162
1163            <diagnostics_after>
1164            Error at line 15: missing semicolon
1165            </diagnostics_after>
1166            "#
1167        .unindent();
1168
1169        assert!(rendered.contains(&expected_diagnostics_section));
1170    }
1171
1172    #[test]
1173    fn test_judge_prompt_without_diagnostics() {
1174        let templates = templates();
1175
1176        // Case 5: No diagnostics check run
1177        let input = JudgeDiffInput {
1178            repository_diff: "diff content goes here".to_string(),
1179            ran_diagnostics_check: false,
1180            diagnostics_before: None,
1181            diagnostics_after: None,
1182            criteria: "Fix all bugs".to_string(),
1183        };
1184
1185        let rendered = templates.render(JUDGE_PROMPT_NAME, &input).unwrap();
1186
1187        // Check for the message when no diagnostics were performed
1188        let diagnostics_message = "No diagnostic checks were performed.";
1189
1190        assert!(rendered.contains(diagnostics_message));
1191        assert!(!rendered.contains("<diagnostics_before>"));
1192        assert!(!rendered.contains("<diagnostics_after>"));
1193    }
1194
1195    const JUDGE_PROMPT_NAME: &str = "judge_prompt";
1196
1197    fn templates() -> Handlebars<'static> {
1198        let mut judge_prompt = include_str!("judge_diff_prompt.hbs").to_string();
1199        language::LineEnding::normalize(&mut judge_prompt);
1200        let mut handlebars = Handlebars::new();
1201        handlebars
1202            .register_template_string(JUDGE_PROMPT_NAME, judge_prompt)
1203            .unwrap();
1204        handlebars
1205    }
1206}