tenant.openagents/omega

1use std::fmt::Write;
2use std::sync::Arc;
3
4use super::symbol_locator::{LocationDisplay, SymbolLocator};
5use crate::{AgentTool, ToolCallEventStream, ToolInput};
6use agent_client_protocol::schema::v1 as acp;
7use gpui::{App, Entity, SharedString, Task};
8use project::Project;
9use schemars::JsonSchema;
10use serde::{Deserialize, Serialize};
11
12/// Finds all references to a symbol across the project using the language server.
13///
14/// Returns a list of locations where the symbol is referenced, including file paths, line numbers, and code snippets for each reference.
15///
16/// Before using this tool, use read_file or grep to find the exact symbol name and line number.
17#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
18pub struct FindReferencesToolInput {
19    /// The symbol to find references of.
20    pub symbol: SymbolLocator,
21}
22
23pub struct FindReferencesTool {
24    project: Entity<Project>,
25}
26
27impl FindReferencesTool {
28    pub fn new(project: Entity<Project>) -> Self {
29        Self { project }
30    }
31}
32
33impl AgentTool for FindReferencesTool {
34    type Input = FindReferencesToolInput;
35    type Output = String;
36
37    const NAME: &'static str = "find_references";
38
39    fn kind() -> acp::ToolKind {
40        acp::ToolKind::Search
41    }
42
43    fn initial_title(
44        &self,
45        input: Result<Self::Input, serde_json::Value>,
46        _cx: &mut App,
47    ) -> SharedString {
48        if let Ok(input) = input {
49            format!("Find references to `{}`", input.symbol.symbol_name).into()
50        } else {
51            "Find references".into()
52        }
53    }
54
55    fn run(
56        self: Arc<Self>,
57        input: ToolInput<Self::Input>,
58        _event_stream: ToolCallEventStream,
59        cx: &mut App,
60    ) -> Task<Result<String, String>> {
61        let project = self.project.clone();
62        cx.spawn(async move |cx| {
63            let input = input
64                .recv()
65                .await
66                .map_err(|e| format!("Failed to receive tool input: {e}"))?;
67
68            let resolved = input.symbol.resolve(&project, cx).await?;
69
70            let references_task = project.update(cx, |project, cx| {
71                project.references(&resolved.buffer, resolved.position, cx)
72            });
73
74            let references = references_task
75                .await
76                .map_err(|e| format!("Find references failed: {e}"))?
77                .unwrap_or_default();
78
79            if references.is_empty() {
80                return Ok(format!(
81                    "No references found for '{}'.",
82                    input.symbol.symbol_name
83                ));
84            }
85
86            let mut output = format!(
87                "Found {} references to `{}`:\n",
88                references.len(),
89                input.symbol.symbol_name
90            );
91
92            for location in &references {
93                let display = location
94                    .buffer
95                    .read_with(cx, |_, cx| LocationDisplay::from_location(location, cx));
96                write!(output, "\n## {display}\n").ok();
97            }
98
99            Ok(output)
100        })
101    }
102}
103