IdeA/crates/infrastructure/src/workspace_tool.rs

//! Bounded, sandboxed process adapter for the plugin workspace execution SDK.

use std::io::{Read, Write};
use std::process::{Command, Stdio};
use std::sync::Arc;
use std::time::{Duration, Instant};

use async_trait::async_trait;
use domain::ports::{
    PluginWorkspaceExecAuditSink, WorkspaceToolExecError, WorkspaceToolExecOutput,
    WorkspaceToolExecSpec, WorkspaceToolRunner,
};
use domain::sandbox::SandboxEnforcer;
use domain::PluginWorkspaceExecAuditEntry;

const MAX_CONCURRENT_WORKSPACE_TOOLS: usize = 4;

/// Local process runner with hard concurrency, time, output and sandbox bounds.
pub struct LocalWorkspaceToolRunner {
    sandbox: Arc<dyn SandboxEnforcer>,
    permits: Arc<tokio::sync::Semaphore>,
}

impl LocalWorkspaceToolRunner {
    /// Builds the adapter with the selected platform sandbox enforcer.
    #[must_use]
    pub fn new(sandbox: Arc<dyn SandboxEnforcer>) -> Self {
        Self {
            sandbox,
            permits: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_WORKSPACE_TOOLS)),
        }
    }
}

#[async_trait]
impl WorkspaceToolRunner for LocalWorkspaceToolRunner {
    async fn run(
        &self,
        spec: WorkspaceToolExecSpec,
    ) -> Result<WorkspaceToolExecOutput, WorkspaceToolExecError> {
        let _permit = Arc::clone(&self.permits)
            .acquire_owned()
            .await
            .map_err(|_| WorkspaceToolExecError::Execution("runner is shutting down".to_owned()))?;
        let sandbox = Arc::clone(&self.sandbox);
        let (tx, rx) = tokio::sync::oneshot::channel();
        std::thread::spawn(move || {
            let result = run_blocking(spec, sandbox.as_ref());
            let _ = tx.send(result);
        });
        rx.await.map_err(|_| {
            WorkspaceToolExecError::Execution("workspace tool runner thread stopped".to_owned())
        })?
    }
}

fn run_blocking(
    spec: WorkspaceToolExecSpec,
    sandbox: &dyn SandboxEnforcer,
) -> Result<WorkspaceToolExecOutput, WorkspaceToolExecError> {
    let started = Instant::now();
    sandbox
        .enforce(&spec.sandbox)
        .map_err(|error| WorkspaceToolExecError::Execution(format!("sandbox: {error}")))?;

    let mut command = Command::new(&spec.command);
    command
        .args(&spec.args)
        .current_dir(spec.cwd.as_str())
        .env_clear()
        .envs(spec.env.iter().cloned())
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    let mut child = command
        .spawn()
        .map_err(|error| WorkspaceToolExecError::Execution(format!("{}: {error}", spec.command)))?;

    let stdout = child
        .stdout
        .take()
        .ok_or_else(|| WorkspaceToolExecError::Execution("stdout pipe unavailable".to_owned()))?;
    let stderr = child
        .stderr
        .take()
        .ok_or_else(|| WorkspaceToolExecError::Execution("stderr pipe unavailable".to_owned()))?;
    let stdout_limit = spec.max_stdout_bytes;
    let stderr_limit = spec.max_stderr_bytes;
    let stdout_reader = std::thread::spawn(move || read_bounded(stdout, stdout_limit));
    let stderr_reader = std::thread::spawn(move || read_bounded(stderr, stderr_limit));

    let stdin_writer = match spec.stdin {
        Some(stdin) => {
            let mut pipe = child.stdin.take().ok_or_else(|| {
                WorkspaceToolExecError::Execution("stdin pipe unavailable".to_owned())
            })?;
            Some(std::thread::spawn(move || pipe.write_all(stdin.as_bytes())))
        }
        None => {
            drop(child.stdin.take());
            None
        }
    };

    let deadline = Duration::from_millis(spec.timeout_ms);
    let (status, timed_out) = loop {
        if let Some(status) = child
            .try_wait()
            .map_err(|error| WorkspaceToolExecError::Execution(error.to_string()))?
        {
            break (status, false);
        }
        if started.elapsed() >= deadline {
            child
                .kill()
                .map_err(|error| WorkspaceToolExecError::Execution(error.to_string()))?;
            let status = child
                .wait()
                .map_err(|error| WorkspaceToolExecError::Execution(error.to_string()))?;
            break (status, true);
        }
        std::thread::sleep(Duration::from_millis(5));
    };

    let (stdout, stdout_truncated) = stdout_reader.join().map_err(|_| {
        WorkspaceToolExecError::Execution("stdout reader thread panicked".to_owned())
    })??;
    let (stderr, stderr_truncated) = stderr_reader.join().map_err(|_| {
        WorkspaceToolExecError::Execution("stderr reader thread panicked".to_owned())
    })??;
    if let Some(writer) = stdin_writer {
        writer
            .join()
            .map_err(|_| {
                WorkspaceToolExecError::Execution("stdin writer thread panicked".to_owned())
            })?
            .or_else(|error| {
                if error.kind() == std::io::ErrorKind::BrokenPipe {
                    Ok(())
                } else {
                    Err(error)
                }
            })
            .map_err(|error| WorkspaceToolExecError::Execution(error.to_string()))?;
    }

    Ok(WorkspaceToolExecOutput {
        exit_code: status.code(),
        signal: exit_signal(&status),
        timed_out,
        stdout,
        stderr,
        stdout_truncated,
        stderr_truncated,
        duration_ms: u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX),
    })
}

