aion_worker/shell/

action.rs

//! A declared command executed as an activity.
//!
//! [`ShellAction`] turns a command declared in an `.awl` worker block into a
//! runnable activity, so an action can do real work without a hand-written
//! Rust crate. The command is parsed once at construction and executed as a
//! direct `execve` with no shell interposed — see [`super::template`] for why
//! that is the security property rather than a convenience.

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};
use tokio::process::Command;

use super::template::{CommandTemplate, SubstitutionError, TemplateError};
use crate::activity::ActivityFailure;
use crate::command_transcript::CommandTranscript;
use crate::context::ActivityContext;
use crate::process::{CancellableCommandOutput, ProcessGroupError, run_cancellable_command};

/// What a declared action produces when its command succeeds.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
pub struct ShellOutcome {
    /// The command's exit status code.
    pub exit_code: i32,
    /// Everything the command wrote to standard output, trailing newline
    /// trimmed so a one-line result is usable directly.
    pub stdout: String,
    /// Everything the command wrote to standard error, trailing newline
    /// trimmed. Captured even on success, because a command that succeeds
    /// while warning is a thing an operator needs to see.
    pub stderr: String,
}

/// How a declared action treats a non-zero exit.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum FailureMode {
    /// A non-zero exit is retryable: the engine may run the command again
    /// under the action's retry policy. This is the default because the
    /// common causes of a failing command — a busy resource, an unreachable
    /// host, a transient permission state — are the ones that succeed on a
    /// second attempt.
    #[default]
    Retryable,
    /// A non-zero exit is terminal: the command will never be retried. Use
    /// when re-running the command could not possibly change the outcome, or
    /// when running it twice would itself be harmful.
    Terminal,
}

/// The one environment variable a declared command inherits from its host.
///
/// A declared command runs with the host's environment CLEARED. That is a
/// security boundary, not tidiness: the host process legitimately holds
/// credentials — store URLs, signing keys, cloud tokens — and a deployed
/// package must never be able to read them by declaring `run "env"`. Nothing
/// crosses unless an operator names it.
///
/// `PATH` is the sole exception, and only because resolving a bare program
/// name requires it: `run "echo hi"` must find `echo`. Its value is the
/// host's, unmodified. An action that will not tolerate even that names its
/// program by absolute path and passes an environment that overrides `PATH`
/// to whatever it wants.
pub(super) const INHERITED_VARIABLE: &str = "PATH";

/// A declared command, ready to run as an activity.
#[derive(Debug, Clone)]
pub struct ShellAction {
    template: CommandTemplate,
    failure_mode: FailureMode,
    environment: BTreeMap<String, String>,
    working_directory: Option<std::path::PathBuf>,
}

impl ShellAction {
    /// Parse `command` into a runnable action.
    ///
    /// # Errors
    ///
    /// Returns [`TemplateError`] when the command is empty, has an unterminated
    /// quote or `${`, contains a `$` naming no parameter, or puts a parameter
    /// in program position.
    pub fn new(command: &str) -> Result<Self, TemplateError> {
        Ok(Self {
            template: CommandTemplate::parse(command)?,
            failure_mode: FailureMode::default(),
            environment: BTreeMap::new(),
            working_directory: None,
        })
    }

    /// Set how a non-zero exit is classified.
    #[must_use]
    pub const fn with_failure_mode(mut self, failure_mode: FailureMode) -> Self {
        self.failure_mode = failure_mode;
        self
    }

    /// Supply the environment the command runs with.
    ///
    /// Additive to nothing: the host's environment is cleared regardless, so
    /// these entries plus [`INHERITED_VARIABLE`] are the command's whole
    /// environment. An entry named `PATH` here replaces the inherited one.
    #[must_use]
    pub fn with_environment(mut self, environment: BTreeMap<String, String>) -> Self {
        self.environment = environment;
        self
    }

