use std::collections::BTreeMap;
use std::path::PathBuf;
use std::process::Stdio;
use std::time::Duration;
use futures::future::BoxFuture;
use serde_json::Value;
use tokio::io::AsyncWriteExt;
use tokio::process::Command;
use crate::error::BoxError;
use super::{HookCtx, HookError, StepHandler};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CommandSpec {
Argv {
argv: Vec<String>,
argv_windows: Option<Vec<String>>,
cwd: Option<PathBuf>,
env: BTreeMap<String, String>,
timeout_sec: Option<u64>,
},
Shell {
shell: ShellKind,
command: String,
cwd: Option<PathBuf>,
env: BTreeMap<String, String>,
timeout_sec: Option<u64>,
},
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ShellKind {
Sh,
Bash,
Pwsh,
Cmd,
Custom { program: String, args: Vec<String> },
}
impl CommandSpec {
fn timeout(&self) -> Option<Duration> {
let secs = match self {
Self::Argv { timeout_sec, .. } | Self::Shell { timeout_sec, .. } => *timeout_sec,
};
secs.map(Duration::from_secs)
}
}
pub struct CommandHandler {
spec: CommandSpec,
}
impl CommandHandler {
#[must_use]
pub fn new(spec: CommandSpec) -> Self {
Self { spec }
}
#[must_use]
pub fn timeout(&self) -> Option<Duration> {
self.spec.timeout()
}
}
impl StepHandler for CommandHandler {
fn handle_step<'a>(
&'a self,
envelope: &'a Value,
ctx: HookCtx<'a>,
) -> BoxFuture<'a, Result<Option<Value>, HookError>> {
Box::pin(async move {
let stdin_payload = serde_json::to_vec(envelope).map_err(|err| {
HookError::HandlerFailed(BoxError::new(io_invalid("serialize step envelope", err)))
})?;
let env_vars = step_env_vars(envelope, &ctx);
let mut cmd = build_command(&self.spec, &env_vars)?;
cmd.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let mut child = cmd
.spawn()
.map_err(|err| HookError::HandlerFailed(BoxError::new(err)))?;
if let Some(mut stdin) = child.stdin.take() {
let write_res = async {
stdin.write_all(&stdin_payload).await?;
stdin.write_all(b"\n").await
}
.await;
match write_res {
Ok(()) => {}
Err(err) if err.kind() == std::io::ErrorKind::BrokenPipe => {}
Err(err) => return Err(HookError::HandlerFailed(BoxError::new(err))),
}
drop(stdin);
}
let cancel = ctx.cancel.clone();
let output = tokio::select! {
() = cancel.cancelled() => return Err(HookError::Timeout),
result = child.wait_with_output() => {
result.map_err(|err| HookError::HandlerFailed(BoxError::new(err)))?
}
};
let stderr_text = String::from_utf8_lossy(&output.stderr).into_owned();
if !stderr_text.is_empty() {
tracing::debug!(target: "defect_agent::hooks::command", stderr = %stderr_text, "command stderr");
}
match output.status.code() {
Some(0) => {
let trimmed = output.stdout.trim_ascii();
if trimmed.is_empty() {
return Ok(None);
}
match serde_json::from_slice::<Value>(trimmed) {
Ok(v) => Ok(Some(v)),
Err(_) => Ok(None),
}
}
Some(2) => {
let mut obj = serde_json::Map::new();
obj.insert("control".to_string(), Value::String("veto".to_string()));
if !stderr_text.is_empty() {
obj.insert(
"additional_context".to_string(),
Value::Array(vec![Value::String(stderr_text)]),
);
}
Ok(Some(Value::Object(obj)))
}
Some(c) => Err(HookError::HandlerFailed(BoxError::new(io_invalid(
format!("hook command exited with status {c}"),
"",
)))),
None => Err(HookError::HandlerFailed(BoxError::new(io_invalid(
"hook command terminated by signal",
"",
)))),
}
})
}
}
fn build_command(
spec: &CommandSpec,
env_vars: &BTreeMap<String, String>,
) -> Result<Command, HookError> {
match spec {
CommandSpec::Argv {
argv,
argv_windows,
cwd,
env,
..
} => {
let chosen = if cfg!(target_os = "windows") {
argv_windows.as_ref().unwrap_or(argv)
} else {
argv
};
let (program, args) = chosen.split_first().ok_or_else(|| {
HookError::Configuration("command handler `argv` must not be empty".into())
})?;
let mut cmd = Command::new(program);
cmd.args(args);
if let Some(dir) = cwd {
cmd.current_dir(dir);
}
for (k, v) in env_vars {
cmd.env(k, v);
}
for (k, v) in env {
cmd.env(k, v);
}
Ok(cmd)
}
CommandSpec::Shell {
shell,
command,
cwd,
env,
..
} => {
let mut cmd = build_shell_command(shell, command);
if let Some(dir) = cwd {
cmd.current_dir(dir);
}
for (k, v) in env_vars {
cmd.env(k, v);
}
for (k, v) in env {
cmd.env(k, v);
}
Ok(cmd)
}
}
}
fn build_shell_command(shell: &ShellKind, command: &str) -> Command {
match shell {
ShellKind::Sh => {
let mut c = Command::new("sh");
c.arg("-c").arg(command);
c
}
ShellKind::Bash => {
let mut c = Command::new("bash");
c.arg("-c").arg(command);
c
}
ShellKind::Pwsh => {
let mut c = Command::new("pwsh");
c.arg("-NoProfile")
.arg("-NonInteractive")
.arg("-Command")
.arg(command);
c
}
ShellKind::Cmd => {
let mut c = Command::new("cmd");
c.arg("/C").arg(command);
c
}
ShellKind::Custom { program, args } => {
let mut c = Command::new(program);
c.args(args).arg(command);
c
}
}
}
fn step_env_vars(envelope: &Value, ctx: &HookCtx<'_>) -> BTreeMap<String, String> {
let mut out = BTreeMap::new();
out.insert(
"DEFECT_SESSION_ID".to_string(),
ctx.session_id.0.to_string(),
);
out.insert(
"DEFECT_CWD".to_string(),
ctx.cwd.to_string_lossy().into_owned(),
);
if let Some(tool) = envelope.get("tool").and_then(Value::as_str) {
out.insert("DEFECT_TOOL_NAME".to_string(), tool.to_string());
}
out
}
fn io_invalid(msg: impl Into<String>, detail: impl std::fmt::Display) -> std::io::Error {
let s = msg.into();
let body = if s.is_empty() {
detail.to_string()
} else if format!("{detail}").is_empty() {
s
} else {
format!("{s}: {detail}")
};
std::io::Error::new(std::io::ErrorKind::InvalidData, body)
}
#[cfg(test)]
mod tests;