fn read_bounded(
    mut reader: impl Read,
    limit: usize,
) -> Result<(Vec<u8>, bool), WorkspaceToolExecError> {
    let mut retained = Vec::with_capacity(limit.min(8 * 1024));
    let mut truncated = false;
    let mut chunk = [0_u8; 8 * 1024];
    loop {
        let read = reader
            .read(&mut chunk)
            .map_err(|error| WorkspaceToolExecError::Execution(error.to_string()))?;
        if read == 0 {
            break;
        }
        let remaining = limit.saturating_sub(retained.len());
        let keep = remaining.min(read);
        retained.extend_from_slice(&chunk[..keep]);
        truncated |= keep < read;
    }
    Ok((retained, truncated))
}

#[cfg(unix)]
fn exit_signal(status: &std::process::ExitStatus) -> Option<String> {
    use std::os::unix::process::ExitStatusExt;
    status.signal().map(|signal| signal.to_string())
}

#[cfg(not(unix))]
fn exit_signal(_status: &std::process::ExitStatus) -> Option<String> {
    None
}

/// Structured diagnostic audit sink that never serializes captured output.
#[derive(Debug, Default, Clone, Copy)]
pub struct DiagnosticPluginWorkspaceExecAudit;

impl PluginWorkspaceExecAuditSink for DiagnosticPluginWorkspaceExecAudit {
    fn record(&self, entry: PluginWorkspaceExecAuditEntry) {
        match serde_json::to_string(&entry) {
            Ok(json) => eprintln!("[plugin-workspace-exec] {json}"),
            Err(error) => eprintln!("[plugin-workspace-exec] audit serialization failed: {error}"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use domain::permission::Posture;
    use domain::sandbox::SandboxPlan;
    use domain::ProjectPath;

    #[test]
    fn bounded_reader_drains_but_retains_only_prefix() {
        let (bytes, truncated) = read_bounded(&b"abcdef"[..], 3).unwrap();
        assert_eq!(bytes, b"abc");
        assert!(truncated);
    }

    fn spec(args: &[&str], timeout_ms: u64, stdout_limit: usize) -> WorkspaceToolExecSpec {
        WorkspaceToolExecSpec {
            command: "sh".to_owned(),
            args: args.iter().map(|value| (*value).to_owned()).collect(),
            cwd: ProjectPath::new("/").unwrap(),
            env: vec![("PATH".to_owned(), "/usr/bin:/bin".to_owned())],
            stdin: None,
            timeout_ms,
            max_stdout_bytes: stdout_limit,
            max_stderr_bytes: 32,
            sandbox: SandboxPlan {
                allowed: Vec::new(),
                default_posture: Posture::Allow,
            },
        }
    }

    #[tokio::test]
    async fn local_runner_times_out_and_reaps_process() {
        let runner = LocalWorkspaceToolRunner::new(Arc::new(crate::sandbox::NoopSandbox));
        let result = runner.run(spec(&["-c", "sleep 1"], 20, 32)).await.unwrap();
        assert!(result.timed_out);
        assert!(result.duration_ms < 1_000);
    }

    #[tokio::test]
    async fn local_runner_timeout_is_not_blocked_by_unread_stdin() {
        let runner = LocalWorkspaceToolRunner::new(Arc::new(crate::sandbox::NoopSandbox));
        let mut request = spec(&["-c", "sleep 1"], 20, 32);
        request.stdin = Some("x".repeat(64_000));
        let result = runner.run(request).await.unwrap();
        assert!(result.timed_out);
        assert!(result.duration_ms < 1_000);
    }

    #[tokio::test]
    async fn local_runner_keeps_streams_separate_and_truncates_output() {
        let runner = LocalWorkspaceToolRunner::new(Arc::new(crate::sandbox::NoopSandbox));
        let result = runner
            .run(spec(&["-c", "printf 123456; printf err >&2"], 1_000, 3))
            .await
            .unwrap();
        assert_eq!(result.stdout, b"123");
        assert_eq!(result.stderr, b"err");
        assert!(result.stdout_truncated);
        assert!(!result.stderr_truncated);
    }
}