agentic_warden/

supervisor.rs

use crate::cli_type::CliType;
use crate::core::models::ProcessTreeInfo;
use crate::core::process_tree::ProcessTreeError;
use crate::error::RegistryError;
use crate::logging::debug;
use crate::logging::warn;
#[cfg(windows)]
use crate::platform::ChildResources;
use crate::platform::{self};
use crate::provider::{AiType, ProviderManager};
use crate::signal;
use crate::storage::TaskStorage;
use crate::task_record::TaskRecord;
use crate::unified_registry::Registry;
use chrono::{DateTime, Utc};
use std::env;
use std::ffi::OsString;
use std::io;
use std::path::PathBuf;
use std::process::{ExitStatus, Stdio};
use std::sync::Arc;
use thiserror::Error;
use tokio::fs::OpenOptions;
use tokio::io::{AsyncRead, AsyncWriteExt, BufWriter};
use tokio::process::Command;
use tokio::sync::Mutex;

#[derive(Debug, Error)]
pub enum ProcessError {
    #[error("IO error: {0}")]
    Io(#[from] io::Error),
    #[error("Registry error: {0}")]
    Registry(#[from] RegistryError),
    #[error("Process tree error: {0}")]
    ProcessTree(#[from] ProcessTreeError),
    #[error("CLI executable not found: {0}")]
    CliNotFound(String),
    #[error("{0}")]
    Other(String),
}

async fn get_cli_command(cli_type: &CliType) -> Result<String, ProcessError> {
    // First try environment variable
    if let Ok(custom_path) = env::var(cli_type.env_var_name()) {
        if custom_path.is_empty() {
            return Err(ProcessError::CliNotFound(format!(
                "{} environment variable is empty",
                cli_type.env_var_name()
            )));
        }
        return Ok(custom_path);
    }

    // Fall back to default command name
    let default_cmd = cli_type.command_name();

    // On Windows, try to find the actual executable path
    if cfg!(windows) {
        let output = Command::new("where")
            .arg(default_cmd)
            .output()
            .await
            .map_err(|_| {
                ProcessError::CliNotFound(format!(
                    "Failed to check if '{}' exists in PATH",
                    default_cmd
                ))
            })?;

        if output.status.success() {
            let stdout = String::from_utf8_lossy(&output.stdout);
            // Prefer .cmd files on Windows, otherwise use first result
            for line in stdout.lines() {
                if line.ends_with(".cmd") || line.ends_with(".bat") || line.ends_with(".exe") {
                    return Ok(line.to_string());
                }
            }
            // Fallback to first line if no Windows executable found
            if let Some(first_line) = stdout.lines().next() {
                return Ok(first_line.to_string());
            }
        }

        return Err(ProcessError::CliNotFound(format!(
            "'{}' not found in PATH. Set {} environment variable or ensure it's in PATH",
            default_cmd,
            cli_type.env_var_name()
        )));
    } else {
        let output = Command::new("which")
            .arg(default_cmd)
            .output()
            .await
            .map_err(|_| {
                ProcessError::CliNotFound(format!(
                    "Failed to check if '{}' exists in PATH",
                    default_cmd
                ))
            })?;

        if !output.status.success() {
            return Err(ProcessError::CliNotFound(format!(
                "'{}' not found in PATH. Set {} environment variable or ensure it's in PATH",
                default_cmd,
                cli_type.env_var_name()
            )));
        }
    }

    Ok(default_cmd.to_string())
}

/// Output handling strategy for CLI execution
enum OutputStrategy {
    /// Mirror output to stdout/stderr
    Mirror,
    /// Capture stdout to buffer
    Capture(Arc<Mutex<Vec<u8>>>),
}

pub async fn execute_cli<S: TaskStorage>(
    registry: &Registry<S>,
    cli_type: &CliType,
    args: &[OsString],
    provider: Option<String>,
) -> Result<i32, ProcessError> {
    execute_cli_internal(
        registry,
        cli_type,
        args,
        provider,
        None,
        OutputStrategy::Mirror,
    )
    .await
    .map(|(exit_code, _)| exit_code)
}

/// Execute CLI and capture stdout output (for code generation)
pub async fn execute_cli_with_output<S: TaskStorage>(
    registry: &Registry<S>,
    cli_type: &CliType,
    args: &[OsString],
    provider: Option<String>,
    timeout: std::time::Duration,
) -> Result<String, ProcessError> {
    let buffer = Arc::new(Mutex::new(Vec::new()));
    let (_, output_opt) = execute_cli_internal(
        registry,
        cli_type,
        args,
        provider,
        Some(timeout),
        OutputStrategy::Capture(buffer.clone()),
    )
    .await?;

    match output_opt {
        Some(output) => Ok(output),
        None => Err(ProcessError::Other(
            "Output capture failed unexpectedly".to_string(),
        )),
    }
}

/// Internal CLI execution with configurable output handling
async fn execute_cli_internal<S: TaskStorage>(
    registry: &Registry<S>,
    cli_type: &CliType,
    args: &[OsString],
    provider: Option<String>,
    timeout: Option<std::time::Duration>,
    output_strategy: OutputStrategy,
) -> Result<(i32, Option<String>), ProcessError> {
    let is_capture_mode = matches!(output_strategy, OutputStrategy::Capture(_));

    platform::init_platform();

    let terminate_wrapper = |pid: u32| {
        platform::terminate_process(pid);
        Ok(())
    };
    registry.sweep_stale_entries(Utc::now(), platform::process_alive, &terminate_wrapper)?;

    // Load provider configuration
    let provider_manager = ProviderManager::new()
        .map_err(|e| ProcessError::Other(format!("Failed to load provider: {}", e)))?;

    // Determine which provider to use
    let (provider_name, provider_config) = if let Some(name) = provider {
        let config = provider_manager
            .get_provider(&name)
            .map_err(|e| ProcessError::Other(e.to_string()))?;
        (name, config)
    } else {
        let (name, config) = provider_manager
            .get_default_provider()
            .ok_or_else(|| ProcessError::Other("No default provider configured".to_string()))?;
        (name, config)
    };

    // Validate compatibility
    let _ai_type = match cli_type {
        CliType::Claude => AiType::Claude,
        CliType::Codex => AiType::Codex,
        CliType::Gemini => AiType::Gemini,
    };

    // Display provider info if not using official
    if provider_name != *"official" {
        eprintln!(
            "Using provider: {} ({})",
            provider_name,
            provider_config.summary()
        );
    }

    let cli_command = get_cli_command(cli_type).await?;

    let mut command = Command::new(&cli_command);
    command.args(args);
    command.stdin(Stdio::null());
    command.stdout(Stdio::piped());
    command.stderr(Stdio::piped());

    // Platform-specific command preparation
    #[cfg(unix)]
    {
        unsafe {
            command.pre_exec(|| {
                let result = libc::setpgid(0, 0);
                if result != 0 {
                    return Err(io::Error::last_os_error());
                }
                #[cfg(target_os = "linux")]
                {
                    let result = libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM);
                    if result != 0 {
                        return Err(io::Error::last_os_error());
                    }
                }
                Ok(())
            });
        }
    }

    // Inject environment variables
    for (key, value) in &provider_config.env {
        command.env(key, value);
    }

    let mut child = command.spawn()?;
    let child_pid = child
        .id()
        .ok_or_else(|| io::Error::other("Failed to get child PID"))?;

    let log_path = match generate_log_path(child_pid) {
        Ok(path) => path,
        Err(err) => {
            platform::terminate_process(child_pid);
            let _ = child.wait();
            return Err(err.into());
        }
    };

    let log_file = match OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(&log_path)
        .await
    {
        Ok(file) => file,
        Err(err) => {
            platform::terminate_process(child_pid);
            let _ = child.wait().await;
            return Err(err.into());
        }
    };

    debug(format!(
        "Started {} process pid={} log={}{}",
        cli_type.display_name(),
        child_pid,
        log_path.display(),
        if is_capture_mode {
            " (capture mode)"
        } else {
            ""
        }
    ));

    #[cfg(windows)]
    {
        let _resources = ChildResources::with_job(None);
    }

    let signal_guard = signal::install(child_pid)?;

    let log_writer = Arc::new(Mutex::new(BufWriter::new(log_file)));
    let mut copy_handles = Vec::new();

    // Handle stdout based on strategy
    if let Some(stdout) = child.stdout.take() {
        match &output_strategy {
            OutputStrategy::Mirror => {
                copy_handles.push(tokio::spawn(spawn_copy(
                    stdout,
                    log_writer.clone(),
                    StreamMirror::Stdout,
                )));
            }
            OutputStrategy::Capture(buffer) => {
                let buffer_clone = buffer.clone();
                let writer_clone = log_writer.clone();
                copy_handles.push(tokio::spawn(async move {
                    spawn_copy_with_capture(stdout, writer_clone, buffer_clone).await
                }));
            }
        }
    }

    if let Some(stderr) = child.stderr.take() {
        copy_handles.push(tokio::spawn(spawn_copy(
            stderr,
            log_writer.clone(),
            StreamMirror::Stderr,
        )));
    }

    let registration_guard = {
        let mut record = TaskRecord::new(
            Utc::now(),
            child_pid.to_string(),
            log_path.to_string_lossy().into_owned(),
            Some(platform::current_pid()),
        );

        // Get process tree information
        match ProcessTreeInfo::current() {
            Ok(tree_info) => match record.clone().with_process_tree_info(tree_info) {
                Ok(updated) => {
                    record = updated;
                }
                Err(err) => {
                    warn(format!("Failed to attach process tree info: {}", err));
                }
            },
            Err(err) => {
                warn(format!("Failed to get process tree info: {}", err));
            }
        }

        if let Err(err) = registry.register(child_pid, &record) {
            platform::terminate_process(child_pid);
            let _ = child.wait().await;
            return Err(err.into());
        }
        Some(RegistrationGuard::new(registry, child_pid))
    };

    // Wait with optional timeout
    let status = if let Some(timeout_duration) = timeout {
        tokio::select! {
            result = child.wait() => result?,
            _ = tokio::time::sleep(timeout_duration) => {
                platform::terminate_process(child_pid);
                let _ = child.wait().await;
                return Err(ProcessError::Other(format!(
                    "CLI execution timed out after {:?}",
                    timeout_duration
                )));
            }
        }
    } else {
        child.wait().await?
    };

    drop(signal_guard);

    for handle in copy_handles {
        match handle.await {
            Ok(result) => result?,
            Err(_) => {
                return Err(io::Error::other("Log writer task failed").into());
            }
        }
    }

    {
        let mut writer = log_writer.lock().await;
        writer.flush().await?;
        writer.get_ref().sync_all().await?;
    }

    // Display log file path to user (not in capture mode)
    if !is_capture_mode {
        eprintln!("完整日志已保存到: {}", log_path.display());
    }

    if let Some(guard) = registration_guard {
        let completed_at = Utc::now();
        let exit_code = status.code();
        let result = match (status.success(), exit_code) {
            (true, _) => Some(if is_capture_mode {
                "codegen_success".to_owned()
            } else {
                "success".to_owned()
            }),
            (false, Some(code)) => Some(format!(
                "{}_failed_with_exit_code_{code}",
                if is_capture_mode { "codegen" } else { "cli" }
            )),
            (false, None) => Some(format!(
                "{}_failed_without_exit_code",
                if is_capture_mode { "codegen" } else { "cli" }
            )),
        };
        let _ = guard.mark_completed(result, exit_code, completed_at);
    }

    // Extract captured output if in capture mode
    let captured_output = if let OutputStrategy::Capture(buffer) = output_strategy {
        let output = buffer.lock().await.clone();
        let output_str = String::from_utf8_lossy(&output).to_string();

        if !status.success() {
            return Err(ProcessError::Other(format!(
                "{} CLI failed with exit code {}: {}",
                cli_type.display_name(),
                extract_exit_code(status),
                output_str
            )));
        }

        Some(output_str)
    } else {
        None
    };

    Ok((extract_exit_code(status), captured_output))
}

/// Generate a secure log file path in runtime directory
///
/// Security considerations:
/// - Uses system temp directory (cross-platform)
/// - Creates directory with restrictive permissions (0700 on Unix)
/// - Ensures logs are only accessible by the current user
/// - Logs are automatically cleaned up on system reboot
fn generate_log_path(pid: u32) -> io::Result<PathBuf> {
    // Use system temp directory as per SPEC design (cross-platform)
    // Linux/macOS: /tmp/.aiw/logs/
    // Windows: %TEMP%\.aiw\logs\
    // Runtime data (logs, temp files) → temp_dir()/.aiw/
    // Persistent config → ~/.aiw/
    let log_dir = std::env::temp_dir().join(".aiw").join("logs");

    // Create the logs directory if it doesn't exist
    if !log_dir.exists() {
        std::fs::create_dir_all(&log_dir)?;

        // Set restrictive permissions on Unix systems (only user can read/write/execute)
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mut perms = std::fs::metadata(&log_dir)?.permissions();
            perms.set_mode(0o700); // rwx------
            std::fs::set_permissions(&log_dir, perms)?;
        }
    }

    Ok(log_dir.join(format!("{pid}.log")))
}

#[derive(Copy, Clone)]
enum StreamMirror {
    Stdout,
    Stderr,
}

impl StreamMirror {
    async fn write(self, data: &[u8]) -> io::Result<()> {
        use tokio::io::AsyncWriteExt;
        match self {
            StreamMirror::Stdout => {
                let mut handle = tokio::io::stdout();
                handle.write_all(data).await?;
                handle.flush().await
            }
            StreamMirror::Stderr => {
                let mut handle = tokio::io::stderr();
                handle.write_all(data).await?;
                handle.flush().await
            }
        }
    }
}

async fn spawn_copy<R>(
    mut reader: R,
    writer: Arc<Mutex<BufWriter<tokio::fs::File>>>,
    mirror: StreamMirror,
) -> io::Result<()>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    use tokio::io::AsyncReadExt;