    /// Run the command in `directory` instead of wherever the host happens to
    /// be.
    ///
    /// Worth setting whenever a command names a relative path. Absent, the
    /// command inherits the host process's working directory, which is a
    /// property of how the host was launched and not something a package
    /// author can see — so a relative path in a declared command means
    /// something the author cannot predict.
    #[must_use]
    pub fn with_working_directory(mut self, directory: impl Into<std::path::PathBuf>) -> Self {
        self.working_directory = Some(directory.into());
        self
    }

    /// The parameter names this command references, sorted.
    ///
    /// A caller can compare these against an action's declared parameters to
    /// catch a reference to a parameter that does not exist.
    #[must_use]
    pub fn referenced_parameters(&self) -> Vec<String> {
        self.template.referenced_parameters()
    }

    /// Run the command with `arguments` bound to its parameter references.
    ///
    /// Cancellation is tree-wide: the command runs in its own process group,
    /// and cancelling the activity terminates that whole group rather than
    /// only the direct child, so a command that spawned children does not
    /// leave them running.
    ///
    /// Output is streamed onto `context`'s transcript seam line by line as the
    /// command writes it (see [`CommandTranscript`]), so an operator can read
    /// what a command step is printing mid-run. The returned [`ShellOutcome`]
    /// is unaffected: it still carries the command's complete stdout and stderr.
    ///
    /// # Errors
    ///
    /// Returns an [`ActivityFailure`] when a referenced parameter is missing or
    /// unrepresentable (terminal — re-running cannot supply it), when the
    /// command cannot be spawned or observed (terminal — the program is absent
    /// or unrunnable), when the command is cancelled (terminal), or when it
    /// exits non-zero (classified by the action's [`FailureMode`]).
    pub async fn run(
        &self,
        arguments: &BTreeMap<String, serde_json::Value>,
        context: &ActivityContext,
    ) -> Result<ShellOutcome, ActivityFailure> {
        let argv = self
            .template
            .render(arguments)
            .map_err(|error| substitution_failure(&error))?;
        let (program, rest) = argv
            .split_first()
            // `CommandTemplate::parse` refuses an empty command, so a rendered
            // argv always has a program. Handled rather than indexed so a
            // future change to the parser cannot turn this into a panic.
            .ok_or_else(|| ActivityFailure::terminal("the declared command rendered no program"))?;

        let mut command = Command::new(program);
        command.args(rest);
        // stdin is CLOSED, never inherited. A declared command that reads
        // standard input would otherwise be reading the host process's — a
        // terminal on an interactive boot, or another activity's leftovers —
        // and would block the activity forever waiting for a line nobody is
        // going to type.
        command.stdin(std::process::Stdio::null());
        // The host's environment does not cross into a declared command. See
        // `INHERITED_VARIABLE`: only PATH does, and only so a bare program
        // name resolves.
        command.env_clear();
        if let Some(path) = std::env::var_os(INHERITED_VARIABLE) {
            command.env(INHERITED_VARIABLE, path);
        }
        for (name, value) in &self.environment {
            command.env(name, value);
        }
        if let Some(directory) = &self.working_directory {
            command.current_dir(directory);
        }

        let transcript = CommandTranscript::new(context);
        match run_cancellable_command(command, context.cancelled(), &transcript).await {
            Ok(CancellableCommandOutput::Completed(output)) => {
                let outcome = ShellOutcome {
                    exit_code: output.status.code().unwrap_or(EXIT_CODE_SIGNALLED),
                    stdout: trim_trailing_newline(&String::from_utf8_lossy(&output.stdout)),
                    stderr: trim_trailing_newline(&String::from_utf8_lossy(&output.stderr)),
                };
                if output.status.success() {
                    Ok(outcome)
                } else {
                    Err(self.exit_failure(program, &outcome))
                }
            }
            Ok(CancellableCommandOutput::Cancelled) => Err(ActivityFailure::terminal(format!(
                "the declared command `{program}` was cancelled and its process group was terminated"
            ))),
            Err(error) => Err(spawn_failure(program, &error)),
        }
    }

