use super::read_bounded;
use agent_workspace_contract::{
CmdOutput, CommandPermitRepository, CommandPermitRequest, CommandSpec, SecretResolver,
WorkspaceCommandExecutor, WORKSPACE_MAX_COMMAND_OUTPUT_BYTES,
};
use anyhow::{Context, Result};
use async_trait::async_trait;
use std::fmt;
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::AsyncRead;
const MAX_COMMAND_TIMEOUT_MS: u64 = 60_000;
#[derive(Clone)]
pub struct LocalHostExecutor {
permits: Arc<dyn CommandPermitRepository>,
secrets: Arc<dyn SecretResolver>,
owner: String,
root: std::path::PathBuf,
global_limit: usize,
tenant_limit: usize,
allow_shell: bool,
}
#[derive(Debug)]
pub struct EnvironmentSecretResolver;
#[async_trait]
impl SecretResolver for EnvironmentSecretResolver {
async fn resolve(
&self,
_tenant_id: &str,
secret: &agent_workspace_contract::SecretRef,
) -> Result<Vec<u8>> {
anyhow::ensure!(
secret.provider == "env",
"local executor only accepts env SecretRef"
);
anyhow::ensure!(
valid_env_name(&secret.key),
"invalid secret environment key"
);
std::env::var_os(&secret.key)
.map(|value| value.to_string_lossy().as_bytes().to_vec())
.context("secret environment variable is unavailable")
}
}
impl fmt::Debug for LocalHostExecutor {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("LocalHostExecutor")
.field("owner", &self.owner)
.field("root", &self.root)
.field("global_limit", &self.global_limit)
.field("tenant_limit", &self.tenant_limit)
.field("allow_shell", &self.allow_shell)
.finish_non_exhaustive()
}
}
impl LocalHostExecutor {
pub fn for_local_development(
permits: Arc<dyn CommandPermitRepository>,
secrets: Arc<dyn SecretResolver>,
owner: impl Into<String>,
root: std::path::PathBuf,
global_limit: usize,
tenant_limit: usize,
allow_shell: bool,
) -> Result<Self> {
anyhow::ensure!(global_limit > 0, "global command limit must be positive");
anyhow::ensure!(
tenant_limit > 0 && tenant_limit <= global_limit,
"tenant command limit must be in 1..=global limit"
);
Ok(Self {
permits,
secrets,
owner: owner.into(),
root,
global_limit,
tenant_limit,
allow_shell,
})
}
}
#[async_trait]
impl WorkspaceCommandExecutor for LocalHostExecutor {
async fn execute(
&self,
tenant_id: &str,
_workspace_id: &str,
command_id: &str,
spec: &CommandSpec,
) -> Result<CmdOutput> {
anyhow::ensure!(!tenant_id.is_empty(), "tenant id is required");
anyhow::ensure!(!spec.argv.is_empty(), "argv must not be empty");
anyhow::ensure!(spec.argv.len() <= 256, "argv item limit exceeded");
anyhow::ensure!(
spec.max_output_bytes <= WORKSPACE_MAX_COMMAND_OUTPUT_BYTES,
"command output limit exceeds policy"
);
anyhow::ensure!(
spec.timeout_ms > 0 && spec.timeout_ms <= MAX_COMMAND_TIMEOUT_MS,
"command timeout exceeds policy"
);
anyhow::ensure!(
!spec.shell || self.allow_shell,
"shell mode is not permitted"
);
let acquired = self
.permits
.acquire_command_permit(CommandPermitRequest {
command_id: command_id.into(),
tenant_id: tenant_id.into(),
owner: self.owner.clone(),
now_ms: now_ms(),
lease_ms: spec.timeout_ms.saturating_add(5_000),
global_limit: self.global_limit,
tenant_limit: self.tenant_limit,
})
.await?;
anyhow::ensure!(acquired, "workspace command quota exceeded");
let result = self.execute_acquired(tenant_id, command_id, spec).await;
let release = self
.permits
.release_command_permit(command_id, &self.owner)
.await;
match (result, release) {
(Ok(output), Ok(())) => Ok(output),
(Err(error), _) => Err(error),
(Ok(_), Err(error)) => Err(error),
}
}
async fn cancel(&self, tenant_id: &str, command_id: &str) -> Result<()> {
anyhow::ensure!(!tenant_id.is_empty(), "tenant id is required");
anyhow::ensure!(
self.permits.cancel_command_permit(command_id).await?,
"command not found"
);
Ok(())
}
}
impl LocalHostExecutor {
async fn execute_acquired(
&self,
tenant_id: &str,
command_id: &str,
spec: &CommandSpec,
) -> Result<CmdOutput> {
let cwd = spec.cwd.as_deref().unwrap_or(".");
let cwd = crate::path_resolver::WorkspacePathResolver::new(&self.root)
.existing(cwd)
.await?;
let mut command = if spec.shell {
let mut value = tokio::process::Command::new("sh");
value.arg("-lc").arg(spec.argv.join(" "));
value
} else {
let mut value = tokio::process::Command::new(&spec.argv[0]);
value.args(&spec.argv[1..]);
value
};
command
.current_dir(cwd)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
for (name, reference) in &spec.env {
anyhow::ensure!(valid_env_name(name), "invalid environment variable name");
let value = self.secrets.resolve(tenant_id, reference).await?;
#[cfg(unix)]
{
use std::os::unix::ffi::OsStringExt;
command.env(name, std::ffi::OsString::from_vec(value));
}
#[cfg(not(unix))]
command.env(
name,
String::from_utf8(value).context("secret environment value is not UTF-8")?,
);
}
#[cfg(unix)]
configure_isolation(&mut command, spec)?;
let mut child = command.spawn().context("spawn local development command")?;
let process_id = child.id().context("spawned command has no process id")?;
let stdout = child.stdout.take().context("stdout not captured")?;
let stderr = child.stderr.take().context("stderr not captured")?;
let stdout_task = tokio::spawn(read_output(stdout, spec.max_output_bytes));
let stderr_task = tokio::spawn(read_output(stderr, spec.max_output_bytes));
let timeout = tokio::time::sleep(Duration::from_millis(spec.timeout_ms));
tokio::pin!(timeout);
let status = loop {
tokio::select! {
result=child.wait()=>break result?,
_=&mut timeout=>{
kill_process_tree(process_id); let _=child.kill().await; let _=child.wait().await;
anyhow::bail!("workspace command timed out");
}
_=tokio::time::sleep(Duration::from_millis(50))=>{
if self.permits.command_canceled(command_id).await?{
kill_process_tree(process_id); let _=child.kill().await; let _=child.wait().await;
anyhow::bail!("workspace command canceled");
}
}
}
};
let (stdout, stdout_truncated) = stdout_task.await.context("join stdout")??;
let (stderr, stderr_truncated) = stderr_task.await.context("join stderr")??;
Ok(CmdOutput {
exit_code: status.code().unwrap_or(-1),
stdout: String::from_utf8_lossy(&stdout).into_owned(),
stderr: String::from_utf8_lossy(&stderr).into_owned(),
stdout_truncated,
stderr_truncated,
})
}
}
async fn read_output<R: AsyncRead + Unpin>(
reader: R,
limit: usize,
) -> std::io::Result<(Vec<u8>, bool)> {
read_bounded(reader, limit).await
}
fn valid_env_name(value: &str) -> bool {
!value.is_empty()
&& value
.bytes()
.all(|value| value == b'_' || value.is_ascii_alphanumeric())
&& !value.bytes().next().unwrap().is_ascii_digit()
}
fn now_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis()
.try_into()
.unwrap_or(u64::MAX)
}
#[cfg(unix)]
fn configure_isolation(command: &mut tokio::process::Command, spec: &CommandSpec) -> Result<()> {
use std::os::unix::process::CommandExt;
let memory = spec.memory_bytes;
let cpu_secs = spec.cpu_millis.div_ceil(1000).max(1);
let processes = u64::from(spec.max_processes.max(1));
command.as_std_mut().process_group(0);
unsafe {
command.as_std_mut().pre_exec(move || {
#[cfg(not(target_os = "macos"))]
{
set_limit(libc::RLIMIT_AS, memory)?;
set_limit(libc::RLIMIT_NPROC, processes)?;
}
#[cfg(target_os = "macos")]
let _ = (memory, processes);
set_limit(libc::RLIMIT_CPU, cpu_secs)?;
Ok(())
});
}
Ok(())
}
#[cfg(unix)]
fn set_limit(resource: RlimitResource, value: u64) -> std::io::Result<()> {
let mut current = libc::rlimit {
rlim_cur: 0,
rlim_max: 0,
};
if unsafe { libc::getrlimit(resource, &mut current) } != 0 {
return Err(std::io::Error::last_os_error());
}
let limit = libc::rlimit {
rlim_cur: (value as libc::rlim_t).min(current.rlim_max),
rlim_max: current.rlim_max,
};
if unsafe { libc::setrlimit(resource, &limit) } == 0 {
Ok(())
} else {
Err(std::io::Error::last_os_error())
}
}
#[cfg(all(unix, any(target_os = "linux", target_os = "android")))]
type RlimitResource = libc::c_uint;
#[cfg(all(unix, not(any(target_os = "linux", target_os = "android"))))]
type RlimitResource = libc::c_int;
#[cfg(unix)]
fn kill_process_tree(pid: u32) {
unsafe {
libc::kill(-(pid as i32), libc::SIGKILL);
}
}
#[cfg(not(unix))]
fn kill_process_tree(_pid: u32) {}