    let mut buffer = [0u8; 8192];
    loop {
        let read = reader.read(&mut buffer).await?;
        if read == 0 {
            break;
        }
        let chunk = &buffer[..read];
        {
            let mut guard = writer.lock().await;
            guard.write_all(chunk).await?;
            guard.flush().await?;
        }
        mirror.write(chunk).await?;
    }
    Ok(())
}

/// Copy stream to log file and capture to buffer (for code generation)
async fn spawn_copy_with_capture<R>(
    mut reader: R,
    writer: Arc<Mutex<BufWriter<tokio::fs::File>>>,
    capture_buffer: Arc<Mutex<Vec<u8>>>,
) -> io::Result<()>
where
    R: AsyncRead + Unpin + Send + 'static,
{
    use tokio::io::AsyncReadExt;

    let mut buffer = [0u8; 8192];
    loop {
        let read = reader.read(&mut buffer).await?;
        if read == 0 {
            break;
        }
        let chunk = &buffer[..read];

        // Write to log file
        {
            let mut guard = writer.lock().await;
            guard.write_all(chunk).await?;
            guard.flush().await?;
        }

        // Capture to buffer
        {
            let mut capture = capture_buffer.lock().await;
            capture.extend_from_slice(chunk);
        }
    }
    Ok(())
}

