aion-cli 0.26.0

//! Compiling the worker document, advertising its surface, and serving it.

use std::collections::BTreeMap;
use std::path::Path;
use std::sync::Arc;

use aion_awl::{CompiledWorkerAction, CompiledWorkerBody, CompiledWorkerDocument, RunCapture};
use aion_package::ActivityDescriptor;
use aion_package::contract::CommandBodyCapture;
use aion_worker::shell::{
    DeclaredCommandAction, ShellAction, resolve_working_directory, shape_command_result,
};
use aion_worker::{ActivityContext, ActivityFailure, Worker, WorkerConfig};
use anyhow::{Context, Result};
use serde_json::Value;

use super::args::{self, AwlWorkerArgs, ConnectionSettings};
use crate::settings::Environment;

/// One action, ready to serve: its executable body.
pub(super) struct ServedAction {
    /// Action name — the activity type dispatched on the queue.
    name: String,
    /// What this action runs.
    body: ServedBody,
}

/// The two executable forms an action's body takes.
///
/// Both are prepared ONCE, at startup. A command the executor cannot accept
/// is a refusal at the desk rather than a terminal failure on the first
/// dispatch that reaches it.
pub(super) enum ServedBody {
    /// A command LINE, parsed here into an argv template.
    Run {
        /// The parsed command.
        command: ShellAction,
        /// What a successful command's result is.
        capture: RunCapture,
    },
    /// A DECLARED command, already emitted by the AWL compiler.
    Command {
        /// The prepared action.
        command: Box<DeclaredCommandAction>,
        /// What a successful command's result is.
        capture: CommandBodyCapture,
    },
}

/// Serve the worker document `args.document` names on the queue it declares.
///
/// # Errors
///
/// Returns an error when the document cannot be read, is a workflow document,
/// does not check cleanly, when a connection setting has no value, or when the
/// worker cannot be built or served.
pub(crate) async fn run(args: &AwlWorkerArgs, endpoint: &str) -> Result<()> {
    let compiled = compile(&args.document)?;
    if args.check {
        return report(&mut std::io::stdout(), &compiled);
    }
    let environment = Environment::from_process();
    let settings = args::resolve(args, &compiled.task_queue, &environment)?;
    build_worker(
        &compiled,
        &settings,
        endpoint,
        crate::awl::document_root(&args.document),
    )?
    .run()
    .await?;
    Ok(())
}

/// Read and compile the document, naming it in every refusal.
pub(super) fn compile(document: &Path) -> Result<CompiledWorkerDocument> {
    let source = std::fs::read_to_string(document)
        .with_context(|| format!("failed to read AWL document {}", document.display()))?;
    aion_awl::compile_worker_document(&source, crate::awl::document_root(document)).map_err(
        |error| {
            anyhow::anyhow!(
                "{} cannot be served:\n{}",
                document.display(),
                crate::awl::compile_diagnostics(document, &error).join("\n")
            )
        },
    )
}

/// The typed descriptor for every action, keyed by action name.
fn descriptors(compiled: &CompiledWorkerDocument) -> BTreeMap<String, ActivityDescriptor> {
    compiled
        .contract
        .actions
        .iter()
        .map(|action| {
            (
                action.name.clone(),
                ActivityDescriptor {
                    name: action.name.clone(),
                    input_schema: action.input_schema.clone(),
                    output_schema: action.output_schema.clone(),
                },
            )
        })
        .collect()
}