    /// Build the failure for a command that ran and exited non-zero.
    fn exit_failure(&self, program: &str, outcome: &ShellOutcome) -> ActivityFailure {
        // The message carries stderr because a failing command's own words are
        // the most useful thing an operator can be handed; without it the
        // failure says only that something exited non-zero.
        let message = if outcome.stderr.is_empty() {
            format!(
                "the declared command `{program}` exited {} with no standard error output",
                outcome.exit_code
            )
        } else {
            format!(
                "the declared command `{program}` exited {}: {}",
                outcome.exit_code, outcome.stderr
            )
        };
        match self.failure_mode {
            FailureMode::Retryable => ActivityFailure::retryable(message),
            FailureMode::Terminal => ActivityFailure::terminal(message),
        }
    }
}

/// Exit code reported when a command was ended by a signal and so has none.
///
/// Matches the long-standing shell convention of `128 + SIGKILL`, which is the
/// number an operator reading the result will already recognise.
const EXIT_CODE_SIGNALLED: i32 = 137;

/// Classify a substitution failure. All are terminal: the same call with the
/// same arguments would fail identically however many times it is retried.
fn substitution_failure(error: &SubstitutionError) -> ActivityFailure {
    ActivityFailure::terminal(error.to_string())
}

/// Classify a spawn or observation failure.
fn spawn_failure(program: &str, error: &ProcessGroupError) -> ActivityFailure {
    // Terminal in every arm: a program that is absent, unrunnable, or whose
    // containment could not be established will be exactly as absent on the
    // next attempt, and retrying a command whose termination could NOT be
    // confirmed risks running it twice concurrently.
    ActivityFailure::terminal(format!(
        "the declared command `{program}` could not be run to completion: {error}"
    ))
}

/// Trim exactly one trailing newline, and its carriage return if present.
///
/// A command's output almost always ends in a newline that belongs to the
/// terminal rather than to the value, so `echo hello` yields `hello` and not
/// `hello\n`. Only one is removed: a command that deliberately emits blank
/// trailing lines keeps them.
pub(super) fn trim_trailing_newline(text: &str) -> String {
    text.strip_suffix('\n')
        .map_or(text, |trimmed| {
            trimmed.strip_suffix('\r').unwrap_or(trimmed)
        })
        .to_owned()
}

#[cfg(test)]
mod tests {
    use super::{FailureMode, ShellAction, trim_trailing_newline};
    use crate::activity::Classification;
    use crate::context::ActivityContext;
    use aion_core::{ActivityId, RunId, WorkflowId};
    use std::collections::BTreeMap;

    fn context() -> (ActivityContext, crate::context::ActivityCancellationHandle) {
        ActivityContext::new(
            WorkflowId::new_v4(),
            RunId::new_v4(),
            ActivityId::from_sequence_position(1),
            1,
        )
    }

    fn arguments(pairs: &[(&str, serde_json::Value)]) -> BTreeMap<String, serde_json::Value> {
        pairs
            .iter()
            .map(|(name, value)| ((*name).to_owned(), value.clone()))
            .collect()
    }

    /// What a test returns. Every fallible step is carried rather than
    /// unwrapped, because the workspace denies panicking accessors in test
    /// code as firmly as in library code.
    type TestResult = Result<(), Box<dyn std::error::Error>>;

    #[tokio::test]
    async fn a_succeeding_command_returns_its_output() -> TestResult {
        let action = ShellAction::new("echo hello")?;
        let (context, _handle) = context();
        let outcome = action.run(&BTreeMap::new(), &context).await?;
        assert_eq!(outcome.exit_code, 0);
        assert_eq!(outcome.stdout, "hello");
        assert_eq!(outcome.stderr, "");
        Ok(())
    }

    #[tokio::test]
    async fn a_parameter_value_reaches_the_program_as_one_argument() -> TestResult {
        let action = ShellAction::new("echo {{greeting}}")?;
        let (context, _handle) = context();
        let outcome = action
            .run(
                &arguments(&[("greeting", serde_json::json!("hello there world"))]),
                &context,
            )
            .await?;
        assert_eq!(outcome.stdout, "hello there world");
        Ok(())
    }