fn extract_exit_code(status: ExitStatus) -> i32 {
    status.code().unwrap_or(1)
}

struct RegistrationGuard<'a, S: TaskStorage> {
    registry: &'a Registry<S>,
    pid: u32,
    active: bool,
}

impl<'a, S: TaskStorage> RegistrationGuard<'a, S> {
    fn new(registry: &'a Registry<S>, pid: u32) -> Self {
        Self {
            registry,
            pid,
            active: true,
        }
    }

    fn mark_completed(
        mut self,
        result: Option<String>,
        exit_code: Option<i32>,
        completed_at: DateTime<Utc>,
    ) -> Result<(), RegistryError> {
        if self.active {
            self.registry
                .mark_completed(self.pid, result, exit_code, completed_at)?;
            self.active = false;
        }
        Ok(())
    }
}

impl<S: TaskStorage> Drop for RegistrationGuard<'_, S> {
    fn drop(&mut self) {
        // 注意：TaskStorage trait不提供remove方法
        // 如果需要清理，应该通过mark_completed或sweep_stale_entries
        // 这里我们什么都不做，让任务记录保留在注册表中
    }
}

/// Start interactive CLI mode (directly launch AI CLI without task prompt)
pub async fn start_interactive_cli<S: TaskStorage>(
    registry: &Registry<S>,
    cli_type: &CliType,
    provider: Option<String>,
) -> Result<i32, ProcessError> {
    platform::init_platform();

    let terminate_wrapper = |pid: u32| {
        platform::terminate_process(pid);
        Ok(())
    };
    registry.sweep_stale_entries(Utc::now(), platform::process_alive, &terminate_wrapper)?;

    // Load provider configuration
    let provider_manager = ProviderManager::new()
        .map_err(|e| ProcessError::Other(format!("Failed to load provider: {}", e)))?;

    // Determine which provider to use
    let (provider_name, provider_config) = if let Some(name) = provider {
        let config = provider_manager
            .get_provider(&name)
            .map_err(|e| ProcessError::Other(e.to_string()))?;
        (name, config)
    } else {
        let (name, config) = provider_manager
            .get_default_provider()
            .ok_or_else(|| ProcessError::Other("No default provider configured".to_string()))?;
        (name, config)
    };

    // Validate compatibility
    let _ai_type = match cli_type {
        CliType::Claude => AiType::Claude,
        CliType::Codex => AiType::Codex,
        CliType::Gemini => AiType::Gemini,
    };

    // Display provider info if not using official
    if provider_name != *"official" {
        eprintln!(
            "Using provider: {} ({})",
            provider_name,
            provider_config.summary()
        );
    }

    let cli_command = get_cli_command(cli_type).await?;

    // Interactive mode: launch CLI with stdin/stdout/stderr inherited
    let mut command = Command::new(&cli_command);

    // Add interactive args (e.g., "exec" for Codex, "-p" for Claude)
    let interactive_args = cli_type.build_interactive_args();
    command.args(&interactive_args);

    command.stdin(Stdio::inherit());
    command.stdout(Stdio::inherit());
    command.stderr(Stdio::inherit());

    // Platform-specific command preparation (Unix: set process group and death signal)
    #[cfg(unix)]
    {
        unsafe {
            command.pre_exec(|| {
                // Set process group ID
                let result = libc::setpgid(0, 0);
                if result != 0 {
                    return Err(io::Error::last_os_error());
                }
                // Set parent death signal on Linux
                #[cfg(target_os = "linux")]
                {
                    let result = libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM);
                    if result != 0 {
                        return Err(io::Error::last_os_error());
                    }
                }
                Ok(())
            });
        }
    }

    // Inject environment variables
    for (key, value) in &provider_config.env {
        command.env(key, value);
    }

    let mut child = command.spawn()?;
    let child_pid = child
        .id()
        .ok_or_else(|| io::Error::other("Failed to get child PID"))?;

    // Register the interactive CLI process
    let log_path = generate_log_path(child_pid)?;
    let record = TaskRecord::new(
        Utc::now(),
        child_pid.to_string(),
        log_path.to_string_lossy().into_owned(),
        Some(platform::current_pid()),
    );

    // Get process tree information
    let record = match ProcessTreeInfo::current() {
        Ok(tree_info) => match record.clone().with_process_tree_info(tree_info) {
            Ok(updated) => updated,
            Err(err) => {
                warn(format!("Failed to attach process tree info: {}", err));
                record
            }
        },
        Err(err) => {
            warn(format!("Failed to get process tree info: {}", err));
            record
        }
    };

    if let Err(err) = registry.register(child_pid, &record) {
        platform::terminate_process(child_pid);
        let _ = child.wait().await;
        return Err(err.into());
    }

    let registration_guard = RegistrationGuard::new(registry, child_pid);
    let signal_guard = signal::install(child_pid)?;

    let status = child.wait().await?;
    drop(signal_guard);

    // Mark as completed
    let completed_at = Utc::now();
    let exit_code = status.code();
    let result = match (status.success(), exit_code) {
        (true, _) => Some("interactive_session_completed".to_owned()),
        (false, Some(code)) => Some(format!("interactive_session_failed_with_exit_code_{code}")),
        (false, None) => Some("interactive_session_failed_without_exit_code".to_owned()),
    };
    let _ = registration_guard.mark_completed(result, exit_code, completed_at);

    Ok(extract_exit_code(status))
}

/// Execute multiple CLI processes (for codex|claude|gemini syntax)
pub async fn execute_multiple_clis<S: TaskStorage>(
    registry: &Registry<S>,
    cli_selector: &crate::cli_type::CliSelector,
    task_prompt: &str,
    provider: Option<String>,
) -> Result<Vec<i32>, ProcessError> {
    let mut exit_codes = Vec::new();

    for cli_type in &cli_selector.types {
        let cli_args = cli_type.build_full_access_args(task_prompt);
        let os_args: Vec<OsString> = cli_args.into_iter().map(|s| s.into()).collect();

        let exit_code = execute_cli(registry, cli_type, &os_args, provider.clone()).await?;
        exit_codes.push(exit_code);
    }

    Ok(exit_codes)
}