use crate::tools::edit_file_tool::*;use crate::{    AgentTool, ContextServerRegistry, EditFileTool, GrepTool, GrepToolInput, ReadFileTool,    ReadFileToolInput, Template, Templates, Thread, ToolCallEventStream, ToolInput,};use Role::*;use anyhow::{Context as _, Result};use client::{Client, RefreshLlmTokenListener, UserStore};use fs::FakeFs;use futures::{FutureExt, StreamExt, future::LocalBoxFuture};use gpui::{AppContext as _, AsyncApp, Entity, TestAppContext, UpdateGlobal as _};use http_client::StatusCode;use language::language_settings::FormatOnSave;use language_model::{    LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,    LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelToolResult, LanguageModelToolResultContent, LanguageModelToolUse,    LanguageModelToolUseId, MessageContent, Role, SelectedModel,};use project::Project;use prompt_store::{ProjectContext, WorktreeContext};use rand::prelude::*;use reqwest_client::ReqwestClient;use serde::Serialize;use serde_json::json;use settings::SettingsStore;use std::{    fmt::{self, Display},    path::{Path, PathBuf},    str::FromStr,    sync::Arc,    time::Duration,};use util::path; #[derive(Serialize)]struct DiffJudgeTemplate {    diff: String,    assertions: &'static str,} impl Template for DiffJudgeTemplate {    const TEMPLATE_NAME: &'static str = "diff_judge.hbs";} #[derive(Clone)]struct EvalInput {    conversation: Vec<LanguageModelRequestMessage>,    input_file_path: PathBuf,    input_content: Option<String>,    assertion: EvalAssertion,} impl EvalInput {    fn new(        conversation: Vec<LanguageModelRequestMessage>,        input_file_path: impl Into<PathBuf>,        input_content: Option<String>,        assertion: EvalAssertion,    ) -> Self {        EvalInput {            conversation,            input_file_path: input_file_path.into(),            input_content,            assertion,        }    }} #[derive(Clone)]struct EvalSample {    text_before: String,    text_after: String,    tool_input: EditFileToolInput,    diff: String,} trait AssertionFn: 'static + Send + Sync {    fn assert<'a>(        &'a self,        sample: &'a EvalSample,        judge_model: Arc<dyn LanguageModel>,        cx: &'a mut TestAppContext,    ) -> LocalBoxFuture<'a, Result<EvalAssertionOutcome>>;} impl<F> AssertionFn for Fwhere    F: 'static        + Send        + Sync        + AsyncFn(            &EvalSample,            Arc<dyn LanguageModel>,            &mut TestAppContext,        ) -> Result<EvalAssertionOutcome>,{    fn assert<'a>(        &'a self,        sample: &'a EvalSample,        judge_model: Arc<dyn LanguageModel>,        cx: &'a mut TestAppContext,    ) -> LocalBoxFuture<'a, Result<EvalAssertionOutcome>> {        (self)(sample, judge_model, cx).boxed_local()    }} #[derive(Clone)]struct EvalAssertion(Arc<dyn AssertionFn>); impl EvalAssertion {    fn new<F>(f: F) -> Self    where        F: 'static            + Send            + Sync            + AsyncFn(                &EvalSample,                Arc<dyn LanguageModel>,                &mut TestAppContext,            ) -> Result<EvalAssertionOutcome>,    {        EvalAssertion(Arc::new(f))    }     fn assert_diff_any(expected_diffs: Vec<impl Into<String>>) -> Self {        let expected_diffs: Vec<String> = expected_diffs.into_iter().map(Into::into).collect();        Self::new(async move |sample, _judge, _cx| {            let matches = expected_diffs.iter().any(|possible_diff| {                language::apply_diff_patch(&sample.text_before, possible_diff)                    .map(|expected| {                        strip_empty_lines(&expected) == strip_empty_lines(&sample.text_after)                    })                    .unwrap_or(false)            });             Ok(EvalAssertionOutcome {                score: if matches { 100 } else { 0 },                message: None,            })        })    }     fn judge_diff(assertions: &'static str) -> Self {        Self::new(async move |sample, judge, cx| {            let prompt = DiffJudgeTemplate {                diff: sample.diff.clone(),                assertions,            }            .render(&Templates::new())            .context("Failed to render diff judge template")?;             let request = LanguageModelRequest {                messages: vec![LanguageModelRequestMessage {                    role: Role::User,                    content: vec![prompt.into()],                    cache: false,                    reasoning_details: None,                }],                thinking_allowed: true,                thinking_effort: judge                    .default_effort_level()                    .map(|effort_level| effort_level.value.to_string()),                ..Default::default()            };            let mut response = retry_on_rate_limit(async || {                Ok(judge                    .stream_completion_text(request.clone(), &cx.to_async())                    .await?)            })            .await?;            let mut output = String::new();            while let Some(chunk) = response.stream.next().await {                let chunk = chunk?;                output.push_str(&chunk);            }             let re = regex::Regex::new(r"<score>(\d+)</score>")                .context("Failed to compile score regex")?;            if let Some(captures) = re.captures(&output)                && let Some(score_match) = captures.get(1)            {                let score = score_match.as_str().parse().unwrap_or(0);                return Ok(EvalAssertionOutcome {                    score,                    message: Some(output),                });            }             anyhow::bail!("No score found in response. Raw output: {output}");        })    }     async fn run(        &self,        input: &EvalSample,        judge_model: Arc<dyn LanguageModel>,        cx: &mut TestAppContext,    ) -> Result<EvalAssertionOutcome> {        self.0.assert(input, judge_model, cx).await    }} #[derive(Clone)]struct EditEvalOutput {    sample: EvalSample,    assertion: EvalAssertionOutcome,} impl Display for EditEvalOutput {    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {        writeln!(f, "Score: {:?}", self.assertion.score)?;        if let Some(message) = self.assertion.message.as_ref() {            writeln!(f, "Message: {}", message)?;        }        writeln!(f, "Diff:\n{}", self.sample.diff)?;        writeln!(f, "Tool Input:\n{:#?}", self.sample.tool_input)?;        Ok(())    }} #[derive(Clone, Debug, Eq, PartialEq, Hash)]struct EvalAssertionOutcome {    score: usize,    message: Option<String>,} struct EditToolTest {    fs: Arc<FakeFs>,    project: Entity<Project>,    model: Arc<dyn LanguageModel>,    judge_model: Arc<dyn LanguageModel>,    model_thinking_effort: Option<String>,} impl EditToolTest {    async fn new(cx: &mut TestAppContext) -> Self {        cx.executor().allow_parking();         let fs = FakeFs::new(cx.executor());        cx.update(|cx| {            let settings_store = SettingsStore::test(cx);            cx.set_global(settings_store);            SettingsStore::update_global(cx, |store: &mut SettingsStore, cx| {                store.update_user_settings(cx, |settings| {                    settings                        .project                        .all_languages                        .defaults                        .ensure_final_newline_on_save = Some(false);                    settings.project.all_languages.defaults.format_on_save =                        Some(FormatOnSave::Off);                });            });             gpui_tokio::init(cx);            let http_client = Arc::new(ReqwestClient::user_agent("agent tests").unwrap());            cx.set_http_client(http_client);            let client = Client::production(cx);            let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));            language_model::init(cx);            RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);            language_models::init(user_store, client, cx);        });         fs.insert_tree("/root", json!({})).await;        let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;        let agent_model = SelectedModel::from_str(            &std::env::var("ZED_AGENT_MODEL")                .unwrap_or("anthropic/claude-sonnet-4-6-latest".into()),        )        .unwrap();        let judge_model = SelectedModel::from_str(            &std::env::var("ZED_JUDGE_MODEL")                .unwrap_or("anthropic/claude-sonnet-4-6-latest".into()),        )        .unwrap();         let authenticate_provider_tasks = cx.update(|cx| {            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {                registry                    .providers()                    .iter()                    .map(|p| p.authenticate(cx))                    .collect::<Vec<_>>()            })        });        let (model, judge_model) = cx            .update(|cx| {                cx.spawn(async move |cx| {                    futures::future::join_all(authenticate_provider_tasks).await;                    let model = Self::load_model(&agent_model, cx).await;                    let judge_model = Self::load_model(&judge_model, cx).await;                    (model.unwrap(), judge_model.unwrap())                })            })            .await;         let model_thinking_effort = model            .default_effort_level()            .map(|effort_level| effort_level.value.to_string());         Self {            fs,            project,            model,            judge_model,            model_thinking_effort,        }    }     async fn load_model(        selected_model: &SelectedModel,        cx: &mut AsyncApp,    ) -> Result<Arc<dyn LanguageModel>> {        cx.update(|cx| {            let registry = LanguageModelRegistry::read_global(cx);            let provider = registry                .provider(&selected_model.provider)                .expect("Provider not found");            provider.authenticate(cx)        })        .await?;        Ok(cx.update(|cx| {            let models = LanguageModelRegistry::read_global(cx);            models                .available_models(cx)                .find(|model| {                    model.provider_id() == selected_model.provider                        && model.id() == selected_model.model                })                .unwrap_or_else(|| panic!("Model {} not found", selected_model.model.0))        }))    }     async fn eval(&self, mut eval: EvalInput, cx: &mut TestAppContext) -> Result<EditEvalOutput> {        eval.conversation            .last_mut()            .context("Conversation must not be empty")?            .cache = true;         // Populate the FakeFs so `resolve_path` / `entry_for_path` can find        // the file in the worktree.        if let Some(input_content) = eval.input_content.as_deref() {            let abs_path = Path::new("/root").join(                eval.input_file_path                    .strip_prefix("root")                    .unwrap_or(&eval.input_file_path),            );            self.fs.insert_file(&abs_path, input_content.into()).await;             // Wait for the worktree to pick up the new file.            cx.run_until_parked();        }         let tools = crate::built_in_tools().collect::<Vec<_>>();         let system_prompt = {            let worktrees = vec![WorktreeContext {                root_name: "root".to_string(),                abs_path: Path::new("/path/to/root").into(),                rules_file: None,            }];            let project_context = ProjectContext::new(worktrees);            let tool_names = tools                .iter()                .map(|tool| tool.name.clone().into())                .collect::<Vec<_>>();            let template = crate::SystemPromptTemplate {                project: &project_context,                available_tools: tool_names,                model_name: None,                date: chrono::Local::now().format("%Y-%m-%d").to_string(),                user_agents_md: None,                sandboxing: false,                is_linux: cfg!(target_os = "linux"),                is_windows: cfg!(target_os = "windows"),            };            let templates = Templates::new();            template.render(&templates)?        };         let has_system_prompt = eval            .conversation            .first()            .is_some_and(|msg| msg.role == Role::System);        let messages = if has_system_prompt {            eval.conversation        } else {            [LanguageModelRequestMessage {                role: Role::System,                content: vec![MessageContent::Text(system_prompt)],                cache: true,                reasoning_details: None,            }]            .into_iter()            .chain(eval.conversation)            .collect::<Vec<_>>()        };         let request = LanguageModelRequest {            messages,            tools,            thinking_allowed: true,            thinking_effort: self.model_thinking_effort.clone(),            ..Default::default()        };         // The model will call the tool as "edit_file" (the production-visible        // name), but the schema is from EditFileTool.        let tool_input =            retry_on_rate_limit(async || self.extract_tool_use(request.clone(), cx).await).await?;         let language_registry = self            .project            .read_with(cx, |project, _cx| project.languages().clone());         let context_server_registry = cx            .new(|cx| ContextServerRegistry::new(self.project.read(cx).context_server_store(), cx));        let thread = cx.new(|cx| {            Thread::new(                self.project.clone(),                cx.new(|_cx| ProjectContext::default()),                context_server_registry,                Templates::new(),                Some(self.model.clone()),                cx,            )        });        let action_log = thread.read_with(cx, |thread, _| thread.action_log().clone());         let tool = Arc::new(EditFileTool::new(            self.project.clone(),            thread.downgrade(),            action_log,            language_registry,        ));         let result = cx            .update(|cx| {                tool.clone().run(                    ToolInput::resolved(tool_input.clone()),                    ToolCallEventStream::test().0,                    cx,                )            })            .await;         let output = match result {            Ok(output) => output,            Err(output) => {                anyhow::bail!("Tool returned error: {}", output);            }        };         let EditFileToolOutput::Success { new_text, .. } = &output else {            anyhow::bail!("Tool returned error output: {}", output);        };         let sample = EvalSample {            tool_input,            diff: language::unified_diff(                eval.input_content.as_deref().unwrap_or_default(),                new_text,            ),            text_before: eval.input_content.unwrap_or_default(),            text_after: new_text.clone(),        };         let assertion = eval            .assertion            .run(&sample, self.judge_model.clone(), cx)            .await?;         Ok(EditEvalOutput { assertion, sample })    }     /// Stream the model completion and extract the first complete tool use    /// whose name matches `EditFileTool::NAME`, parsed as `EditFileToolInput`.    async fn extract_tool_use(        &self,        request: LanguageModelRequest,        cx: &mut TestAppContext,    ) -> Result<EditFileToolInput> {        let model = self.model.clone();        let events = cx            .update(|cx| {                let async_cx = cx.to_async();                cx.foreground_executor()                    .spawn(async move { model.stream_completion(request, &async_cx).await })            })            .await            .map_err(|err| anyhow::anyhow!("completion error: {}", err))?;         let mut streamed_text = String::new();        let mut stop_reason = None;        let mut parse_errors = Vec::new();         let mut events = events.fuse();        while let Some(event) = events.next().await {

Lines 1–500 of 1487. Load more