    #[tokio::test]
    async fn a_hostile_value_is_inert_because_no_shell_ever_sees_it() -> TestResult {
        // If a shell were interposed, this would execute a second command.
        // Because argv goes straight to execve, it is printed as text.
        let action = ShellAction::new("echo {{value}}")?;
        let (context, _handle) = context();
        let outcome = action
            .run(
                &arguments(&[("value", serde_json::json!("hi; echo PWNED"))]),
                &context,
            )
            .await?;
        assert_eq!(outcome.stdout, "hi; echo PWNED");
        assert!(
            !outcome.stdout.contains("PWNED\n"),
            "the injected command must never have run"
        );
        Ok(())
    }

    #[tokio::test]
    async fn a_nonzero_exit_is_retryable_by_default_and_carries_stderr() -> TestResult {
        let action = ShellAction::new("sh -c 'echo trouble >&2; exit 3'")?;
        let (context, _handle) = context();
        let Err(failure) = action.run(&BTreeMap::new(), &context).await else {
            return Err("a non-zero exit must fail the activity".into());
        };
        assert_eq!(failure.classification(), &Classification::Retryable);
        assert!(
            failure.message().contains("trouble"),
            "the failure must carry the command's own words: {}",
            failure.message()
        );
        assert!(failure.message().contains('3'), "the exit code is reported");
        Ok(())
    }

    #[tokio::test]
    async fn a_nonzero_exit_is_terminal_when_the_action_says_so() -> TestResult {
        let action = ShellAction::new("sh -c 'exit 1'")?.with_failure_mode(FailureMode::Terminal);
        let (context, _handle) = context();
        let Err(failure) = action.run(&BTreeMap::new(), &context).await else {
            return Err("a non-zero exit must fail the activity".into());
        };
        assert_eq!(failure.classification(), &Classification::Terminal);
        Ok(())
    }

    #[tokio::test]
    async fn a_missing_parameter_fails_terminally_before_anything_runs() -> TestResult {
        let action = ShellAction::new("echo {{absent}}")?;
        let (context, _handle) = context();
        let Err(failure) = action.run(&BTreeMap::new(), &context).await else {
            return Err("a missing parameter must fail the activity".into());
        };
        assert_eq!(failure.classification(), &Classification::Terminal);
        assert!(failure.message().contains("absent"));
        Ok(())
    }

    #[tokio::test]
    async fn an_absent_program_fails_terminally() -> TestResult {
        let action = ShellAction::new("aion-no-such-program-exists-anywhere")?;
        let (context, _handle) = context();
        let Err(failure) = action.run(&BTreeMap::new(), &context).await else {
            return Err("an absent program must fail the activity".into());
        };
        assert_eq!(failure.classification(), &Classification::Terminal);
        Ok(())
    }

    #[tokio::test]
    async fn cancellation_stops_the_command_and_fails_terminally() -> TestResult {
        let action = ShellAction::new("sleep 30")?;
        let (context, handle) = context();
        let run = tokio::spawn(async move {
            let (context, _keep) = (context, ());
            action.run(&BTreeMap::new(), &context).await
        });
        // Give the command a moment to actually be running before cancelling,
        // so this exercises termination rather than a pre-start short circuit.
        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        handle.cancel();
        let Err(failure) = run.await? else {
            return Err("a cancelled command must fail the activity".into());
        };
        assert_eq!(failure.classification(), &Classification::Terminal);
        assert!(failure.message().contains("cancelled"));
        Ok(())
    }

