Skip to main content

tycode_core/modules/execution/
mod.rs

1pub mod config;
2
3use std::path::PathBuf;
4use std::sync::Arc;
5use std::time::Duration;
6use std::{env, process::Stdio};
7
8use anyhow::{anyhow, Result};
9use serde::Serialize;
10use serde_json::{json, Value};
11use tokio::process::Command;
12
13use std::collections::VecDeque;
14
15use crate::chat::events::{ToolExecutionResult, ToolRequest as ToolRequestEvent, ToolRequestType};
16use crate::file::access::FileAccessManager;
17use crate::module::Module;
18use crate::module::PromptComponent;
19use crate::module::{ContextComponent, ContextComponentId};
20use crate::settings::SettingsManager;
21use crate::tools::r#trait::{
22    ContinuationPreference, ToolCallHandle, ToolCategory, ToolExecutor, ToolOutput, ToolRequest,
23};
24use crate::tools::ToolName;
25
26use config::{CommandExecutionMode, ExecutionConfig, RunBuildTestOutputMode};
27
28const BLOCKED_COMMANDS: &[&str] = &["rm", "rmdir", "dd", "shred", "mkfs", "fdisk", "parted"];
29
30// === Command Outputs Context Component ===
31
32pub const COMMAND_OUTPUTS_ID: ContextComponentId = ContextComponentId("command_outputs");
33
34/// A stored command output with its command and result.
35#[derive(Debug, Clone)]
36pub struct CommandOutput {
37    pub command: String,
38    pub output: String,
39    pub exit_code: Option<i32>,
40}
41
42/// Manages command output history and provides context rendering.
43/// Stores a fixed-size buffer of recent command outputs.
44pub struct CommandOutputsManager {
45    outputs: std::sync::RwLock<VecDeque<CommandOutput>>,
46    max_outputs: usize,
47}
48
49impl CommandOutputsManager {
50    pub fn new(max_outputs: usize) -> Self {
51        Self {
52            outputs: std::sync::RwLock::new(VecDeque::with_capacity(max_outputs)),
53            max_outputs,
54        }
55    }
56
57    /// Add a command output to the buffer.
58    /// If buffer is full, oldest output is removed.
59    pub fn add_output(&self, command: String, output: String, exit_code: Option<i32>) {
60        let mut outputs = self.outputs.write().unwrap();
61        if outputs.len() >= self.max_outputs {
62            outputs.pop_front();
63        }
64        outputs.push_back(CommandOutput {
65            command,
66            output,
67            exit_code,
68        });
69    }
70
71    /// Clear all stored outputs.
72    pub fn clear(&self) {
73        self.outputs.write().unwrap().clear();
74    }
75
76    /// Get the number of stored outputs.
77    pub fn len(&self) -> usize {
78        self.outputs.read().unwrap().len()
79    }
80
81    /// Check if buffer is empty.
82    pub fn is_empty(&self) -> bool {
83        self.outputs.read().unwrap().is_empty()
84    }
85}
86
87#[async_trait::async_trait(?Send)]
88impl ContextComponent for CommandOutputsManager {
89    fn id(&self) -> ContextComponentId {
90        COMMAND_OUTPUTS_ID
91    }
92
93    async fn build_context_section(&self) -> Option<String> {
94        let outputs: Vec<CommandOutput> = {
95            let mut guard = self.outputs.write().unwrap();
96            guard.drain(..).collect()
97        };
98
99        if outputs.is_empty() {
100            return None;
101        }
102
103        let mut result = String::from("Recent Command Outputs:\n");
104        for output in outputs.iter() {
105            result.push_str(&format!("\n$ {}\n", output.command));
106            if let Some(code) = output.exit_code {
107                result.push_str(&format!("Exit code: {}\n", code));
108            }
109            if !output.output.is_empty() {
110                result.push_str(&output.output);
111                if !output.output.ends_with('\n') {
112                    result.push('\n');
113                }
114            }
115        }
116        Some(result)
117    }
118}
119
120#[derive(Debug, Clone, Serialize)]
121pub struct CommandResult {
122    pub command: String,
123    pub code: i32,
124    pub out: String,
125    pub err: String,
126}
127
128pub async fn run_cmd(
129    dir: PathBuf,
130    cmd: String,
131    timeout: Duration,
132    execution_mode: CommandExecutionMode,
133) -> Result<CommandResult> {
134    let path = env::var("PATH")?;
135    tracing::info!(?path, ?dir, ?cmd, ?execution_mode, "Attempting to run_cmd");
136
137    let child = match execution_mode {
138        CommandExecutionMode::Direct => {
139            let parts = shell_words::split(&cmd)
140                .map_err(|e| anyhow::anyhow!("Failed to parse command: {e:?}"))?;
141            if parts.is_empty() {
142                return Err(anyhow::anyhow!("Empty command"));
143            }
144            let program = &parts[0];
145            let args: Vec<&str> = parts[1..].iter().map(|s| s.as_str()).collect();
146
147            Command::new(program)
148                .args(args)
149                .current_dir(&dir)
150                .stdout(Stdio::piped())
151                .stderr(Stdio::piped())
152                .kill_on_drop(true)
153                .spawn()?
154        }
155        CommandExecutionMode::Bash => Command::new("bash")
156            .args(["-c", &cmd])
157            .current_dir(&dir)
158            .stdout(Stdio::piped())
159            .stderr(Stdio::piped())
160            .kill_on_drop(true)
161            .spawn()?,
162    };
163
164    let output = tokio::time::timeout(timeout, async {
165        let output = child.wait_with_output().await?;
166        Ok::<_, std::io::Error>(output)
167    })
168    .await??;
169
170    let code = output.status.code().unwrap_or(1);
171    let out = String::from_utf8_lossy(&output.stdout).to_string();
172    let err = String::from_utf8_lossy(&output.stderr).to_string();
173
174    Ok(CommandResult {
175        command: cmd,
176        code,
177        out,
178        err,
179    })
180}
181
182pub struct ExecutionModule {
183    inner: Arc<ExecutionModuleInner>,
184}
185
186struct ExecutionModuleInner {
187    command_outputs_manager: Arc<CommandOutputsManager>,
188    access: FileAccessManager,
189    settings: SettingsManager,
190}
191
192impl ExecutionModule {
193    pub fn new(workspace_roots: Vec<PathBuf>, settings: SettingsManager) -> Result<Self> {
194        let inner = Arc::new(ExecutionModuleInner {
195            command_outputs_manager: Arc::new(CommandOutputsManager::new(10)),
196            access: FileAccessManager::new(workspace_roots)?,
197            settings,
198        });
199        Ok(Self { inner })
200    }
201}
202
203impl Module for ExecutionModule {
204    fn prompt_components(&self) -> Vec<Arc<dyn PromptComponent>> {
205        vec![]
206    }
207
208    fn context_components(&self) -> Vec<Arc<dyn ContextComponent>> {
209        vec![self.inner.command_outputs_manager.clone()]
210    }
211
212    fn tools(&self) -> Vec<Arc<dyn ToolExecutor>> {
213        vec![Arc::new(RunBuildTestTool {
214            inner: self.inner.clone(),
215        })]
216    }
217
218    fn session_state(&self) -> Option<Arc<dyn crate::module::SessionStateComponent>> {
219        None
220    }
221
222    fn settings_namespace(&self) -> Option<&'static str> {
223        Some("execution")
224    }
225
226    fn settings_json_schema(&self) -> Option<schemars::schema::RootSchema> {
227        Some(schemars::schema_for!(ExecutionConfig))
228    }
229}
230
231pub struct RunBuildTestTool {
232    inner: Arc<ExecutionModuleInner>,
233}
234
235impl RunBuildTestTool {
236    pub fn tool_name() -> ToolName {
237        ToolName::new("run_build_test")
238    }
239}
240
241struct RunBuildTestHandle {
242    command: String,
243    working_directory: PathBuf,
244    timeout_seconds: u64,
245    tool_use_id: String,
246    command_outputs_manager: Arc<CommandOutputsManager>,
247    output_mode: RunBuildTestOutputMode,
248    execution_mode: CommandExecutionMode,
249    max_output_bytes: Option<usize>,
250}
251
252/// Compact output by keeping first half and last half with truncation marker.
253fn compact_output(output: &str, max_bytes: usize) -> String {
254    if output.len() <= max_bytes {
255        return output.to_string();
256    }
257
258    let half = max_bytes / 2;
259    let start_end = output.floor_char_boundary(half);
260    let end_start_target = output.len().saturating_sub(half);
261    let end_start = output.ceil_char_boundary(end_start_target);
262
263    let start = &output[..start_end];
264    let end = &output[end_start..];
265    let omitted = output.len() - start.len() - end.len();
266
267    format!(
268        "{}\n... [output truncated: {} bytes omitted] ...\n{}",
269        start, omitted, end
270    )
271}
272
273#[async_trait::async_trait(?Send)]
274impl ToolCallHandle for RunBuildTestHandle {
275    fn tool_request(&self) -> ToolRequestEvent {
276        ToolRequestEvent {
277            tool_call_id: self.tool_use_id.clone(),
278            tool_name: "run_build_test".to_string(),
279            tool_type: ToolRequestType::RunCommand {
280                command: self.command.clone(),
281                working_directory: self.working_directory.to_string_lossy().to_string(),
282            },
283        }
284    }
285
286    async fn execute(self: Box<Self>) -> ToolOutput {
287        let timeout = Duration::from_secs(self.timeout_seconds);
288
289        let result = match run_cmd(
290            self.working_directory.clone(),
291            self.command.clone(),
292            timeout,
293            self.execution_mode.clone(),
294        )
295        .await
296        {
297            Ok(r) => r,
298            Err(e) => {
299                let error_msg = format!("Command execution failed: {e:?}");
300                return ToolOutput::Result {
301                    content: error_msg.clone(),
302                    is_error: true,
303                    continuation: ContinuationPreference::Continue,
304                    ui_result: ToolExecutionResult::Error {
305                        short_message: "Command failed".to_string(),
306                        detailed_message: error_msg,
307                    },
308                };
309            }
310        };
311
312        let combined_output = if result.err.is_empty() {
313            result.out.clone()
314        } else if result.out.is_empty() {
315            result.err.clone()
316        } else {
317            format!("{}\n{}", result.out, result.err)
318        };
319
320        let combined_output = match self.max_output_bytes {
321            Some(max) if combined_output.len() > max => compact_output(&combined_output, max),
322            _ => combined_output,
323        };
324
325        self.command_outputs_manager.add_output(
326            self.command.clone(),
327            combined_output,
328            Some(result.code),
329        );
330
331        let is_error = result.code != 0;
332        let content = match (&self.output_mode, is_error) {
333            (RunBuildTestOutputMode::ToolResponse, _) => json!({
334                "exit_code": result.code,
335                "stdout": result.out,
336                "stderr": result.err,
337            })
338            .to_string(),
339            (RunBuildTestOutputMode::Context, true) => json!({
340                "exit_code": result.code,
341                "status": "failed",
342                "message": "Command failed. See context section for output."
343            })
344            .to_string(),
345            (RunBuildTestOutputMode::Context, false) => json!({
346                "exit_code": result.code,
347                "status": "success",
348                "message": "Command executed. See context section for output."
349            })
350            .to_string(),
351        };
352
353        ToolOutput::Result {
354            content,
355            is_error,
356            continuation: ContinuationPreference::Continue,
357            ui_result: ToolExecutionResult::RunCommand {
358                exit_code: result.code,
359                stdout: result.out,
360                stderr: result.err,
361            },
362        }
363    }
364}
365
366#[async_trait::async_trait(?Send)]
367impl ToolExecutor for RunBuildTestTool {
368    fn name(&self) -> String {
369        "run_build_test".to_string()
370    }
371
372    fn description(&self) -> String {
373        let config: ExecutionConfig = self.inner.settings.get_module_config("execution");
374        match config.execution_mode {
375            CommandExecutionMode::Direct => {
376                "Run build, test, or execution commands (cargo build, npm test, python main.py) - NOT for file operations (no cat/ls/grep/find); use dedicated file tools instead. Shell features like pipes (cmd | grep) and redirects (cmd > file) will fail or behave unexpectedly.".to_string()
377            }
378            CommandExecutionMode::Bash => {
379                "Run build, test, or execution commands (cargo build, npm test, python main.py) - NOT for file operations (no cat/ls/grep/find); use dedicated file tools instead.".to_string()
380            }
381        }
382    }
383
384    fn input_schema(&self) -> Value {
385        json!({
386            "type": "object",
387            "properties": {
388                "command": {
389                    "type": "string",
390                    "description": "The command to execute"
391                },
392                "working_directory": {
393                    "type": "string",
394                    "description": "The directory to run the command in. Must be within a workspace root. Must be an absolute path."
395                },
396                "timeout_seconds": {
397                    "type": "integer",
398                    "description": "Maximum seconds to wait for command completion",
399                    "minimum": 1,
400                    "maximum": 300
401                }
402            },
403            "required": ["command", "timeout_seconds", "working_directory"]
404        })
405    }
406
407    fn category(&self) -> ToolCategory {
408        ToolCategory::Execution
409    }
410
411    async fn process(&self, request: &ToolRequest) -> Result<Box<dyn ToolCallHandle>> {
412        let command_str = request
413            .arguments
414            .get("command")
415            .and_then(|v| v.as_str())
416            .ok_or_else(|| anyhow!("Missing 'command' argument"))?;
417
418        let timeout_seconds = request
419            .arguments
420            .get("timeout_seconds")
421            .and_then(|v| v.as_u64())
422            .ok_or_else(|| anyhow!("Missing 'timeout_seconds' argument"))?;
423
424        let working_directory = request
425            .arguments
426            .get("working_directory")
427            .and_then(|v| v.as_str())
428            .ok_or_else(|| anyhow!("Missing 'working_directory' argument"))?;
429        let resolved_working_directory = self.inner.access.resolve(working_directory)?;
430
431        let parts: Vec<&str> = command_str.split_whitespace().collect();
432        if parts.is_empty() {
433            return Err(anyhow!("Empty command"));
434        }
435
436        let cmd = parts[0];
437        if BLOCKED_COMMANDS.contains(&cmd) || cmd.starts_with("mkfs.") {
438            let msg = if cmd == "rm" || cmd == "rmdir" {
439                format!("Command '{cmd}' is blocked for safety. Use the delete_file tool instead.")
440            } else {
441                format!("Command '{cmd}' is blocked for safety.")
442            };
443            return Err(anyhow!(msg));
444        }
445
446        let config: ExecutionConfig = self.inner.settings.get_module_config("execution");
447        let output_mode = config.output_mode.clone();
448        let execution_mode = config.execution_mode.clone();
449
450        Ok(Box::new(RunBuildTestHandle {
451            command: command_str.to_string(),
452            working_directory: resolved_working_directory,
453            timeout_seconds,
454            tool_use_id: request.tool_use_id.clone(),
455            command_outputs_manager: self.inner.command_outputs_manager.clone(),
456            output_mode,
457            execution_mode,
458            max_output_bytes: config.max_output_bytes,
459        }))
460    }
461}