tenant.openagents/omega

1use super::*;
2use gpui::{App, SharedString, Task};
3use std::future;
4use std::sync::Mutex;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::time::Duration;
7
8/// A streaming tool that echoes its input, used to test streaming tool
9/// lifecycle (e.g. partial delivery and cleanup when the LLM stream ends
10/// before `is_input_complete`).
11#[derive(JsonSchema, Serialize, Deserialize)]
12pub struct StreamingEchoToolInput {
13    /// The text to echo.
14    pub text: String,
15}
16
17pub struct StreamingEchoTool {
18    wait_until_complete_rx: Mutex<Option<oneshot::Receiver<()>>>,
19}
20
21impl StreamingEchoTool {
22    pub fn new() -> Self {
23        Self {
24            wait_until_complete_rx: Mutex::new(None),
25        }
26    }
27
28    pub fn with_wait_until_complete(mut self, receiver: oneshot::Receiver<()>) -> Self {
29        self.wait_until_complete_rx = Mutex::new(Some(receiver));
30        self
31    }
32}
33
34impl AgentTool for StreamingEchoTool {
35    type Input = StreamingEchoToolInput;
36    type Output = String;
37
38    const NAME: &'static str = "streaming_echo";
39
40    fn supports_input_streaming() -> bool {
41        true
42    }
43
44    fn kind() -> acp::ToolKind {
45        acp::ToolKind::Other
46    }
47
48    fn initial_title(
49        &self,
50        _input: Result<Self::Input, serde_json::Value>,
51        _cx: &mut App,
52    ) -> SharedString {
53        "Streaming Echo".into()
54    }
55
56    fn run(
57        self: Arc<Self>,
58        input: ToolInput<Self::Input>,
59        _event_stream: ToolCallEventStream,
60        cx: &mut App,
61    ) -> Task<Result<String, String>> {
62        let wait_until_complete_rx = self.wait_until_complete_rx.lock().unwrap().take();
63        cx.spawn(async move |_cx| {
64            let input = input.recv().await.map_err(|e| e.to_string())?;
65            if let Some(rx) = wait_until_complete_rx {
66                rx.await.ok();
67            }
68            Ok(input.text)
69        })
70    }
71}
72
73#[derive(JsonSchema, Serialize, Deserialize)]
74pub struct StreamingJsonErrorContextToolInput {
75    /// The text to echo.
76    pub text: String,
77}
78
79pub struct StreamingJsonErrorContextTool;
80
81impl AgentTool for StreamingJsonErrorContextTool {
82    type Input = StreamingJsonErrorContextToolInput;
83    type Output = String;
84
85    const NAME: &'static str = "streaming_json_error_context";
86
87    fn supports_input_streaming() -> bool {
88        true
89    }
90
91    fn kind() -> acp::ToolKind {
92        acp::ToolKind::Other
93    }
94
95    fn initial_title(
96        &self,
97        _input: Result<Self::Input, serde_json::Value>,
98        _cx: &mut App,
99    ) -> SharedString {
100        "Streaming JSON Error Context".into()
101    }
102
103    fn run(
104        self: Arc<Self>,
105        mut input: ToolInput<Self::Input>,
106        _event_stream: ToolCallEventStream,
107        cx: &mut App,
108    ) -> Task<Result<String, String>> {
109        cx.spawn(async move |_cx| {
110            let mut last_partial_text = None;
111
112            loop {
113                match input.next().await {
114                    Ok(ToolInputPayload::Partial(partial)) => {
115                        if let Some(text) = partial.get("text").and_then(|value| value.as_str()) {
116                            last_partial_text = Some(text.to_string());
117                        }
118                    }
119                    Ok(ToolInputPayload::Full(input)) => return Ok(input.text),
120                    Ok(ToolInputPayload::InvalidJson { error_message }) => {
121                        let partial_text = last_partial_text.unwrap_or_default();
122                        return Err(format!(
123                            "Saw partial text '{partial_text}' before invalid JSON: {error_message}"
124                        ));
125                    }
126                    Err(error) => {
127                        return Err(error.to_string());
128                    }
129                }
130            }
131        })
132    }
133}
134
135/// A streaming tool that echoes its input, used to test streaming tool
136/// lifecycle (e.g. partial delivery and cleanup when the LLM stream ends
137/// before `is_input_complete`).
138#[derive(JsonSchema, Serialize, Deserialize)]
139pub struct StreamingFailingEchoToolInput {
140    /// The text to echo.
141    pub text: String,
142}
143
144pub struct StreamingFailingEchoTool {
145    pub receive_chunks_until_failure: usize,
146}
147
148impl AgentTool for StreamingFailingEchoTool {
149    type Input = StreamingFailingEchoToolInput;
150
151    type Output = String;
152
153    const NAME: &'static str = "streaming_failing_echo";
154
155    fn kind() -> acp::ToolKind {
156        acp::ToolKind::Other
157    }
158
159    fn supports_input_streaming() -> bool {
160        true
161    }
162
163    fn initial_title(
164        &self,
165        _input: Result<Self::Input, serde_json::Value>,
166        _cx: &mut App,
167    ) -> SharedString {
168        "echo".into()
169    }
170
171    fn run(
172        self: Arc<Self>,
173        mut input: ToolInput<Self::Input>,
174        _event_stream: ToolCallEventStream,
175        cx: &mut App,
176    ) -> Task<Result<Self::Output, Self::Output>> {
177        cx.spawn(async move |_cx| {
178            for _ in 0..self.receive_chunks_until_failure {
179                let _ = input.next().await;
180            }
181            Err("failed".into())
182        })
183    }
184}
185
186/// A tool that echoes its input
187#[derive(JsonSchema, Serialize, Deserialize)]
188pub struct EchoToolInput {
189    /// The text to echo.
190    pub text: String,
191}
192
193pub struct EchoTool;
194
195impl AgentTool for EchoTool {
196    type Input = EchoToolInput;
197    type Output = String;
198
199    const NAME: &'static str = "echo";
200
201    fn kind() -> acp::ToolKind {
202        acp::ToolKind::Other
203    }
204
205    fn initial_title(
206        &self,
207        _input: Result<Self::Input, serde_json::Value>,
208        _cx: &mut App,
209    ) -> SharedString {
210        "Echo".into()
211    }
212
213    fn run(
214        self: Arc<Self>,
215        input: ToolInput<Self::Input>,
216        _event_stream: ToolCallEventStream,
217        cx: &mut App,
218    ) -> Task<Result<String, String>> {
219        cx.spawn(async move |_cx| {
220            let input = input.recv().await.map_err(|e| e.to_string())?;
221            Ok(input.text)
222        })
223    }
224}
225
226/// A tool that waits for a specified delay
227#[derive(JsonSchema, Serialize, Deserialize)]
228pub struct DelayToolInput {
229    /// The delay in milliseconds.
230    ms: u64,
231}
232
233pub struct DelayTool;
234
235impl AgentTool for DelayTool {
236    type Input = DelayToolInput;
237    type Output = String;
238
239    const NAME: &'static str = "delay";
240
241    fn initial_title(
242        &self,
243        input: Result<Self::Input, serde_json::Value>,
244        _cx: &mut App,
245    ) -> SharedString {
246        if let Ok(input) = input {
247            format!("Delay {}ms", input.ms).into()
248        } else {
249            "Delay".into()
250        }
251    }
252
253    fn kind() -> acp::ToolKind {
254        acp::ToolKind::Other
255    }
256
257    fn run(
258        self: Arc<Self>,
259        input: ToolInput<Self::Input>,
260        _event_stream: ToolCallEventStream,
261        cx: &mut App,
262    ) -> Task<Result<String, String>>
263    where
264        Self: Sized,
265    {
266        let executor = cx.background_executor().clone();
267        cx.foreground_executor().spawn(async move {
268            let input = input.recv().await.map_err(|e| e.to_string())?;
269            executor.timer(Duration::from_millis(input.ms)).await;
270            Ok("Ding".to_string())
271        })
272    }
273}
274
275#[derive(JsonSchema, Serialize, Deserialize)]
276pub struct ToolRequiringPermissionInput {}
277
278pub struct ToolRequiringPermission;
279
280impl AgentTool for ToolRequiringPermission {
281    type Input = ToolRequiringPermissionInput;
282    type Output = String;
283
284    const NAME: &'static str = "tool_requiring_permission";
285
286    fn kind() -> acp::ToolKind {
287        acp::ToolKind::Other
288    }
289
290    fn initial_title(
291        &self,
292        _input: Result<Self::Input, serde_json::Value>,
293        _cx: &mut App,
294    ) -> SharedString {
295        "This tool requires permission".into()
296    }
297
298    fn run(
299        self: Arc<Self>,
300        input: ToolInput<Self::Input>,
301        event_stream: ToolCallEventStream,
302        cx: &mut App,
303    ) -> Task<Result<String, String>> {
304        cx.spawn(async move |cx| {
305            let _input = input.recv().await.map_err(|e| e.to_string())?;
306
307            let authorize = cx.update(|cx| {
308                let context = crate::ToolPermissionContext::new(Self::NAME, vec![String::new()]);
309                event_stream.authorize("Authorize?", context, cx)
310            });
311            authorize.await.map_err(|e| e.to_string())?;
312            Ok("Allowed".to_string())
313        })
314    }
315}
316
317/// A second tool that also requires permission, used to verify that
318/// permission decisions scoped to one tool don't leak into prompts for a
319/// different tool.
320#[derive(JsonSchema, Serialize, Deserialize)]
321pub struct ToolRequiringPermission2Input {}
322
323pub struct ToolRequiringPermission2;
324
325impl AgentTool for ToolRequiringPermission2 {
326    type Input = ToolRequiringPermission2Input;
327    type Output = String;
328
329    const NAME: &'static str = "tool_requiring_permission_2";
330
331    fn kind() -> acp::ToolKind {
332        acp::ToolKind::Other
333    }
334
335    fn initial_title(
336        &self,
337        _input: Result<Self::Input, serde_json::Value>,
338        _cx: &mut App,
339    ) -> SharedString {
340        "This tool also requires permission".into()
341    }
342
343    fn run(
344        self: Arc<Self>,
345        input: ToolInput<Self::Input>,
346        event_stream: ToolCallEventStream,
347        cx: &mut App,
348    ) -> Task<Result<String, String>> {
349        cx.spawn(async move |cx| {
350            let _input = input.recv().await.map_err(|e| e.to_string())?;
351
352            let authorize = cx.update(|cx| {
353                let context = crate::ToolPermissionContext::new(Self::NAME, vec![String::new()]);
354                event_stream.authorize("Authorize?", context, cx)
355            });
356            authorize.await.map_err(|e| e.to_string())?;
357            Ok("Allowed".to_string())
358        })
359    }
360}
361
362#[derive(JsonSchema, Serialize, Deserialize)]
363pub struct InfiniteToolInput {}
364
365pub struct InfiniteTool;
366
367impl AgentTool for InfiniteTool {
368    type Input = InfiniteToolInput;
369    type Output = String;
370
371    const NAME: &'static str = "infinite";
372
373    fn kind() -> acp::ToolKind {
374        acp::ToolKind::Other
375    }
376
377    fn initial_title(
378        &self,
379        _input: Result<Self::Input, serde_json::Value>,
380        _cx: &mut App,
381    ) -> SharedString {
382        "Infinite Tool".into()
383    }
384
385    fn run(
386        self: Arc<Self>,
387        input: ToolInput<Self::Input>,
388        _event_stream: ToolCallEventStream,
389        cx: &mut App,
390    ) -> Task<Result<String, String>> {
391        cx.foreground_executor().spawn(async move {
392            let _input = input.recv().await.map_err(|e| e.to_string())?;
393            future::pending::<()>().await;
394            unreachable!()
395        })
396    }
397}
398
399/// A tool that loops forever but properly handles cancellation via `select!`,
400/// similar to how edit_file_tool handles cancellation.
401#[derive(JsonSchema, Serialize, Deserialize)]
402pub struct CancellationAwareToolInput {}
403
404pub struct CancellationAwareTool {
405    pub was_cancelled: Arc<AtomicBool>,
406}
407
408impl CancellationAwareTool {
409    pub fn new() -> (Self, Arc<AtomicBool>) {
410        let was_cancelled = Arc::new(AtomicBool::new(false));
411        (
412            Self {
413                was_cancelled: was_cancelled.clone(),
414            },
415            was_cancelled,
416        )
417    }
418}
419
420impl AgentTool for CancellationAwareTool {
421    type Input = CancellationAwareToolInput;
422    type Output = String;
423
424    const NAME: &'static str = "cancellation_aware";
425
426    fn kind() -> acp::ToolKind {
427        acp::ToolKind::Other
428    }
429
430    fn initial_title(
431        &self,
432        _input: Result<Self::Input, serde_json::Value>,
433        _cx: &mut App,
434    ) -> SharedString {
435        "Cancellation Aware Tool".into()
436    }
437
438    fn run(
439        self: Arc<Self>,
440        input: ToolInput<Self::Input>,
441        event_stream: ToolCallEventStream,
442        cx: &mut App,
443    ) -> Task<Result<String, String>> {
444        cx.foreground_executor().spawn(async move {
445            let _input = input.recv().await.map_err(|e| e.to_string())?;
446            // Wait for cancellation - this tool does nothing but wait to be cancelled
447            event_stream.cancelled_by_user().await;
448            self.was_cancelled.store(true, Ordering::SeqCst);
449            Err("Tool cancelled by user".to_string())
450        })
451    }
452}
453
454/// A tool that takes an object with map from letters to random words starting with that letter.
455/// All fiealds are required! Pass a word for every letter!
456#[derive(JsonSchema, Serialize, Deserialize)]
457pub struct WordListInput {
458    /// Provide a random word that starts with A.
459    a: Option<String>,
460    /// Provide a random word that starts with B.
461    b: Option<String>,
462    /// Provide a random word that starts with C.
463    c: Option<String>,
464    /// Provide a random word that starts with D.
465    d: Option<String>,
466    /// Provide a random word that starts with E.
467    e: Option<String>,
468    /// Provide a random word that starts with F.
469    f: Option<String>,
470    /// Provide a random word that starts with G.
471    g: Option<String>,
472}
473
474pub struct WordListTool;
475
476impl AgentTool for WordListTool {
477    type Input = WordListInput;
478    type Output = String;
479
480    const NAME: &'static str = "word_list";
481
482    fn kind() -> acp::ToolKind {
483        acp::ToolKind::Other
484    }
485
486    fn initial_title(
487        &self,
488        _input: Result<Self::Input, serde_json::Value>,
489        _cx: &mut App,
490    ) -> SharedString {
491        "List of random words".into()
492    }
493
494    fn run(
495        self: Arc<Self>,
496        input: ToolInput<Self::Input>,
497        _event_stream: ToolCallEventStream,
498        cx: &mut App,
499    ) -> Task<Result<String, String>> {
500        cx.spawn(async move |_cx| {
501            let _input = input.recv().await.map_err(|e| e.to_string())?;
502            Ok("ok".to_string())
503        })
504    }
505}
506