    #[tokio::test]
    async fn the_hosts_environment_does_not_cross_into_a_declared_command() -> TestResult {
        // A host process legitimately holds credentials. A deployed package
        // must not be able to read them by declaring `run "env"`. This is the
        // test that catches a regression to inherited environment.
        //
        // It reads a variable the host ALREADY has rather than setting one:
        // `std::env::set_var` is `unsafe` (process-global, racy against every
        // other thread) and the workspace denies unsafe outright. Any existing
        // variable other than the deliberately-inherited PATH proves the same
        // property.
        let Some(present) = std::env::vars_os()
            .filter_map(|(name, _)| name.into_string().ok())
            .find(|name| name != "PATH" && !name.is_empty() && !name.contains('='))
        else {
            tracing::info!(
                "skipping: the host has no environment variable besides PATH to prove \
                 non-inheritance with"
            );
            return Ok(());
        };
        let action = ShellAction::new(&format!("sh -c 'echo \"[${{{present}:-absent}}]\"'"))?;
        let (context, _handle) = context();
        let outcome = action.run(&BTreeMap::new(), &context).await?;
        assert_eq!(
            outcome.stdout, "[absent]",
            "the host's `{present}` leaked into a declared command"
        );
        Ok(())
    }

    #[tokio::test]
    async fn an_operator_supplied_variable_does_reach_the_command() -> TestResult {
        // The other half: clearing must not make the environment unusable.
        let mut environment = BTreeMap::new();
        environment.insert("DECLARED_GREETING".to_owned(), "supplied".to_owned());
        let action = ShellAction::new("sh -c 'echo \"[$DECLARED_GREETING]\"'")?
            .with_environment(environment);
        let (context, _handle) = context();
        let outcome = action.run(&BTreeMap::new(), &context).await?;
        assert_eq!(outcome.stdout, "[supplied]");
        Ok(())
    }

    #[tokio::test]
    async fn a_command_runs_in_its_declared_working_directory() -> TestResult {
        let action = ShellAction::new("pwd")?.with_working_directory("/");
        let (context, _handle) = context();
        let outcome = action.run(&BTreeMap::new(), &context).await?;
        assert_eq!(outcome.stdout, "/");
        Ok(())
    }

    #[tokio::test]
    async fn a_command_reading_stdin_gets_end_of_file_rather_than_the_hosts() -> TestResult {
        // Inherited stdin would hang the activity forever on an interactive
        // boot. Closed stdin makes `cat` finish immediately with nothing.
        let action = ShellAction::new("cat")?;
        let (context, _handle) = context();
        let outcome = action.run(&BTreeMap::new(), &context).await?;
        assert_eq!(outcome.exit_code, 0);
        assert_eq!(outcome.stdout, "");
        Ok(())
    }

    #[tokio::test]
    async fn awkward_values_survive_as_exactly_one_argument_each() -> TestResult {
        // Cally's misparse question, not the injection one: a value that a
        // careless argv build would split, truncate, or lose entirely. Each is
        // printed back with a delimiter so a silent split is visible.
        for (value, expected) in [
            ("two\nlines", "[two\nlines]"),
            ("", "[]"),
            ("   leading and trailing   ", "[   leading and trailing   ]"),
            ("tab\there", "[tab\there]"),
            ("quote\"inside", "[quote\"inside]"),
            ("single'inside", "[single'inside]"),
            ("back\\slash", "[back\\slash]"),
        ] {
            let action = ShellAction::new("printf [%s] {{value}}")?;
            let (context, _handle) = context();
            let outcome = action
                .run(&arguments(&[("value", serde_json::json!(value))]), &context)
                .await?;
            assert_eq!(
                outcome.stdout, expected,
                "value {value:?} did not arrive as exactly one argument"
            );
        }
        Ok(())
    }

