test_tools.rs

  1use super::*;
  2use agent_settings::AgentSettings;
  3use gpui::{App, SharedString, Task};
  4use std::future;
  5use std::sync::atomic::{AtomicBool, Ordering};
  6
  7/// A tool that echoes its input
  8#[derive(JsonSchema, Serialize, Deserialize)]
  9pub struct EchoToolInput {
 10    /// The text to echo.
 11    pub text: String,
 12}
 13
 14pub struct EchoTool;
 15
 16impl AgentTool for EchoTool {
 17    type Input = EchoToolInput;
 18    type Output = String;
 19
 20    const NAME: &'static str = "echo";
 21
 22    fn kind() -> acp::ToolKind {
 23        acp::ToolKind::Other
 24    }
 25
 26    fn initial_title(
 27        &self,
 28        _input: Result<Self::Input, serde_json::Value>,
 29        _cx: &mut App,
 30    ) -> SharedString {
 31        "Echo".into()
 32    }
 33
 34    fn run(
 35        self: Arc<Self>,
 36        input: Self::Input,
 37        _event_stream: ToolCallEventStream,
 38        _cx: &mut App,
 39    ) -> Task<Result<String, String>> {
 40        Task::ready(Ok(input.text))
 41    }
 42}
 43
 44/// A tool that waits for a specified delay
 45#[derive(JsonSchema, Serialize, Deserialize)]
 46pub struct DelayToolInput {
 47    /// The delay in milliseconds.
 48    ms: u64,
 49}
 50
 51pub struct DelayTool;
 52
 53impl AgentTool for DelayTool {
 54    type Input = DelayToolInput;
 55    type Output = String;
 56
 57    const NAME: &'static str = "delay";
 58
 59    fn initial_title(
 60        &self,
 61        input: Result<Self::Input, serde_json::Value>,
 62        _cx: &mut App,
 63    ) -> SharedString {
 64        if let Ok(input) = input {
 65            format!("Delay {}ms", input.ms).into()
 66        } else {
 67            "Delay".into()
 68        }
 69    }
 70
 71    fn kind() -> acp::ToolKind {
 72        acp::ToolKind::Other
 73    }
 74
 75    fn run(
 76        self: Arc<Self>,
 77        input: Self::Input,
 78        _event_stream: ToolCallEventStream,
 79        cx: &mut App,
 80    ) -> Task<Result<String, String>>
 81    where
 82        Self: Sized,
 83    {
 84        let executor = cx.background_executor().clone();
 85        cx.foreground_executor().spawn(async move {
 86            executor.timer(Duration::from_millis(input.ms)).await;
 87            Ok("Ding".to_string())
 88        })
 89    }
 90}
 91
 92#[derive(JsonSchema, Serialize, Deserialize)]
 93pub struct ToolRequiringPermissionInput {}
 94
 95pub struct ToolRequiringPermission;
 96
 97impl AgentTool for ToolRequiringPermission {
 98    type Input = ToolRequiringPermissionInput;
 99    type Output = String;
100
101    const NAME: &'static str = "tool_requiring_permission";
102
103    fn kind() -> acp::ToolKind {
104        acp::ToolKind::Other
105    }
106
107    fn initial_title(
108        &self,
109        _input: Result<Self::Input, serde_json::Value>,
110        _cx: &mut App,
111    ) -> SharedString {
112        "This tool requires permission".into()
113    }
114
115    fn run(
116        self: Arc<Self>,
117        _input: Self::Input,
118        event_stream: ToolCallEventStream,
119        cx: &mut App,
120    ) -> Task<Result<String, String>> {
121        let settings = AgentSettings::get_global(cx);
122        let decision = decide_permission_from_settings(Self::NAME, &[String::new()], settings);
123
124        let authorize = match decision {
125            ToolPermissionDecision::Allow => None,
126            ToolPermissionDecision::Deny(reason) => {
127                return Task::ready(Err(reason));
128            }
129            ToolPermissionDecision::Confirm => {
130                let context = crate::ToolPermissionContext::new(
131                    "tool_requiring_permission",
132                    vec![String::new()],
133                );
134                Some(event_stream.authorize("Authorize?", context, cx))
135            }
136        };
137
138        cx.foreground_executor().spawn(async move {
139            if let Some(authorize) = authorize {
140                authorize.await.map_err(|e| e.to_string())?;
141            }
142            Ok("Allowed".to_string())
143        })
144    }
145}
146
147#[derive(JsonSchema, Serialize, Deserialize)]
148pub struct InfiniteToolInput {}
149
150pub struct InfiniteTool;
151
152impl AgentTool for InfiniteTool {
153    type Input = InfiniteToolInput;
154    type Output = String;
155
156    const NAME: &'static str = "infinite";
157
158    fn kind() -> acp::ToolKind {
159        acp::ToolKind::Other
160    }
161
162    fn initial_title(
163        &self,
164        _input: Result<Self::Input, serde_json::Value>,
165        _cx: &mut App,
166    ) -> SharedString {
167        "Infinite Tool".into()
168    }
169
170    fn run(
171        self: Arc<Self>,
172        _input: Self::Input,
173        _event_stream: ToolCallEventStream,
174        cx: &mut App,
175    ) -> Task<Result<String, String>> {
176        cx.foreground_executor().spawn(async move {
177            future::pending::<()>().await;
178            unreachable!()
179        })
180    }
181}
182
183/// A tool that loops forever but properly handles cancellation via `select!`,
184/// similar to how edit_file_tool handles cancellation.
185#[derive(JsonSchema, Serialize, Deserialize)]
186pub struct CancellationAwareToolInput {}
187
188pub struct CancellationAwareTool {
189    pub was_cancelled: Arc<AtomicBool>,
190}
191
192impl CancellationAwareTool {
193    pub fn new() -> (Self, Arc<AtomicBool>) {
194        let was_cancelled = Arc::new(AtomicBool::new(false));
195        (
196            Self {
197                was_cancelled: was_cancelled.clone(),
198            },
199            was_cancelled,
200        )
201    }
202}
203
204impl AgentTool for CancellationAwareTool {
205    type Input = CancellationAwareToolInput;
206    type Output = String;
207
208    const NAME: &'static str = "cancellation_aware";
209
210    fn kind() -> acp::ToolKind {
211        acp::ToolKind::Other
212    }
213
214    fn initial_title(
215        &self,
216        _input: Result<Self::Input, serde_json::Value>,
217        _cx: &mut App,
218    ) -> SharedString {
219        "Cancellation Aware Tool".into()
220    }
221
222    fn run(
223        self: Arc<Self>,
224        _input: Self::Input,
225        event_stream: ToolCallEventStream,
226        cx: &mut App,
227    ) -> Task<Result<String, String>> {
228        cx.foreground_executor().spawn(async move {
229            // Wait for cancellation - this tool does nothing but wait to be cancelled
230            event_stream.cancelled_by_user().await;
231            self.was_cancelled.store(true, Ordering::SeqCst);
232            Err("Tool cancelled by user".to_string())
233        })
234    }
235}
236
237/// A tool that takes an object with map from letters to random words starting with that letter.
238/// All fiealds are required! Pass a word for every letter!
239#[derive(JsonSchema, Serialize, Deserialize)]
240pub struct WordListInput {
241    /// Provide a random word that starts with A.
242    a: Option<String>,
243    /// Provide a random word that starts with B.
244    b: Option<String>,
245    /// Provide a random word that starts with C.
246    c: Option<String>,
247    /// Provide a random word that starts with D.
248    d: Option<String>,
249    /// Provide a random word that starts with E.
250    e: Option<String>,
251    /// Provide a random word that starts with F.
252    f: Option<String>,
253    /// Provide a random word that starts with G.
254    g: Option<String>,
255}
256
257pub struct WordListTool;
258
259impl AgentTool for WordListTool {
260    type Input = WordListInput;
261    type Output = String;
262
263    const NAME: &'static str = "word_list";
264
265    fn kind() -> acp::ToolKind {
266        acp::ToolKind::Other
267    }
268
269    fn initial_title(
270        &self,
271        _input: Result<Self::Input, serde_json::Value>,
272        _cx: &mut App,
273    ) -> SharedString {
274        "List of random words".into()
275    }
276
277    fn run(
278        self: Arc<Self>,
279        _input: Self::Input,
280        _event_stream: ToolCallEventStream,
281        _cx: &mut App,
282    ) -> Task<Result<String, String>> {
283        Task::ready(Ok("ok".to_string()))
284    }
285}