bzzz-core 0.1.0

//! Local Runtime Implementation
//!
//! Executes Agents and Swarms in subprocesses.
//!
//! This is a REAL executor - not a placeholder. It:
//! - Spawns actual subprocesses to execute agents
//! - Tracks process state via stored Child handles
//! - Waits for real process completion
//! - Kills processes on stop/cancel

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;

use async_trait::async_trait;
use tokio::process::Child;
use tokio::sync::{Mutex, RwLock};
use tokio::time::Instant;

use crate::{
    AgentSpec, ArtifactId, ExecutionContext, ExecutionHandle, ExecutionMetrics, ExecutionResult,
    ResourceLimits, Run, RunError, RunId, RunStatus, RuntimeAdapter, RuntimeKind, StatusResult,
    SwarmFile,
};

/// Local runtime adapter configuration
#[derive(Debug, Clone)]
pub struct LocalConfig {
    /// Base working directory
    pub base_dir: PathBuf,
    /// Default resource limits
    pub default_limits: ResourceLimits,
    /// Enable process isolation
    pub isolate_processes: bool,
}

impl Default for LocalConfig {
    fn default() -> Self {
        LocalConfig {
            base_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        }
    }
}

/// Active execution in local runtime - holds the actual subprocess handle
struct ActiveExecution {
    run_id: RunId,
    status: RunStatus,
    started_at: Instant,
    artifacts: Vec<ArtifactId>,
    /// The actual subprocess handle - None if already waited/completed
    child: Arc<Mutex<Option<Child>>>,
    /// PID for status checking (stored separately since Child might be taken)
    pid: u32,
    /// Exit status once known
    exit_status: Option<std::process::ExitStatus>,
    /// Captured stdout output, set once the process completes
    captured_output: Option<serde_json::Value>,
    /// Captured stderr (truncated to 8 KB), set once the process completes
    captured_stderr: Option<String>,
}

/// Local runtime adapter - REAL subprocess execution
pub struct LocalRuntime {
    config: LocalConfig,
    active_executions: Arc<RwLock<HashMap<String, ActiveExecution>>>,
}

impl LocalRuntime {
    /// Create a new local runtime with default config
    pub fn new() -> Self {
        LocalRuntime::with_config(LocalConfig::default())
    }

