CHORUS/chrs-exec/src/lib.rs

//! chrs-exec: Isolated Task Execution Engine for CHORUS.

use bollard::container::{Config, CreateContainerOptions, StartContainerOptions, LogOutput, LogsOptions};
use bollard::Docker;
use bollard::image::CreateImageOptions;
use bollard::models::HostConfig;
use futures_util::stream::TryStreamExt;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use std::collections::HashMap;
use chrono::Utc;
use uuid::Uuid;

/// Represents a request for task execution.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TaskRequest {
    pub language: String,
    pub code: Option<String>,          // Optional raw code to execute
    pub agent_prompt: Option<String>,  // NEW: Optional prompt for the internal agent (opencode)
    pub workspace_path: Option<String>, // NEW: Path to mount into the container
    pub timeout_secs: u64,
}

/// Results of a task execution.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TaskResult {
    pub exit_code: i64,
    pub stdout: String,
    pub stderr: String,
    pub duration_ms: u64,
}

#[derive(Debug, Error)]
pub enum ExecError {
    #[error("Docker error: {0}")]
    Docker(#[from] bollard::errors::Error),
    #[error("Task timeout")]
    Timeout,
    #[error("Image not found: {0}")]
    ImageNotFound(String),
}

/// Executes tasks within isolated Docker containers.
pub struct DockerExecutor {
    docker: Docker,
    image_map: HashMap<String, String>,
}

impl DockerExecutor {
    pub fn new() -> Result<Self, ExecError> {
        let docker = Docker::connect_with_local_defaults()?;
        let mut image_map = HashMap::new();
        image_map.insert("rust".into(), "anthonyrawlins/chorus-rust-dev:latest".into());
        image_map.insert("python".into(), "anthonyrawlins/chorus-python-dev:latest".into());
        image_map.insert("base".into(), "anthonyrawlins/chorus-base:latest".into());
        
        Ok(Self { docker, image_map })
    }

    /// Executes a task using the "Agent Inception" pattern if a prompt is provided.
    pub async fn execute(&self, task: TaskRequest) -> Result<TaskResult, ExecError> {
        let start_time = Utc::now();
        let image = self.image_map.get(&task.language)
            .or_else(|| self.image_map.get("base"))
            .unwrap();

        // 1. Ensure image is pulled
        self.docker.create_image(
            Some(CreateImageOptions { from_image: image.as_str(), ..Default::default() }),
            None,
            None
        ).try_collect::<Vec<_>>().await?;

        // 2. Prepare Command (Agent Inception vs Raw Script)
        // Using Vec<String> to avoid lifetime issues with temporary formatted strings
        let cmd: Vec<String> = if let Some(prompt) = &task.agent_prompt {
            vec!["opencode".to_string(), "run".to_string(), "--prompt".to_string(), prompt.clone()]
        } else if let Some(code) = &task.code {
            vec![
                "bash".to_string(), 
                "-c".to_string(), 
                format!("echo '{}' > task.code && bash task.code", code.replace("'", "'\\''"))
            ]
        } else {
            vec!["ls".to_string(), "-la".to_string()]
        };

        // 3. Create Container with optional Volume mount
        let container_name = format!("chrs-task-{}", Uuid::new_v4());
        let mut binds = Vec::new();
        if let Some(path) = &task.workspace_path {
            binds.push(format!("{}:/workspace", path));
        }

        let host_config = HostConfig {
            memory: Some(2 * 1024 * 1024 * 1024), // 2GB
            nano_cpus: Some(2_000_000_000),      // 2 Cores
            network_mode: Some("host".into()),   // Allow access to local Gitea for this test
            binds: Some(binds),
            ..Default::default()
        };

        let config = Config {
            image: Some(image.as_str()),
            cmd: Some(cmd.iter().map(|s| s.as_str()).collect()),
            working_dir: Some("/workspace"),
            host_config: Some(host_config),
            ..Default::default()
        };

        self.docker.create_container(
            Some(CreateContainerOptions { name: container_name.as_str(), ..Default::default() }),
            config
        ).await?;

        // 4. Start and Wait
        self.docker.start_container(&container_name, None::<StartContainerOptions<String>>).await?;
        
        // 5. Collect Logs
        let logs = self.docker.logs(
            &container_name,
            Some(LogsOptions::<String> {
                stdout: true,
                stderr: true,
                follow: true,
                ..Default::default()
            })
        ).try_collect::<Vec<_>>().await?;

        let mut stdout = String::new();
        let mut stderr = String::new();

        for log in logs {
            match log {
                LogOutput::StdOut { message } => stdout.push_str(&String::from_utf8_lossy(&message)),
                LogOutput::StdErr { message } => stderr.push_str(&String::from_utf8_lossy(&message)),
                _ => {}
            }
        }

        // 6. Wait for container to exit and get exit code
        let wait_stream = self.docker.wait_container(&container_name, None::<bollard::container::WaitContainerOptions<String>>);
        let wait_res = wait_stream.try_collect::<Vec<_>>().await?;
        let exit_code = wait_res.first().map(|r| r.status_code).unwrap_or(0);

        // 7. Cleanup
        let _ = self.docker.stop_container(&container_name, None).await;
        let _ = self.docker.remove_container(&container_name, None).await;

        let end_time = Utc::now();
        let duration = end_time.signed_duration_since(start_time).num_milliseconds() as u64;

        Ok(TaskResult {
            exit_code,
            stdout,
            stderr,
            duration_ms: duration,
        })
    }
}