Skip to main content

distri_types/
ui_tool_renderers.rs

1use anyhow::Result;
2use std::sync::Arc;
3
4use crate::{Part, ToolUiContext, ToolUiMessage, ToolUiMessageType, UiToolRender};
5
6/// UI renderer for search/grep tools
7#[derive(Debug)]
8pub struct SearchToolRenderer;
9
10impl UiToolRender for SearchToolRenderer {
11    fn get_tool_name(&self) -> String {
12        "search".to_string()
13    }
14
15    fn render_tool_start(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
16        let query = context
17            .tool_call
18            .input
19            .get("query")
20            .and_then(|q| q.as_str())
21            .unwrap_or("unknown");
22
23        let message = format!("🔍 **Searching for:** `{}`", query);
24
25        Ok(ToolUiMessage {
26            message_type: ToolUiMessageType::ToolStart,
27            parts: vec![Part::Text(message)],
28        })
29    }
30
31    fn render_tool_end(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
32        let tool_response = context
33            .tool_response
34            .as_ref()
35            .ok_or_else(|| anyhow::anyhow!("Tool response required for tool_end message"))?;
36
37        let mut parts = vec![Part::Text("✅ **Search completed**".to_string())];
38        parts.extend(tool_response.parts.clone());
39
40        Ok(ToolUiMessage {
41            message_type: ToolUiMessageType::ToolEnd,
42            parts,
43        })
44    }
45
46    fn render_tool_error(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
47        let error_msg = context
48            .error
49            .as_ref()
50            .cloned()
51            .unwrap_or_else(|| "Search failed".to_string());
52
53        let message = format!("❌ **Search failed**\n\n```\n{}\n```", error_msg);
54
55        Ok(ToolUiMessage {
56            message_type: ToolUiMessageType::ToolError,
57            parts: vec![Part::Text(message)],
58        })
59    }
60}
61
62/// UI renderer for file operations (read, write, etc.)
63#[derive(Debug)]
64pub struct FileToolRenderer;
65
66impl UiToolRender for FileToolRenderer {
67    fn get_tool_name(&self) -> String {
68        "file_read".to_string() // Can be used for multiple file tools
69    }
70
71    fn render_tool_start(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
72        let tool_name = &context.tool_call.tool_name;
73        let path = context
74            .tool_call
75            .input
76            .get("path")
77            .and_then(|p| p.as_str())
78            .unwrap_or("unknown");
79
80        let action = match tool_name.as_str() {
81            name if name.contains("read") => "📖 Reading",
82            name if name.contains("write") => "✏️ Writing",
83            name if name.contains("delete") => "🗑️ Deleting",
84            name if name.contains("copy") => "📋 Copying",
85            name if name.contains("move") => "🔄 Moving",
86            _ => "📁 Processing",
87        };
88
89        let message = format!("{} **{}**", action, path);
90
91        Ok(ToolUiMessage {
92            message_type: ToolUiMessageType::ToolStart,
93            parts: vec![Part::Text(message)],
94        })
95    }
96
97    fn render_tool_end(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
98        let tool_response = context
99            .tool_response
100            .as_ref()
101            .ok_or_else(|| anyhow::anyhow!("Tool response required for tool_end message"))?;
102
103        let tool_name = &context.tool_call.tool_name;
104        let action = match tool_name.as_str() {
105            name if name.contains("read") => "Read",
106            name if name.contains("write") => "Wrote",
107            name if name.contains("delete") => "Deleted",
108            name if name.contains("copy") => "Copied",
109            name if name.contains("move") => "Moved",
110            _ => "Processed",
111        };
112
113        let mut parts = vec![Part::Text(format!("✅ **{} completed**", action))];
114        parts.extend(tool_response.parts.clone());
115
116        Ok(ToolUiMessage {
117            message_type: ToolUiMessageType::ToolEnd,
118            parts,
119        })
120    }
121
122    fn render_tool_error(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
123        let error_msg = context
124            .error
125            .as_ref()
126            .cloned()
127            .unwrap_or_else(|| "File operation failed".to_string());
128
129        let message = format!("❌ **File operation failed**\n\n```\n{}\n```", error_msg);
130
131        Ok(ToolUiMessage {
132            message_type: ToolUiMessageType::ToolError,
133            parts: vec![Part::Text(message)],
134        })
135    }
136}
137
138/// UI renderer for code execution tools
139#[derive(Debug)]
140pub struct CodeExecutionToolRenderer;
141
142impl UiToolRender for CodeExecutionToolRenderer {
143    fn get_tool_name(&self) -> String {
144        "distri_execute_code".to_string()
145    }
146
147    fn render_tool_start(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
148        let code = context
149            .tool_call
150            .input
151            .get("code")
152            .and_then(|c| c.as_str())
153            .unwrap_or("");
154
155        let language = context
156            .tool_call
157            .input
158            .get("language")
159            .and_then(|l| l.as_str())
160            .unwrap_or("javascript");
161
162        let preview = if code.len() > 100 {
163            format!("{}...", &code[..100])
164        } else {
165            code.to_string()
166        };
167
168        let message = format!(
169            "⚡ **Executing {} code**\n\n```{}\n{}\n```",
170            language, language, preview
171        );
172
173        Ok(ToolUiMessage {
174            message_type: ToolUiMessageType::ToolStart,
175            parts: vec![Part::Text(message)],
176        })
177    }
178
179    fn render_tool_end(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
180        let tool_response = context
181            .tool_response
182            .as_ref()
183            .ok_or_else(|| anyhow::anyhow!("Tool response required for tool_end message"))?;
184
185        let mut parts = vec![Part::Text("✅ **Code execution completed**".to_string())];
186        parts.extend(tool_response.parts.clone());
187
188        Ok(ToolUiMessage {
189            message_type: ToolUiMessageType::ToolEnd,
190            parts,
191        })
192    }
193
194    fn render_tool_error(&self, context: &ToolUiContext) -> Result<ToolUiMessage> {
195        let error_msg = context
196            .error
197            .as_ref()
198            .cloned()
199            .unwrap_or_else(|| "Code execution failed".to_string());
200
201        let message = format!("❌ **Code execution failed**\n\n```\n{}\n```", error_msg);
202
203        Ok(ToolUiMessage {
204            message_type: ToolUiMessageType::ToolError,
205            parts: vec![Part::Text(message)],
206        })
207    }
208
209    fn supports_progress(&self) -> bool {
210        true
211    }
212
213    fn render_tool_progress(&self, context: &ToolUiContext) -> Result<Option<ToolUiMessage>> {
214        if let Some(progress_info) = &context.progress_info
215            && let Some(status) = progress_info.get("status").and_then(|s| s.as_str()) {
216                let message = format!("⏳ **Code execution:** {}", status);
217
218                return Ok(Some(ToolUiMessage {
219                    message_type: ToolUiMessageType::ToolProgress,
220                    parts: vec![Part::Text(message)],
221                }));
222            }
223
224        Ok(None)
225    }
226}
227
228/// Helper function to register common tool renderers
229pub fn register_common_renderers(registry: &mut crate::ToolUiRenderRegistry) {
230    registry.register("search".to_string(), Arc::new(SearchToolRenderer));
231    registry.register("file_read".to_string(), Arc::new(FileToolRenderer));
232    registry.register("file_write".to_string(), Arc::new(FileToolRenderer));
233    registry.register("file_delete".to_string(), Arc::new(FileToolRenderer));
234    registry.register("file_copy".to_string(), Arc::new(FileToolRenderer));
235    registry.register("file_move".to_string(), Arc::new(FileToolRenderer));
236    registry.register(
237        "distri_execute_code".to_string(),
238        Arc::new(CodeExecutionToolRenderer),
239    );
240}