    /// Create a new local runtime with custom config
    pub fn with_config(config: LocalConfig) -> Self {
        LocalRuntime {
            config,
            active_executions: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Get the base directory
    pub fn base_dir(&self) -> &PathBuf {
        &self.config.base_dir
    }

    /// Create an execution context for a Swarm
    pub async fn create_swarm_context(
        &self,
        swarm: &SwarmFile,
    ) -> Result<ExecutionContext, RunError> {
        let ctx = ExecutionContext::new(
            format!("local-swarm-{}", swarm.id.as_str()),
            RuntimeKind::Local,
        )
        .with_working_dir(self.config.base_dir.clone())
        .with_limits(self.config.default_limits.clone());

        Ok(ctx)
    }

    /// Build a command from AgentSpec runtime config
    fn build_command(
        spec: &AgentSpec,
        working_dir: &Option<PathBuf>,
        is_background: bool,
    ) -> Option<tokio::process::Command> {
        // Get command from runtime.config, or use a default command
        let command_str = spec
            .runtime
            .config
            .get("command")
            .map(|s| s.as_str())
            .unwrap_or("echo done"); // Default to a simple no-op command

        // Parse command (simple: split by spaces, first is program, rest are args)
        let parts: Vec<&str> = command_str.split_whitespace().collect();
        if parts.is_empty() {
            return None;
        }

        let program = parts[0];
        let args = &parts[1..];

        let mut cmd = tokio::process::Command::new(program);
        cmd.args(args);

        // Pipe stdout so we can capture worker output for data flow
        cmd.stdout(std::process::Stdio::piped());

        // Pipe stderr so we can capture error output for diagnostics
        cmd.stderr(std::process::Stdio::piped());

        // Critical: kill_on_drop behavior differs for foreground vs background
        // - foreground: kill_on_drop(true) prevents orphan processes if CLI crashes
        // - background: kill_on_drop(false) lets the process survive after CLI exits
        cmd.kill_on_drop(!is_background);

        // Set working directory if specified
        if let Some(dir) = working_dir {
            cmd.current_dir(dir);
        }

        // Set environment variables from context (if any in runtime.config)
        for (key, value) in &spec.runtime.config {
            if key != "command" && key != "image" {
                cmd.env(key, value);
            }
        }

        Some(cmd)
    }

    /// Check if a process with given PID is still running
    #[cfg(unix)]
    fn is_process_running(pid: u32) -> bool {
        use std::process::Command;
        // On Unix, we can check with ps
        let output = Command::new("ps").arg("-p").arg(pid.to_string()).output();
        match output {
            Ok(o) => o.status.success(),
            Err(_) => false,
        }
    }

    #[cfg(not(unix))]
    fn is_process_running(_pid: u32) -> bool {
        // On non-Unix, we rely on Child handle status
        true // Assume running if we can't check
    }

    /// Load AgentSpec from Run target
    fn load_spec_from_run(&self, run: &Run, ctx: &ExecutionContext) -> Result<AgentSpec, RunError> {
        match &run.target {
            crate::RunTarget::Agent { spec_path } => {
                // Resolve path relative to context's working_dir (for SwarmFile execution)
                // Falls back to base_dir if working_dir not set
                let base = ctx.working_dir.as_ref().unwrap_or(&self.config.base_dir);
                let full_path = if spec_path.is_absolute() {
                    spec_path.clone()
                } else {
                    base.join(spec_path)
                };

                // Try to load from file, or create a default spec
                if full_path.exists() {
                    AgentSpec::from_yaml_file(&full_path)
                } else {
                    // Create a default spec with the path's stem as the ID
                    let name = full_path
                        .file_stem()
                        .and_then(|s| s.to_str())
                        .unwrap_or("agent");
                    Ok(AgentSpec::new(name, RuntimeKind::Local))
                }
            }
            crate::RunTarget::Swarm { swarmfile_path: _ } => {
                // For swarm execution, we'd need to load workers
                Err(RunError::InvalidConfig {
                    message: "Swarm execution requires worker loading".into(),
                })
            }
            crate::RunTarget::A2AAgent { .. } => Err(RunError::InvalidConfig {
                message: "A2A targets must use A2ARuntime, not LocalRuntime".into(),
            }),
        }
    }
}

impl Default for LocalRuntime {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl RuntimeAdapter for LocalRuntime {
    fn kind(&self) -> RuntimeKind {
        RuntimeKind::Local
    }

    async fn create(&self, spec: &AgentSpec) -> Result<ExecutionContext, RunError> {
        let ctx = ExecutionContext::new(format!("local-{}", spec.id.as_str()), RuntimeKind::Local)
            .with_working_dir(self.config.base_dir.clone())
            .with_limits(self.config.default_limits.clone());

        Ok(ctx)
    }

    async fn execute(
        &self,
        ctx: &ExecutionContext,
        run: &Run,
    ) -> Result<ExecutionHandle, RunError> {
        // For execute() (blocking), we spawn the process and wait immediately
        let spec = self.load_spec_from_run(run, ctx)?;

        // Foreground execution: kill_on_drop(true) prevents orphan processes
        let mut cmd = Self::build_command(&spec, &ctx.working_dir, false).ok_or_else(|| {
            RunError::InvalidConfig {
                message: "No 'command' specified in AgentSpec runtime.config".into(),
            }
        })?;

        // Pass input parameters via BZZZ_INPUT environment variable
        if let Some(input) = &run.input {
            if let Ok(json_str) = serde_json::to_string(input) {
                cmd.env("BZZZ_INPUT", &json_str);
            }
        }

        let started_at = Instant::now();

        // Spawn the subprocess
        let child = cmd.spawn().map_err(|e| RunError::StartupFailed {
            message: format!("Failed to spawn process: {}", e),
        })?;

        let pid = child.id().unwrap_or(0);

        // Wait for completion and capture stdout
        let output = child
            .wait_with_output()
            .await
            .map_err(|e| RunError::ExecutionFailed {
                message: format!("Failed to wait for process: {}", e),
            })?;

        let exit_status = output.status;
        let stderr_raw = String::from_utf8_lossy(&output.stderr);
        let stderr_truncated = super::truncate_stderr(&stderr_raw).to_string();
        let captured_stderr = if stderr_truncated.trim().is_empty() {
            None
        } else {
            Some(stderr_truncated.clone())
        };
        let captured_output = super::parse_worker_output(
            &String::from_utf8_lossy(&output.stdout),
            &stderr_truncated,
        );

        let _elapsed = started_at.elapsed();
        let status = if exit_status.success() {
            RunStatus::Completed
        } else {
            RunStatus::Failed
        };

        // Store completed execution for potential status queries
        let execution = ActiveExecution {
            run_id: run.id.clone(),
            status,
            started_at,
            artifacts: Vec::new(),
            child: Arc::new(Mutex::new(None)), // Child already consumed
            pid,
            exit_status: Some(exit_status),
            captured_output,
            captured_stderr,
        };

        {
            let mut executions = self.active_executions.write().await;
            executions.insert(run.id.as_str().to_string(), execution);
        }

        Ok(ExecutionHandle::new(
            run.id.clone(),
            RuntimeKind::Local,
            format!("pid:{}", pid),
        ))
    }

    async fn execute_background(
        &self,
        ctx: &ExecutionContext,
        run: &Run,
    ) -> Result<ExecutionHandle, RunError> {
        // REAL background execution - spawn subprocess and track it
        // kill_on_drop(false) lets the process survive after CLI exits
        let spec = self.load_spec_from_run(run, ctx)?;

        let mut cmd = Self::build_command(&spec, &ctx.working_dir, true).ok_or_else(|| {
            RunError::InvalidConfig {
                message: "No 'command' specified in AgentSpec runtime.config".into(),
            }
        })?;

        // Pass input parameters via BZZZ_INPUT environment variable
        if let Some(input) = &run.input {
            if let Ok(json_str) = serde_json::to_string(input) {
                cmd.env("BZZZ_INPUT", &json_str);
            }
        }

        let started_at = Instant::now();

        // Spawn the subprocess - it runs independently
        let child = cmd.spawn().map_err(|e| RunError::StartupFailed {
            message: format!("Failed to spawn process: {}", e),
        })?;

        let pid = child.id().unwrap_or(0);

        // Store the execution with the live Child handle
        let execution = ActiveExecution {
            run_id: run.id.clone(),
            status: RunStatus::Running,
            started_at,
            artifacts: Vec::new(),
            child: Arc::new(Mutex::new(Some(child))),
            pid,
            exit_status: None,
            captured_output: None,
            captured_stderr: None,
        };

        {
            let mut executions = self.active_executions.write().await;
            executions.insert(run.id.as_str().to_string(), execution);
        }

        Ok(ExecutionHandle::new(
            run.id.clone(),
            RuntimeKind::Local,
            format!("pid:{}", pid),
        ))
    }

    async fn status(&self, handle: &ExecutionHandle) -> Result<StatusResult, RunError> {
        // Get the data we need without holding the lock for long
        let (run_id, status, started_at, artifacts, child_arc, pid, exit_status_opt) = {
            let executions = self.active_executions.read().await;

            match executions.get(handle.run_id.as_str()) {
                Some(execution) => (
                    execution.run_id.clone(),
                    execution.status,
                    execution.started_at,
                    execution.artifacts.clone(),
                    execution.child.clone(),
                    execution.pid,
                    execution.exit_status,
                ),
                None => {
                    return Err(RunError::InvalidConfig {
                        message: format!("Execution not found: {}", handle.run_id.as_str()),
                    });
                }
            }
            // Read lock released here
        };

        // Determine actual status by checking process state
        let actual_status = if exit_status_opt.is_some() {
            // Already completed - use stored status
            status
        } else if status == RunStatus::Cancelled {
            // Was cancelled
            RunStatus::Cancelled
        } else {
            // Check if process is still running
            let child_guard = child_arc.lock().await;
            if child_guard.is_some() {
                // Still have the handle - check with PID
                if Self::is_process_running(pid) {
                    RunStatus::Running
                } else {
                    // Process exited but we haven't waited yet
                    RunStatus::Running // Will be updated on wait()
                }
            } else {
                // Child handle was taken (by wait())
                status
            }
        };

        Ok(StatusResult {
            run_id,
            status: actual_status,
            current_step: None,
            progress: 0,
            elapsed_ms: started_at.elapsed().as_millis() as u64,
            artifacts,
        })
    }

    async fn destroy(&self, _ctx: &ExecutionContext) -> Result<(), RunError> {
        // Local runtime doesn't need context cleanup
        Ok(())
    }

    async fn cancel(&self, handle: &ExecutionHandle) -> Result<(), RunError> {
        // Get the child handle without holding the read lock
        let child_arc = {
            let executions = self.active_executions.read().await;

            match executions.get(handle.run_id.as_str()) {
                Some(execution) => execution.child.clone(),
                None => {
                    return Err(RunError::InvalidConfig {
                        message: format!("Execution not found: {}", handle.run_id.as_str()),
                    });
                }
            }
            // Read lock released here
        };

        // Now lock the child mutex and kill the process
        let mut child_guard = child_arc.lock().await;
        if let Some(mut child) = child_guard.take() {
            // Kill the subprocess
            child.kill().await.map_err(|e| RunError::RuntimeError {
                message: format!("Failed to kill process: {}", e),
            })?;

            // Wait for the killed process to finish reaping
            // This ensures the process is truly dead and we get the exit status
            let exit_status = child.wait().await.ok();

            drop(child_guard);

            // Update status to Cancelled (not Failed - this was a user-initiated stop)
            let mut executions = self.active_executions.write().await;
            if let Some(exec) = executions.get_mut(handle.run_id.as_str()) {
                exec.status = RunStatus::Cancelled;
                exec.exit_status = exit_status;
            }
        }
        Ok(())
    }

    async fn wait(&self, handle: &ExecutionHandle) -> Result<ExecutionResult, RunError> {
        // First, get the child handle without holding the lock for too long
        let (run_id, started_at, artifacts, child_arc) = {
            let executions = self.active_executions.read().await;

            match executions.get(handle.run_id.as_str()) {
                Some(execution) => {
                    (
                        execution.run_id.clone(),
                        execution.started_at,
                        execution.artifacts.clone(),
                        execution.child.clone(), // Clone the Arc<Mutex>
                    )
                }
                None => {
                    return Err(RunError::InvalidConfig {
                        message: format!("Execution not found: {}", handle.run_id.as_str()),
                    });
                }
            }
            // Read lock released here
        };

        // Now lock the child mutex and wait
        let mut child_guard = child_arc.lock().await;

        if let Some(child) = child_guard.take() {
            // We have the child handle - wait for it and capture stdout
            let output = child
                .wait_with_output()
                .await
                .map_err(|e| RunError::ExecutionFailed {
                    message: format!("Failed to wait for process: {}", e),
                })?;

            let exit_status = output.status;
            let stderr_raw = String::from_utf8_lossy(&output.stderr);
            let stderr_truncated = super::truncate_stderr(&stderr_raw).to_string();
            let captured_stderr = if stderr_truncated.trim().is_empty() {
                None
            } else {
                Some(stderr_truncated.clone())
            };
            let captured_output = super::parse_worker_output(
                &String::from_utf8_lossy(&output.stdout),
                &stderr_truncated,
            );
            let elapsed = started_at.elapsed();
            let status = if exit_status.success() {
                RunStatus::Completed
            } else {
                RunStatus::Failed
            };

            // Update execution state
            {
                let mut executions = self.active_executions.write().await;
                if let Some(exec) = executions.get_mut(handle.run_id.as_str()) {
                    exec.status = status;
                    exec.exit_status = Some(exit_status);
                    exec.captured_output = captured_output.clone();
                    exec.captured_stderr = captured_stderr.clone();
                }
            }

            let error = if status == RunStatus::Failed {
                let base_msg = format!(
                    "Process exited with code: {}",
                    exit_status.code().unwrap_or(-1)
                );
                let msg = if let Some(ref stderr) = captured_stderr {
                    format!("{}
stderr: {}", base_msg, stderr)
                } else {
                    base_msg
                };
                Some(RunError::ExecutionFailed { message: msg })
            } else {
                None
            };

            Ok(ExecutionResult {
                run_id,
                status,
                artifacts,
                error,
                metrics: ExecutionMetrics {
                    wall_time_ms: elapsed.as_millis() as u64,
                    cpu_time_ms: 0,
                    peak_memory_bytes: 0,
                    retries: 0,
                    selection_trace: None,
                },
                output: captured_output,
            })
        } else {
            // Child already consumed (foreground execute()) - return stored result
            let executions = self.active_executions.read().await;
            match executions.get(handle.run_id.as_str()) {
                Some(execution) => {
                    let elapsed = execution.started_at.elapsed();

                    let error = if execution.status == RunStatus::Failed {
                        let base_msg = format!(
                            "Process exited with code: {}",
                            execution.exit_status.and_then(|s| s.code()).unwrap_or(-1)
                        );
                        let msg = if let Some(ref stderr) = execution.captured_stderr {
                            format!("{}
stderr: {}", base_msg, stderr)
                        } else {
                            base_msg
                        };
                        Some(RunError::ExecutionFailed { message: msg })
                    } else if execution.status == RunStatus::Cancelled {
                        Some(RunError::Cancelled {
                            reason: "Process was cancelled by user".into(),
                        })
                    } else {
                        None
                    };

                    Ok(ExecutionResult {
                        run_id: execution.run_id.clone(),
                        status: execution.status,
                        artifacts: execution.artifacts.clone(),
                        error,
                        metrics: ExecutionMetrics {
                            wall_time_ms: elapsed.as_millis() as u64,
                            cpu_time_ms: 0,
                            peak_memory_bytes: 0,
                            retries: 0,
                            selection_trace: None,
                        },
                        output: execution.captured_output.clone(),
                    })
                }
                None => Err(RunError::InvalidConfig {
                    message: format!("Execution not found: {}", handle.run_id.as_str()),
                }),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::time::{timeout, Duration};

    #[test]
    fn test_local_runtime_creation() {
        let runtime = LocalRuntime::new();
        assert_eq!(runtime.kind(), RuntimeKind::Local);
    }

    #[test]
    fn test_local_runtime_config() {
        let config = LocalConfig {
            base_dir: PathBuf::from("/tmp/test"),
            default_limits: ResourceLimits {
                max_cpu_cores: 2,
                max_memory_bytes: 1024 * 1024,
                max_time_seconds: 60,
            },
            isolate_processes: true,
        };

        let runtime = LocalRuntime::with_config(config);
        assert_eq!(runtime.base_dir(), &PathBuf::from("/tmp/test"));
    }

    #[tokio::test]
    async fn test_local_runtime_create_context() {
        let runtime = LocalRuntime::new();
        let spec = AgentSpec::new("test-agent", RuntimeKind::Local);

        let ctx = runtime.create(&spec).await.unwrap();
        assert!(ctx.id.starts_with("local-"));
        assert_eq!(ctx.runtime_kind, RuntimeKind::Local);
    }

    /// Test helper: create an AgentSpec YAML file with a command
    fn create_test_agent_spec(dir: &std::path::Path, name: &str, command: &str) -> AgentSpec {
        let mut spec = AgentSpec::new(name, RuntimeKind::Local);
        spec.runtime.config.insert("command".into(), command.into());
        spec.spec_path = Some(dir.join(format!("{}.yaml", name)));

        // Write the YAML
        spec.to_yaml_file(&spec.spec_path.clone().unwrap()).unwrap();
        spec
    }

    #[tokio::test]
    async fn test_real_execute_success() {
        // Create temp directory for test
        let temp_dir = std::env::temp_dir();
        let test_dir = temp_dir.join("bzzz-test-real-execute");
        std::fs::create_dir_all(&test_dir).unwrap();

        // Create agent spec that runs "echo hello"
        let spec = create_test_agent_spec(&test_dir, "echo-agent", "echo hello");

        let config = LocalConfig {
            base_dir: test_dir.clone(),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        };
        let runtime = LocalRuntime::with_config(config);

        let ctx = runtime.create(&spec).await.unwrap();
        let run = Run::new(
            crate::RunTarget::Agent {
                spec_path: spec.spec_path.clone().unwrap(),
            },
            RuntimeKind::Local,
        );

        // Execute should run the command and wait
        let handle = runtime.execute(&ctx, &run).await.unwrap();
        assert_eq!(handle.runtime_kind, RuntimeKind::Local);

        // Check status - should be Completed
        let status = runtime.status(&handle).await.unwrap();
        assert_eq!(status.status, RunStatus::Completed);

        // Cleanup
        std::fs::remove_dir_all(&test_dir).ok();
    }

    #[tokio::test]
    async fn test_real_execute_failure() {
        let temp_dir = std::env::temp_dir();
        let test_dir = temp_dir.join("bzzz-test-real-execute-fail");
        std::fs::create_dir_all(&test_dir).unwrap();

        // Create agent spec that runs a failing command
        // Using "false" which always exits with 1
        let spec = create_test_agent_spec(&test_dir, "fail-agent", "false");

        let config = LocalConfig {
            base_dir: test_dir.clone(),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        };
        let runtime = LocalRuntime::with_config(config);

        let ctx = runtime.create(&spec).await.unwrap();
        let run = Run::new(
            crate::RunTarget::Agent {
                spec_path: spec.spec_path.clone().unwrap(),
            },
            RuntimeKind::Local,
        );

        // Execute should run the failing command
        let handle = runtime.execute(&ctx, &run).await.unwrap();

        // Status should be Failed
        let status = runtime.status(&handle).await.unwrap();
        assert_eq!(status.status, RunStatus::Failed);

        // Cleanup
        std::fs::remove_dir_all(&test_dir).ok();
    }

    #[tokio::test]
    async fn test_real_background_execute() {
        let temp_dir = std::env::temp_dir();
        let test_dir = temp_dir.join("bzzz-test-real-bg");
        std::fs::create_dir_all(&test_dir).unwrap();

        // Create agent spec that runs a longer command
        let spec = create_test_agent_spec(&test_dir, "sleep-agent", "sleep 0.1");

        let config = LocalConfig {
            base_dir: test_dir.clone(),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        };
        let runtime = LocalRuntime::with_config(config);

        let ctx = runtime.create(&spec).await.unwrap();
        let run = Run::new(
            crate::RunTarget::Agent {
                spec_path: spec.spec_path.clone().unwrap(),
            },
            RuntimeKind::Local,
        );

        // Execute in background - process should start
        let handle = runtime.execute_background(&ctx, &run).await.unwrap();
        assert_eq!(handle.runtime_kind, RuntimeKind::Local);

        // Immediately check status - should be Running
        let status = runtime.status(&handle).await.unwrap();
        // Might be Running or Completed depending on timing
        assert!(status.status == RunStatus::Running || status.status == RunStatus::Completed);

        // Wait for completion
        let result = runtime.wait(&handle).await.unwrap();
        assert_eq!(result.status, RunStatus::Completed);

        // Cleanup
        std::fs::remove_dir_all(&test_dir).ok();
    }

    #[tokio::test]
    async fn test_real_cancel() {
        let temp_dir = std::env::temp_dir();
        let test_dir = temp_dir.join("bzzz-test-real-cancel");
        std::fs::create_dir_all(&test_dir).unwrap();

        // Create agent spec that sleeps for a long time (but we'll cancel it quickly)
        let spec = create_test_agent_spec(&test_dir, "long-sleep-agent", "sleep 30");

        let config = LocalConfig {
            base_dir: test_dir.clone(),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        };
        let runtime = LocalRuntime::with_config(config);

        let ctx = runtime.create(&spec).await.unwrap();
        let run = Run::new(
            crate::RunTarget::Agent {
                spec_path: spec.spec_path.clone().unwrap(),
            },
            RuntimeKind::Local,
        );

        // Execute in background
        let handle = runtime.execute_background(&ctx, &run).await.unwrap();

        // Check it's running
        let status = runtime.status(&handle).await.unwrap();
        assert_eq!(status.status, RunStatus::Running);

        // Cancel it - REAL kill
        runtime.cancel(&handle).await.unwrap();

        // Wait should now complete quickly (process was killed)
        let result = timeout(Duration::from_secs(2), runtime.wait(&handle))
            .await
            .expect("wait should complete quickly after cancel")
            .unwrap();

        // The process was cancelled, so it should have Cancelled status
        assert_eq!(result.status, RunStatus::Cancelled);
        assert!(result.error.is_some());

        // Cleanup
        std::fs::remove_dir_all(&test_dir).ok();
    }

    #[tokio::test]
    async fn test_foreground_captures_json_stdout() {
        let temp_dir = std::env::temp_dir();
        let test_dir = temp_dir.join("bzzz-test-fg-json");
        std::fs::create_dir_all(&test_dir).unwrap();

        // Command that outputs valid JSON
        let spec = create_test_agent_spec(&test_dir, "json-agent", r#"echo {"key":"value"}"#);

        let config = LocalConfig {
            base_dir: test_dir.clone(),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        };
        let runtime = LocalRuntime::with_config(config);
        let ctx = runtime.create(&spec).await.unwrap();
        let run = Run::new(
            crate::RunTarget::Agent {
                spec_path: spec.spec_path.clone().unwrap(),
            },
            RuntimeKind::Local,
        );

        let handle = runtime.execute(&ctx, &run).await.unwrap();
        let result = runtime.wait(&handle).await.unwrap();

        assert_eq!(result.status, RunStatus::Completed);
        assert!(result.output.is_some());

        std::fs::remove_dir_all(&test_dir).ok();
    }

    #[tokio::test]
    async fn test_background_captures_stdout() {
        let temp_dir = std::env::temp_dir();
        let test_dir = temp_dir.join("bzzz-test-bg-stdout");
        std::fs::create_dir_all(&test_dir).unwrap();

        // Command that outputs plain text
        let spec = create_test_agent_spec(&test_dir, "plain-agent", "echo hello");

        let config = LocalConfig {
            base_dir: test_dir.clone(),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        };
        let runtime = LocalRuntime::with_config(config);
        let ctx = runtime.create(&spec).await.unwrap();
        let run = Run::new(
            crate::RunTarget::Agent {
                spec_path: spec.spec_path.clone().unwrap(),
            },
            RuntimeKind::Local,
        );

        let handle = runtime.execute_background(&ctx, &run).await.unwrap();
        let result = runtime.wait(&handle).await.unwrap();

        assert_eq!(result.status, RunStatus::Completed);
        assert!(result.output.is_some());
        assert_eq!(
            result.output.unwrap(),
            serde_json::json!({"stdout": "hello"})
        );

        std::fs::remove_dir_all(&test_dir).ok();
    }

    #[tokio::test]
    async fn test_real_wait_timeout() {
        let temp_dir = std::env::temp_dir();
        let test_dir = temp_dir.join("bzzz-test-real-wait");
        std::fs::create_dir_all(&test_dir).unwrap();

        // Create agent spec
        let spec = create_test_agent_spec(&test_dir, "quick-agent", "echo test");

        let config = LocalConfig {
            base_dir: test_dir.clone(),
            default_limits: ResourceLimits::default(),
            isolate_processes: false,
        };
        let runtime = LocalRuntime::with_config(config);

        let ctx = runtime.create(&spec).await.unwrap();
        let run = Run::new(
            crate::RunTarget::Agent {
                spec_path: spec.spec_path.clone().unwrap(),
            },
            RuntimeKind::Local,
        );

        // Execute in background
        let handle = runtime.execute_background(&ctx, &run).await.unwrap();

        // Wait should complete quickly
        let result = timeout(Duration::from_secs(5), runtime.wait(&handle))
            .await
            .expect("wait should not timeout")
            .unwrap();

        assert_eq!(result.status, RunStatus::Completed);

        // Cleanup
        std::fs::remove_dir_all(&test_dir).ok();
    }
}