/// Write what `--check` reports: the queue, and per action the descriptor this
/// worker would advertise plus the command it would run.
///
/// The descriptors ARE the admission decision on a contract-bearing queue —
/// an operator diffs them against the deployed package contract — so they are
/// canonically encoded, key order coming from the schemas' own content rather
/// than from this build's map representation. The command and its capture ride
/// alongside because they are the execution decision, and a report that showed
/// only the schemas would leave an operator unable to see what will actually
/// run. Reported by the same shape `aion worker agent --check` uses, so the
/// two verbs' surfaces read alike.
///
/// # Errors
///
/// Returns an error when the report cannot be encoded or the writer cannot be
/// written to (a closed pipe).
pub(super) fn report(
    writer: &mut impl std::io::Write,
    compiled: &CompiledWorkerDocument,
) -> Result<()> {
    let descriptors = descriptors(compiled);
    let mut actions = Vec::with_capacity(compiled.actions.len());
    for action in &compiled.actions {
        // Refused rather than reported as null: this report exists to be
        // diffed against a deployed contract, and a missing schema shown as an
        // absent value would read as "the queue expects nothing here".
        let descriptor = descriptors.get(&action.name).with_context(|| {
            format!(
                "no derived descriptor for action `{}`; the action surface is incomplete",
                action.name
            )
        })?;
        let (form, capture, command) = match &action.body {
            CompiledWorkerBody::Run {
                command, capture, ..
            } => ("run", capture_word(*capture), serde_json::json!(command)),
            CompiledWorkerBody::Command { command, capture } => (
                "runs command",
                command_capture_word(*capture),
                serde_json::to_value(command.as_ref())
                    .context("failed to encode a declared command body")?,
            ),
        };
        actions.push(serde_json::json!({
            "name": action.name,
            "form": form,
            "capture": capture,
            "command": command,
            "input_schema": descriptor.input_schema,
            "output_schema": descriptor.output_schema,
        }));
    }
    let report = serde_json::json!({
        "task_queue": compiled.task_queue,
        "actions": actions,
    });
    serde_json::to_writer_pretty(&mut *writer, &aion_package::CanonicalJson::new(report))
        .context("failed to encode the advertised action surface")?;
    writeln!(writer).context("failed to write the advertised action surface")?;
    Ok(())
}

/// How a `run` body's capture reads in the `--check` report.
const fn capture_word(capture: RunCapture) -> &'static str {
    match capture {
        RunCapture::Outcome => "outcome",
        RunCapture::Text => "text",
        RunCapture::Json => "json",
    }
}

/// How a `runs command` body's capture reads in the `--check` report.
const fn command_capture_word(capture: CommandBodyCapture) -> &'static str {
    match capture {
        CommandBodyCapture::Text => "text",
        CommandBodyCapture::Json => "json",
    }
}

/// Build the worker, registering every action the document declares.
///
/// Serving is all-or-nothing by construction here: the document IS the served
/// set, so there is no way to advertise part of a queue — the failure the
/// manifest worker has to check for cannot be expressed.
pub(super) fn build_worker(
    compiled: &CompiledWorkerDocument,
    settings: &ConnectionSettings,
    endpoint: &str,
    workspace_root: &Path,
) -> Result<Worker> {
    let mut descriptors = descriptors(compiled);
    let config = WorkerConfig::builder()
        .endpoint(crate::normalize_endpoint(endpoint))
        .task_queue(&compiled.task_queue)
        .identity(&settings.identity)
        .max_concurrency(settings.concurrency)
        .reconnect_initial_backoff(settings.initial_backoff)
        .reconnect_max_backoff(settings.max_backoff)
        .reconnect_max_attempts(settings.max_attempts)
        .build()?;
    let mut builder = Worker::builder(config);
    for action in &compiled.actions {
        let name = action.name.clone();
        let descriptor = descriptors.remove(&name).with_context(|| {
            format!("no derived descriptor for action `{name}`; the action surface is incomplete")
        })?;
        let served = Arc::new(served_action(action, workspace_root)?);
        builder = builder.register_activity_with_descriptor(
            name,
            descriptor,
            move |input: Value, context| {
                let served = Arc::clone(&served);
                Box::pin(async move { execute(&served, &input, context).await })
            },
        )?;
    }
    builder.build().map_err(Into::into)
}

