tenant.openagents/omega

1use crate::{AgentTool, Template, Templates, TerminalTool, TerminalToolInput};
2use Role::*;
3use anyhow::{Context as _, Result};
4use client::{Client, RefreshLlmTokenListener, UserStore};
5use futures::{FutureExt as _, StreamExt};
6use gpui::{AppContext as _, AsyncApp, TestAppContext};
7use http_client::StatusCode;
8use language_model::{
9    LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
10    LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
11    SelectedModel,
12};
13use prompt_store::{ProjectContext, WorktreeContext};
14use rand::prelude::*;
15use reqwest_client::ReqwestClient;
16use settings::SettingsStore;
17use std::{
18    fmt::{self, Display},
19    path::Path,
20    str::FromStr,
21    sync::Arc,
22    time::Duration,
23};
24
25#[derive(Clone)]
26struct EvalInput {
27    conversation: Vec<LanguageModelRequestMessage>,
28    assertion: CommandAssertion,
29}
30
31impl EvalInput {
32    fn new(conversation: Vec<LanguageModelRequestMessage>, assertion: CommandAssertion) -> Self {
33        Self {
34            conversation,
35            assertion,
36        }
37    }
38}
39
40#[derive(Clone, Debug, Eq, PartialEq, Hash)]
41struct EvalAssertionOutcome {
42    score: usize,
43    message: Option<String>,
44}
45
46type AssertionFn = Arc<dyn Fn(&TerminalToolInput) -> EvalAssertionOutcome + Send + Sync + 'static>;
47
48#[derive(Clone)]
49struct CommandAssertion {
50    description: &'static str,
51    check: AssertionFn,
52}
53
54impl CommandAssertion {
55    fn new(
56        description: &'static str,
57        check: impl Fn(&TerminalToolInput) -> EvalAssertionOutcome + Send + Sync + 'static,
58    ) -> Self {
59        Self {
60            description,
61            check: Arc::new(check),
62        }
63    }
64
65    /// Passes when the command is a git command and every git subcommand that
66    /// could block on a pty (pager or editor) is guarded with the appropriate
67    /// environment variable or flag.
68    ///
69    /// This is intentionally permissive about *which* git subcommand the model
70    /// chooses — for an indirect prompt like "combine my last 3 commits", the
71    /// model is free to first investigate with `git log` or jump straight to
72    /// `git rebase -i`. Either is fine, as long as whatever it picks won't
73    /// hang on a pager or editor.
74    fn git_pty_safe(description: &'static str) -> Self {
75        Self::new(description, |input| {
76            let cmd = input.command.as_str();
77            let words: Vec<&str> = cmd.split_whitespace().collect();
78
79            if !words.contains(&"git") {
80                return EvalAssertionOutcome {
81                    score: 0,
82                    message: Some(format!("Expected a `git` command, got: {cmd}")),
83                };
84            }
85
86            // Subcommands that pipe their output through a pager by default,
87            // and so will hang on `less` unless one of these escape hatches is
88            // present somewhere in the command:
89            const PAGER_SUBCMDS: &[&str] = &["log", "diff", "show", "blame"];
90            const PAGER_GUARDS: &[&str] = &["--no-pager", "GIT_PAGER=cat", "PAGER=cat"];
91
92            // Subcommands that may invoke an interactive editor and so will
93            // hang unless one of these escape hatches is present:
94            const EDITOR_SUBCMDS: &[&str] = &["rebase", "commit", "merge", "tag"];
95            const EDITOR_GUARDS: &[&str] =
96                &["GIT_EDITOR=true", "GIT_EDITOR=:", "EDITOR=true", "EDITOR=:"];
97
98            let has_pager_guard = PAGER_GUARDS.iter().any(|guard| cmd.contains(guard));
99            let has_editor_guard = EDITOR_GUARDS.iter().any(|guard| cmd.contains(guard));
100
101            for subcmd in PAGER_SUBCMDS {
102                if words.contains(subcmd) && !has_pager_guard {
103                    return EvalAssertionOutcome {
104                        score: 0,
105                        message: Some(format!(
106                            "`git {subcmd}` is missing a pager guard \
107                             (one of {PAGER_GUARDS:?}). Command: {cmd}"
108                        )),
109                    };
110                }
111            }
112
113            for subcmd in EDITOR_SUBCMDS {
114                if words.contains(subcmd) && !has_editor_guard {
115                    return EvalAssertionOutcome {
116                        score: 0,
117                        message: Some(format!(
118                            "`git {subcmd}` is missing an editor guard \
119                             (one of {EDITOR_GUARDS:?}). Command: {cmd}"
120                        )),
121                    };
122                }
123            }
124
125            EvalAssertionOutcome {
126                score: 100,
127                message: None,
128            }
129        })
130    }
131}
132
133struct EvalOutput {
134    tool_input: TerminalToolInput,
135    assertion: EvalAssertionOutcome,
136    assertion_description: &'static str,
137}
138
139impl Display for EvalOutput {
140    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
141        writeln!(f, "Score: {}", self.assertion.score)?;
142        writeln!(f, "Assertion: {}", self.assertion_description)?;
143        if let Some(message) = self.assertion.message.as_ref() {
144            writeln!(f, "Message: {}", message)?;
145        }
146        writeln!(f, "Tool input: {:#?}", self.tool_input)?;
147        Ok(())
148    }
149}
150
151struct TerminalToolTest {
152    model: Arc<dyn LanguageModel>,
153    model_thinking_effort: Option<String>,
154}
155
156impl TerminalToolTest {
157    async fn new(cx: &mut TestAppContext) -> Self {
158        cx.executor().allow_parking();
159
160        cx.update(|cx| {
161            let settings_store = SettingsStore::test(cx);
162            cx.set_global(settings_store);
163
164            gpui_tokio::init(cx);
165            let http_client = Arc::new(ReqwestClient::user_agent("agent tests").unwrap());
166            cx.set_http_client(http_client);
167            let client = Client::production(cx);
168            let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
169            language_model::init(cx);
170            RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
171            language_models::init(user_store, client, cx);
172        });
173
174        let agent_model = SelectedModel::from_str(
175            &std::env::var("ZED_AGENT_MODEL")
176                .unwrap_or("anthropic/claude-sonnet-4-6-latest".into()),
177        )
178        .unwrap();
179
180        let authenticate_provider_tasks = cx.update(|cx| {
181            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
182                registry
183                    .providers()
184                    .iter()
185                    .map(|p| p.authenticate(cx))
186                    .collect::<Vec<_>>()
187            })
188        });
189
190        let model = cx
191            .update(|cx| {
192                cx.spawn(async move |cx| {
193                    futures::future::join_all(authenticate_provider_tasks).await;
194                    load_model(&agent_model, cx).await.unwrap()
195                })
196            })
197            .await;
198
199        let model_thinking_effort = model
200            .default_effort_level()
201            .map(|effort_level| effort_level.value.to_string());
202
203        Self {
204            model,
205            model_thinking_effort,
206        }
207    }
208
209    async fn eval(&self, mut eval: EvalInput, cx: &mut TestAppContext) -> Result<EvalOutput> {
210        eval.conversation
211            .last_mut()
212            .context("Conversation must not be empty")?
213            .cache = true;
214
215        let tools = crate::built_in_tools().collect::<Vec<_>>();
216
217        let system_prompt = {
218            let worktrees = vec![WorktreeContext {
219                root_name: "root".to_string(),
220                abs_path: Path::new("/path/to/root").into(),
221                rules_file: None,
222            }];
223            let project_context = ProjectContext::new(worktrees);
224            let tool_names = tools
225                .iter()
226                .map(|tool| tool.name.clone().into())
227                .collect::<Vec<_>>();
228            let template = crate::SystemPromptTemplate {
229                project: &project_context,
230                available_tools: tool_names,
231                model_name: None,
232                date: chrono::Local::now().format("%Y-%m-%d").to_string(),
233                user_agents_md: None,
234                sandboxing: false,
235                is_linux: cfg!(target_os = "linux"),
236                is_windows: cfg!(target_os = "windows"),
237            };
238            template.render(&Templates::new())?
239        };
240
241        let has_system_prompt = eval
242            .conversation
243            .first()
244            .is_some_and(|msg| msg.role == Role::System);
245        let messages = if has_system_prompt {
246            eval.conversation
247        } else {
248            [LanguageModelRequestMessage {
249                role: Role::System,
250                content: vec![MessageContent::Text(system_prompt)],
251                cache: true,
252                reasoning_details: None,
253            }]
254            .into_iter()
255            .chain(eval.conversation)
256            .collect::<Vec<_>>()
257        };
258
259        let request = LanguageModelRequest {
260            messages,
261            tools,
262            thinking_allowed: true,
263            thinking_effort: self.model_thinking_effort.clone(),
264            ..Default::default()
265        };
266
267        let tool_input =
268            retry_on_rate_limit(async || extract_tool_use(&self.model, request.clone(), cx).await)
269                .await?;
270
271        let assertion = (eval.assertion.check)(&tool_input);
272        Ok(EvalOutput {
273            tool_input,
274            assertion,
275            assertion_description: eval.assertion.description,
276        })
277    }
278}
279
280async fn load_model(
281    selected_model: &SelectedModel,
282    cx: &mut AsyncApp,
283) -> Result<Arc<dyn LanguageModel>> {
284    cx.update(|cx| {
285        let registry = LanguageModelRegistry::read_global(cx);
286        let provider = registry
287            .provider(&selected_model.provider)
288            .expect("Provider not found");
289        provider.authenticate(cx)
290    })
291    .await?;
292    Ok(cx.update(|cx| {
293        let models = LanguageModelRegistry::read_global(cx);
294        models
295            .available_models(cx)
296            .find(|model| {
297                model.provider_id() == selected_model.provider && model.id() == selected_model.model
298            })
299            .unwrap_or_else(|| panic!("Model {} not found", selected_model.model.0))
300    }))
301}
302
303/// Stream the model completion and extract the first complete tool use whose
304/// name matches `TerminalTool::NAME`, parsed as `TerminalToolInput`.
305async fn extract_tool_use(
306    model: &Arc<dyn LanguageModel>,
307    request: LanguageModelRequest,
308    cx: &mut TestAppContext,
309) -> Result<TerminalToolInput> {
310    let model = model.clone();
311    let events = cx
312        .update(|cx| {
313            let async_cx = cx.to_async();
314            cx.foreground_executor()
315                .spawn(async move { model.stream_completion(request, &async_cx).await })
316        })
317        .await
318        .map_err(|err| anyhow::anyhow!("completion error: {}", err))?;
319
320    let mut streamed_text = String::new();
321    let mut stop_reason = None;
322    let mut parse_errors = Vec::new();
323
324    let mut events = events.fuse();
325    while let Some(event) = events.next().await {
326        match event {
327            Ok(LanguageModelCompletionEvent::ToolUse(tool_use))
328                if tool_use.is_input_complete && tool_use.name.as_ref() == TerminalTool::NAME =>
329            {
330                let input: TerminalToolInput = tool_use
331                    .input
332                    .parse()
333                    .context("Failed to parse tool input as TerminalToolInput")?;
334                return Ok(input);
335            }
336            Ok(LanguageModelCompletionEvent::Text(text)) => {
337                if streamed_text.len() < 2_000 {
338                    streamed_text.push_str(&text);
339                }
340            }
341            Ok(LanguageModelCompletionEvent::Stop(reason)) => {
342                stop_reason = Some(reason);
343            }
344            Ok(LanguageModelCompletionEvent::ToolUseJsonParseError {
345                tool_name,
346                raw_input,
347                json_parse_error,
348                ..
349            }) if tool_name.as_ref() == TerminalTool::NAME => {
350                parse_errors.push(format!("{json_parse_error}\nRaw input:\n{raw_input:?}"));
351            }
352            Err(err) => {
353                return Err(anyhow::anyhow!("completion error: {}", err));
354            }
355            _ => {}
356        }
357    }
358
359    let streamed_text = streamed_text.trim();
360    let streamed_text_suffix = if streamed_text.is_empty() {
361        String::new()
362    } else {
363        format!("\nStreamed text:\n{streamed_text}")
364    };
365    let stop_reason_suffix = stop_reason
366        .map(|reason| format!("\nStop reason: {reason:?}"))
367        .unwrap_or_default();
368    let parse_errors_suffix = if parse_errors.is_empty() {
369        String::new()
370    } else {
371        format!("\nTool parse errors:\n{}", parse_errors.join("\n"))
372    };
373
374    anyhow::bail!(
375        "Stream ended without a terminal tool use{stop_reason_suffix}{parse_errors_suffix}{streamed_text_suffix}"
376    )
377}
378
379async fn retry_on_rate_limit<R>(mut request: impl AsyncFnMut() -> Result<R>) -> Result<R> {
380    const MAX_RETRIES: usize = 20;
381    let mut attempt = 0;
382
383    loop {
384        attempt += 1;
385        let response = request().await;
386
387        if attempt >= MAX_RETRIES {
388            return response;
389        }
390
391        let retry_delay = match &response {
392            Ok(_) => None,
393            Err(err) => match err.downcast_ref::<LanguageModelCompletionError>() {
394                Some(err) => match &err {
395                    LanguageModelCompletionError::RateLimitExceeded { retry_after, .. }
396                    | LanguageModelCompletionError::ServerOverloaded { retry_after, .. } => {
397                        Some(retry_after.unwrap_or(Duration::from_secs(5)))
398                    }
399                    LanguageModelCompletionError::UpstreamProviderError {
400                        status,
401                        retry_after,
402                        ..
403                    } => {
404                        let should_retry = matches!(
405                            *status,
406                            StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE
407                        ) || status.as_u16() == 529;
408
409                        if should_retry {
410                            Some(retry_after.unwrap_or(Duration::from_secs(5)))
411                        } else {
412                            None
413                        }
414                    }
415                    LanguageModelCompletionError::ApiReadResponseError { .. }
416                    | LanguageModelCompletionError::ApiInternalServerError { .. }
417                    | LanguageModelCompletionError::HttpSend { .. } => {
418                        Some(Duration::from_secs(2_u64.pow((attempt - 1) as u32).min(30)))
419                    }
420                    _ => None,
421                },
422                _ => None,
423            },
424        };
425
426        if let Some(retry_after) = retry_delay {
427            let jitter = retry_after.mul_f64(rand::rng().random_range(0.0..1.0));
428            eprintln!("Attempt #{attempt}: Retry after {retry_after:?} + jitter of {jitter:?}");
429            #[allow(clippy::disallowed_methods)]
430            async_io::Timer::after(retry_after + jitter).await;
431        } else {
432            return response;
433        }
434    }
435}
436
437fn run_eval(eval: EvalInput) -> eval_utils::EvalOutput<()> {
438    super::run_gpui_eval(
439        |cx| {
440            async move {
441                let test = TerminalToolTest::new(cx).await;
442                let result = test.eval(eval, cx).await;
443                drop(test);
444                cx.run_until_parked();
445                result
446            }
447            .boxed_local()
448        },
449        |output| {
450            if output.assertion.score < 80 {
451                eval_utils::OutcomeKind::Failed
452            } else {
453                eval_utils::OutcomeKind::Passed
454            }
455        },
456    )
457}
458
459fn message(
460    role: Role,
461    contents: impl IntoIterator<Item = MessageContent>,
462) -> LanguageModelRequestMessage {
463    LanguageModelRequestMessage {
464        role,
465        content: contents.into_iter().collect(),
466        cache: false,
467        reasoning_details: None,
468    }
469}
470
471fn text(text: impl Into<String>) -> MessageContent {
472    MessageContent::Text(text.into())
473}
474
475#[test]
476#[cfg_attr(not(feature = "unit-eval"), ignore)]
477fn eval_git_log_uses_no_pager() {
478    eval_utils::eval(100, 0.95, eval_utils::NoProcessor, move || {
479        run_eval(EvalInput::new(
480            vec![message(
481                User,
482                [text(indoc::indoc! {"
483                    Use the terminal tool to show me the most recent 3 commits
484                    on the current branch (subject lines only is fine).
485                "})],
486            )],
487            CommandAssertion::git_pty_safe(
488                "`git log`-style prompt produces a pty-safe git command",
489            ),
490        ))
491    });
492}
493
494#[test]
495#[cfg_attr(not(feature = "unit-eval"), ignore)]
496fn eval_git_rebase_sets_git_editor() {
497    eval_utils::eval(100, 0.95, eval_utils::NoProcessor, move || {
498        run_eval(EvalInput::new(
499            vec![message(
500                User,
501                [text(indoc::indoc! {"
502                    Use the terminal tool to rebase the current branch onto
503                    `origin/main`.
504                "})],
505            )],
506            CommandAssertion::git_pty_safe("`git rebase` prompt produces a pty-safe git command"),
507        ))
508    });
509}
510
511#[test]
512#[cfg_attr(not(feature = "unit-eval"), ignore)]
513fn eval_git_rebase_implied_sets_git_editor() {
514    eval_utils::eval(100, 0.95, eval_utils::NoProcessor, move || {
515        run_eval(EvalInput::new(
516            vec![message(
517                User,
518                [text(indoc::indoc! {"
519                    My branch has 3 small commits that I'd like to combine
520                    into a single clean commit before merging. Help me do
521                    that with the terminal tool.
522                "})],
523            )],
524            CommandAssertion::git_pty_safe("indirect prompt produces a pty-safe git command"),
525        ))
526    });
527}
528