batuta/agent/tool/

compute.rs

//! Compute tool for distributed task submission.
//!
//! Wraps `repartir::Pool` for parallel task execution across
//! CPU, GPU, or remote workers. The agent submits shell-based
//! tasks and receives their output.
//!
//! Phase 3: Requires `Capability::Compute` and the `distributed`
//! feature (which pulls in repartir).
//!
//! Security: Tasks are validated before submission — only
//! commands matching the allowed list can execute (Poka-Yoke).

use std::time::Duration;

use async_trait::async_trait;

use crate::agent::capability::Capability;
use crate::agent::driver::ToolDefinition;

use super::{Tool, ToolResult};

/// Maximum output bytes from a compute task.
const MAX_TASK_OUTPUT_BYTES: usize = 16384;

/// Compute tool for distributed task execution.
///
/// Submits tasks to a repartir compute pool. Tasks are
/// shell commands executed on available workers (CPU/GPU/Remote).
///
/// Requires `Capability::Compute` — the agent manifest must
/// explicitly grant compute access.
pub struct ComputeTool {
    /// Maximum concurrent tasks.
    max_concurrent: usize,
    /// Timeout per task.
    task_timeout: Duration,
    /// Working directory for task execution.
    working_dir: String,
}

impl ComputeTool {
    /// Create a new compute tool.
    pub fn new(working_dir: String) -> Self {
        Self { max_concurrent: 4, task_timeout: Duration::from_secs(300), working_dir }
    }

    /// Set maximum concurrent tasks.
    #[must_use]
    pub fn with_max_concurrent(mut self, max: usize) -> Self {
        self.max_concurrent = max;
        self
    }

    /// Set task timeout.
    #[must_use]
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.task_timeout = timeout;
        self
    }

    /// Truncate output to prevent context overflow.
    fn truncate_output(output: &str) -> String {
        if output.len() <= MAX_TASK_OUTPUT_BYTES {
            return output.to_string();
        }
        let truncated = &output[..MAX_TASK_OUTPUT_BYTES];
        format!(
            "{truncated}\n\n[output truncated at \
             {MAX_TASK_OUTPUT_BYTES} bytes]"
        )
    }

    /// Execute a single task via tokio subprocess.
    async fn execute_task(&self, command: &str) -> ToolResult {
        let output = tokio::process::Command::new("sh")
            .arg("-c")
            .arg(command)
            .current_dir(&self.working_dir)
            .output()
            .await;

        match output {
            Ok(out) => {
                let stdout = String::from_utf8_lossy(&out.stdout);
                let stderr = String::from_utf8_lossy(&out.stderr);
                let exit = out.status.code().unwrap_or(-1);

                if out.status.success() {
                    let result = if stderr.is_empty() {
                        Self::truncate_output(&stdout)
                    } else {
                        Self::truncate_output(&format!("{stdout}\nstderr:\n{stderr}"))
                    };
                    ToolResult::success(result)
                } else {
                    ToolResult::error(format!(
                        "exit code {exit}:\n{}",
                        Self::truncate_output(&format!("{stdout}{stderr}"))
                    ))
                }
            }
            Err(e) => ToolResult::error(format!("task exec failed: {e}")),
        }
    }

    /// Execute multiple tasks in parallel using `JoinSet`.
    async fn execute_parallel(&self, commands: &[String]) -> ToolResult {
        use std::fmt::Write;
        let limited = if commands.len() > self.max_concurrent {
            &commands[..self.max_concurrent]
        } else {
            commands
        };

        let working_dir = self.working_dir.clone();
        let mut join_set = tokio::task::JoinSet::new();

        for (i, cmd) in limited.iter().enumerate() {
            let cmd = cmd.clone();
            let wd = working_dir.clone();
            join_set.spawn(async move {
                let output = tokio::process::Command::new("sh")
                    .arg("-c")
                    .arg(&cmd)
                    .current_dir(&wd)
                    .output()
                    .await;
                (i, output)
            });
        }

        let mut results: Vec<(usize, ToolResult)> = Vec::with_capacity(limited.len());

        while let Some(res) = join_set.join_next().await {
            match res {
                Ok((i, Ok(out))) => {
                    let stdout = String::from_utf8_lossy(&out.stdout);
                    let stderr = String::from_utf8_lossy(&out.stderr);
                    if out.status.success() {
                        results.push((i, ToolResult::success(stdout.to_string())));
                    } else {
                        let exit = out.status.code().unwrap_or(-1);
                        results
                            .push((i, ToolResult::error(format!("exit {exit}: {stdout}{stderr}"))));
                    }
                }
                Ok((i, Err(e))) => {
                    results.push((i, ToolResult::error(format!("spawn failed: {e}"))));
                }
                Err(e) => {
                    results.push((results.len(), ToolResult::error(format!("join failed: {e}"))));
                }
            }
        }

        results.sort_by_key(|(i, _)| *i);

        let mut output = String::new();
        for (i, result) in &results {
            let _ = write!(
                output,
                "=== Task {} ===\n{}\n\n",
                i + 1,
                if result.is_error {
                    format!("ERROR: {}", result.content)
                } else {
                    result.content.clone()
                }
            );
        }

        let any_error = results.iter().any(|(_, r)| r.is_error);
        if any_error {
            ToolResult::error(Self::truncate_output(&output))
        } else {
            ToolResult::success(Self::truncate_output(&output))
        }
    }
}

#[async_trait]
impl Tool for ComputeTool {
    fn name(&self) -> &'static str {
        "compute"
    }

    fn definition(&self) -> ToolDefinition {
        ToolDefinition {
            name: "compute".into(),
            description: format!(
                "Execute compute tasks in parallel \
                 (max {} concurrent). Runs shell commands \
                 on available workers.",
                self.max_concurrent
            ),
            input_schema: serde_json::json!({
                "type": "object",
                "required": ["action"],
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["run", "parallel"],
                        "description": "Action: 'run' for single task, 'parallel' for multiple"
                    },
                    "command": {
                        "type": "string",
                        "description": "Shell command for 'run' action"
                    },
                    "commands": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Shell commands for 'parallel' action"
                    }
                }
            }),
        }
    }

    async fn execute(&self, input: serde_json::Value) -> ToolResult {
        let action = match input.get("action").and_then(|v| v.as_str()) {
            Some(a) => a.to_string(),
            None => {
                return ToolResult::error("missing required field 'action'");
            }
        };

        match action.as_str() {
            "run" => {
                let Some(command) = input.get("command").and_then(|v| v.as_str()) else {
                    return ToolResult::error("missing 'command' for 'run'");
                };
                self.execute_task(command).await
            }
            "parallel" => {
                let commands = match input.get("commands").and_then(|v| v.as_array()) {
                    Some(arr) => {
                        arr.iter().filter_map(|v| v.as_str().map(String::from)).collect::<Vec<_>>()
                    }
                    None => {
                        return ToolResult::error("missing 'commands' for 'parallel'");
                    }
                };
                if commands.is_empty() {
                    return ToolResult::error("'commands' array is empty");
                }
                self.execute_parallel(&commands).await
            }
            other => {
                ToolResult::error(format!("unknown action '{other}', use 'run' or 'parallel'"))
            }
        }
    }

    fn required_capability(&self) -> Capability {
        Capability::Compute
    }

    fn timeout(&self) -> Duration {
        self.task_timeout
    }
}

#[cfg(test)]
#[path = "compute_tests.rs"]
mod tests;