/// Prepare one action's body at startup.
///
/// A DECLARED command's working directory is resolved HERE, once, against the
/// same root the recipe surface uses — the document's own directory — because
/// a `{workspace_root}` in a worker document means "where this worker was
/// pointed", and asking the question per dispatch would let the answer change
/// underneath a running queue.
pub(super) fn served_action(
    action: &CompiledWorkerAction,
    workspace_root: &Path,
) -> Result<ServedAction> {
    let body = match &action.body {
        CompiledWorkerBody::Run {
            command, capture, ..
        } => ServedBody::Run {
            command: ShellAction::new(command).with_context(|| {
                // The checker rules on command hygiene, so reaching this means
                // the document and the executor disagree — named as the defect
                // it is rather than surfacing later as a dispatch failure.
                format!(
                    "action `{}` declares a command the executor cannot parse",
                    action.name
                )
            })?,
            capture: *capture,
        },
        CompiledWorkerBody::Command { command, capture } => {
            let prepared = DeclaredCommandAction::new(command.as_ref().clone());
            let prepared = match prepared.declared_working_directory() {
                // The SAME rules the server applies, through the same
                // function: absolute, valid UTF-8, no NUL, and created before
                // anything is dispatched into it. Two executors that agreed
                // about the argv and disagreed about the world the process
                // runs in would be two different bodies wearing one
                // declaration.
                Some(declared) => {
                    let resolved = resolve_working_directory(declared, workspace_root)
                        .with_context(|| {
                            format!(
                                "action `{}` declares a working directory this worker cannot                                  resolve",
                                action.name
                            )
                        })?;
                    prepared.with_working_directory(resolved)
                }
                None => prepared,
            };
            ServedBody::Command {
                command: Box::new(prepared),
                capture: *capture,
            }
        }
    };
    Ok(ServedAction {
        name: action.name.clone(),
        body,
    })
}

/// Run one dispatch and shape its result per the action's capture.
///
/// Failure classification is the executor's for every capture: a non-zero
/// exit is retryable and carries the command's own stderr, a spawn failure or
/// a cancellation is terminal. The capture decides only what a SUCCESSFUL
/// command's result is, so the three forms cannot drift into three failure
/// vocabularies.
pub(super) async fn execute(
    served: &ServedAction,
    input: &Value,
    context: &ActivityContext,
) -> Result<Value, ActivityFailure> {
    let arguments = decode_arguments(&served.name, input)?;
    match &served.body {
        ServedBody::Run { command, capture } => {
            let outcome = command.run(&arguments, context).await?;
            match capture {
                RunCapture::Outcome => serde_json::to_value(&outcome).map_err(|error| {
                    ActivityFailure::terminal(format!(
                        "action `{}` could not encode its command outcome: {error}",
                        served.name
                    ))
                }),
                RunCapture::Text => Ok(Value::String(outcome.stdout)),
                RunCapture::Json => serde_json::from_str(&outcome.stdout).map_err(|error| {
                    ActivityFailure::terminal(format!(
                        "action `{}` declares a `run json` body and its command printed output \
                         that is not valid JSON: {error}",
                        served.name
                    ))
                }),
            }
        }
        // The capture is honoured through the SAME function the server calls,
        // so the two executors cannot answer differently about what a
        // declared-command action returns.
        ServedBody::Command { command, capture } => {
            let outcome = command.run(&arguments, context).await?;
            shape_command_result(&served.name, *capture, outcome)
        }
    }
}

/// Decode the dispatch input into the action's named arguments.
///
/// A declared action's parameters are named in the document, so the input
/// must be a JSON object; anything else cannot bind to `$name` references.
/// Retrying cannot change the input, so the refusal is terminal.
pub(super) fn decode_arguments(
    action: &str,
    input: &Value,
) -> Result<BTreeMap<String, Value>, ActivityFailure> {
    match input {
        Value::Object(members) => Ok(members.clone().into_iter().collect()),
        other => Err(ActivityFailure::terminal(format!(
            "action `{action}` was dispatched input that is not an object binding its \
             parameters by name; got {}",
            json_kind(other)
        ))),
    }
}

/// A JSON value's kind, named for a refusal message.
const fn json_kind(value: &Value) -> &'static str {
    match value {
        Value::Null => "null",
        Value::Bool(_) => "a boolean",
        Value::Number(_) => "a number",
        Value::String(_) => "a string",
        Value::Array(_) => "an array",
        Value::Object(_) => "an object",
    }
}