tenant.openagents/omega

1use crate::{AgentTool, ToolCallEventStream, ToolInput};
2use agent_client_protocol::schema::v1 as acp;
3use futures::{Future, FutureExt as _};
4use gpui::{App, AsyncApp, Entity, Task};
5use language::{DiagnosticSeverity, OffsetRangeExt};
6use project::Project;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9use std::path::Path;
10use std::{fmt::Write, sync::Arc};
11use ui::SharedString;
12use util::markdown::MarkdownInlineCode;
13
14type Result<T, E = String> = core::result::Result<T, E>;
15
16/// Get errors and warnings for the project or a specific file.
17///
18/// This tool can be invoked after a series of edits to determine if further edits are necessary, or if the user asks to fix errors or warnings in their codebase.
19///
20/// When a path is provided, shows all diagnostics for that specific file.
21/// When no path is provided, shows a summary of error and warning counts for all files in the project.
22///
23/// This tool attempts to refresh diagnostics before returning.
24/// If refreshing diagnostics fails (for example, if the language server does not support pull-based diagnostics), it will return any diagnostics already present.
25/// Note that, in this case, the results may be out-of-date, and may or may not reflect the most recent edits.
26/// If this happens, do not attempt to re-run this tool in the hope that refreshing will later succeed. Failures are typically persistent.
27///
28/// <example>
29/// To get diagnostics for a specific file:
30/// {
31///     "path": "src/main.rs"
32/// }
33///
34/// To get a project-wide diagnostic summary:
35/// {}
36/// </example>
37///
38/// <guidelines>
39/// - If you think you can fix a diagnostic, make 1-2 attempts and then give up.
40/// - Don't remove code you've generated just because you can't fix an error. The user can help you fix it.
41/// </guidelines>
42#[derive(Debug, Serialize, Deserialize, JsonSchema)]
43pub struct DiagnosticsToolInput {
44    /// The path to get diagnostics for. If not provided, returns a project-wide summary.
45    ///
46    /// This path should never be absolute, and the first component
47    /// of the path should always be a root directory in a project.
48    ///
49    /// <example>
50    /// If the project has the following root directories:
51    ///
52    /// - lorem
53    /// - ipsum
54    ///
55    /// If you wanna access diagnostics for `dolor.txt` in `ipsum`, you should use the path `ipsum/dolor.txt`.
56    /// </example>
57    pub path: Option<String>,
58}
59
60pub struct DiagnosticsTool {
61    project: Entity<Project>,
62}
63
64impl DiagnosticsTool {
65    pub fn new(project: Entity<Project>) -> Self {
66        Self { project }
67    }
68}
69
70async fn with_cancellation<T>(f: impl Future<Output = T>, s: &ToolCallEventStream) -> Result<T> {
71    futures::select! {
72        result = f.fuse() => Ok(result),
73        _ = s.cancelled_by_user().fuse() => {
74            Err("Diagnostics cancelled by user".to_string())
75        }
76    }
77}
78
79fn freshness_message(refreshed: bool) -> &'static str {
80    if refreshed {
81        "Diagnostics successfully refreshed."
82    } else {
83        "Failed to refresh diagnostics. Diagnostics may be stale."
84    }
85}
86
87/// Attempt to pull fresh diagnostics from the LSP before reading them.
88///
89/// Returns `Ok(true)` if diagnostics were successfully refreshed,
90/// `Ok(false)` if the pull failed (callers should fall through to
91/// read cached diagnostics), or `Err` if cancelled by the user.
92async fn pull_diagnostics(
93    project: &Entity<Project>,
94    path: Option<&Path>,
95    event_stream: &ToolCallEventStream,
96    cx: &mut AsyncApp,
97) -> Result<bool, String> {
98    match path {
99        Some(path) => {
100            let open_buffer_task = project.update(cx, |project, cx| {
101                let Some(project_path) = project.find_project_path(path, cx) else {
102                    return Err(format!("Could not find path {} in project", path.display()));
103                };
104                Ok(project.open_buffer(project_path, cx))
105            })?;
106
107            let buffer = with_cancellation(open_buffer_task, event_stream)
108                .await?
109                .map_err(|e| e.to_string())?;
110
111            let lsp_store = project.read_with(cx, |project, _cx| project.lsp_store());
112            let pull_task = lsp_store.update(cx, |lsp_store, cx| {
113                lsp_store.pull_diagnostics_for_buffer(buffer, cx)
114            });
115            let pull_result = with_cancellation(pull_task, event_stream).await?;
116            if let Err(error) = &pull_result {
117                log::warn!("Failed to pull diagnostics, using cached: {error:#}");
118            }
119            Ok(pull_result.is_ok())
120        }
121        None => {
122            let lsp_store = project.read_with(cx, |project, _cx| project.lsp_store());
123            let pull_task = lsp_store.update(cx, |lsp_store, cx| {
124                lsp_store.pull_workspace_diagnostics_once(cx)
125            });
126            let succeeded = with_cancellation(pull_task, event_stream).await?;
127            if !succeeded {
128                log::warn!("Failed to pull workspace diagnostics, using cached");
129            }
130            Ok(succeeded)
131        }
132    }
133}
134
135impl AgentTool for DiagnosticsTool {
136    type Input = DiagnosticsToolInput;
137    type Output = String;
138
139    const NAME: &'static str = "diagnostics";
140
141    fn kind() -> acp::ToolKind {
142        acp::ToolKind::Read
143    }
144
145    fn initial_title(
146        &self,
147        input: Result<Self::Input, serde_json::Value>,
148        _cx: &mut App,
149    ) -> SharedString {
150        if let Some(path) = input.ok().and_then(|input| match input.path {
151            Some(path) if !path.is_empty() => Some(path),
152            _ => None,
153        }) {
154            format!("Check diagnostics for {}", MarkdownInlineCode(&path)).into()
155        } else {
156            "Check project diagnostics".into()
157        }
158    }
159
160    fn run(
161        self: Arc<Self>,
162        input: ToolInput<Self::Input>,
163        event_stream: ToolCallEventStream,
164        cx: &mut App,
165    ) -> Task<Result<Self::Output, Self::Output>> {
166        let project = self.project.clone();
167        cx.spawn(async move |cx| {
168            let input = input.recv().await.map_err(|e| e.to_string())?;
169
170            match input.path {
171                Some(ref path) if !path.is_empty() => {
172                    let refreshed =
173                        pull_diagnostics(&project, Some(Path::new(path)), &event_stream, cx)
174                            .await?;
175
176                    let open_buffer_task = project.update(cx, |project, cx| {
177                        let Some(project_path) = project.find_project_path(path, cx) else {
178                            return Err(format!("Could not find path {path} in project"));
179                        };
180                        Ok(project.open_buffer(project_path, cx))
181                    })?;
182
183                    let buffer = with_cancellation(open_buffer_task, &event_stream)
184                        .await?
185                        .map_err(|e| e.to_string())?;
186
187                    let mut output = String::new();
188                    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
189
190                    for (_, group) in snapshot.diagnostic_groups(None) {
191                        let entry = &group.entries[group.primary_ix];
192                        let range = entry.range.to_point(&snapshot);
193                        let severity = match entry.diagnostic.severity {
194                            DiagnosticSeverity::ERROR => "error",
195                            DiagnosticSeverity::WARNING => "warning",
196                            _ => continue,
197                        };
198
199                        writeln!(
200                            output,
201                            "{} at line {}: {}",
202                            severity,
203                            range.start.row + 1,
204                            entry.diagnostic.message
205                        )
206                        .ok();
207                    }
208
209                    let freshness = freshness_message(refreshed);
210                    if output.is_empty() {
211                        Ok(format!(
212                            "{freshness}\n\nFile doesn't have errors or warnings!"
213                        ))
214                    } else {
215                        Ok(format!("{freshness}\n\n{output}"))
216                    }
217                }
218                _ => {
219                    let refreshed = pull_diagnostics(&project, None, &event_stream, cx).await?;
220
221                    let (output, has_diagnostics) = project.read_with(cx, |project, cx| {
222                        let mut output = String::new();
223                        let mut has_diagnostics = false;
224
225                        for (project_path, _, summary) in project.diagnostic_summaries(true, cx) {
226                            if summary.error_count > 0 || summary.warning_count > 0 {
227                                let Some(worktree) =
228                                    project.worktree_for_id(project_path.worktree_id, cx)
229                                else {
230                                    continue;
231                                };
232
233                                has_diagnostics = true;
234                                output.push_str(&format!(
235                                    "{}: {} error(s), {} warning(s)\n",
236                                    worktree.read(cx).absolutize(&project_path.path).display(),
237                                    summary.error_count,
238                                    summary.warning_count
239                                ));
240                            }
241                        }
242
243                        (output, has_diagnostics)
244                    });
245
246                    let freshness = freshness_message(refreshed);
247                    if has_diagnostics {
248                        Ok(format!("{freshness}\n\n{output}"))
249                    } else {
250                        Ok(format!(
251                            "{freshness}\n\nNo errors or warnings found in the project."
252                        ))
253                    }
254                }
255            }
256        })
257    }
258}
259