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