    #[tokio::test]
    async fn a_value_that_looks_like_a_flag_is_still_passed_as_a_value() -> TestResult {
        // A leading `-` is the misparse Apollo's law points at: the ARGV is
        // correct (one element), but the PROGRAM may read it as an option.
        // The executor cannot fix that — only the author can, by writing
        // `--` in the declaration. Both halves are pinned here so the
        // behaviour is recorded rather than discovered.
        let exposed = ShellAction::new("printf %s {{value}}")?;
        let (exposed_context, _exposed_handle) = context();
        let outcome = exposed
            .run(
                &arguments(&[("value", serde_json::json!("-n"))]),
                &exposed_context,
            )
            .await;
        // `printf %s -n` — the program's own parse decides; the executor
        // delivered exactly what was declared either way.
        assert!(
            outcome.is_ok(),
            "the executor must deliver the value and let the program parse it"
        );

        // With `--` the author has told the program where options end, and the
        // value is unambiguously a value.
        let guarded = ShellAction::new("printf -- [%s] {{value}}")?;
        let (guarded_context, _guarded_handle) = context();
        let guarded_outcome = guarded
            .run(
                &arguments(&[("value", serde_json::json!("-n"))]),
                &guarded_context,
            )
            .await?;
        assert_eq!(guarded_outcome.stdout, "[-n]");
        Ok(())
    }

    /// THE JOIN THIS BUILD MAKES, on the real declared-action path: a command
    /// writing to BOTH streams has those lines on the activity's transcript seam
    /// while it is still running — the same seam, envelope and stream key an
    /// agent step publishes on.
    ///
    /// The run future is polled concurrently with the reads, so a capture that
    /// only materialized at exit would resolve the run arm first and fail the
    /// test by name. The command then sleeps until the test cancels it, which
    /// also holds the pre-existing cancellation contract under streaming.
    #[tokio::test]
    async fn a_declared_command_streams_both_streams_before_it_exits() -> TestResult {
        use aion_core::{RunId, WorkflowId};

        let (sender, mut events) = tokio::sync::mpsc::unbounded_channel();
        let (context, cancellation) = ActivityContext::with_transcript(
            WorkflowId::new_v4(),
            RunId::new_v4(),
            ActivityId::from_sequence_position(9),
            3,
            sender,
        );
        let action = ShellAction::new("sh -c 'echo working; echo warning >&2; sleep 30'")?;
        let no_arguments = BTreeMap::new();
        let run = action.run(&no_arguments, &context);
        tokio::pin!(run);

        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
        let mut observed: Vec<(String, String)> = Vec::new();
        while !observed.iter().any(|(role, _)| role.contains("stdout"))
            || !observed.iter().any(|(role, _)| role.contains("stderr"))
        {
            tokio::select! {
                biased;
                outcome = &mut run => {
                    drop(outcome);
                    return Err(format!(
                        "the command completed before its output reached the transcript: \
                         {observed:?}"
                    )
                    .into());
                }
                event = events.recv() => {
                    let event = event.ok_or("the transcript seam closed mid-command")?;
                    assert_eq!(event.attempt, 3, "the event carries the attempt it belongs to");
                    assert_eq!(
                        event.activity_id,
                        ActivityId::from_sequence_position(9),
                        "the event is keyed to the activity that is running"
                    );
                    let aion_core::ActivityEventKind::Message { text, .. } = event.kind else {
                        return Err("an output line must be a Message".into());
                    };
                    observed.push((event.agent_role, text));
                }
                () = tokio::time::sleep_until(deadline) => {
                    return Err(format!(
                        "the running command's output never reached the transcript: {observed:?}"
                    )
                    .into());
                }
            }
        }

        assert!(
            observed.contains(&("command stdout".to_owned(), "working".to_owned())),
            "stdout must arrive labelled by its stream: {observed:?}"
        );
        assert!(
            observed.contains(&("command stderr".to_owned(), "warning".to_owned())),
            "stderr must arrive labelled by its stream: {observed:?}"
        );

        cancellation.cancel();
        let Err(failure) = run.await else {
            return Err("a cancelled command must fail the activity".into());
        };
        assert_eq!(failure.classification(), &Classification::Terminal);
        Ok(())
    }

    #[test]
    fn only_one_trailing_newline_is_trimmed() {
        assert_eq!(trim_trailing_newline("hello\n"), "hello");
        assert_eq!(trim_trailing_newline("hello\r\n"), "hello");
        assert_eq!(trim_trailing_newline("hello\n\n"), "hello\n");
        assert_eq!(trim_trailing_newline("hello"), "hello");
        assert_eq!(trim_trailing_newline(""), "");
    }
}