aion-server 0.31.0

use std::net::{SocketAddr, TcpListener as StdTcpListener};
use std::path::{Path, PathBuf};
use std::process::{Child, Command};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use aion_core::Event;
use aion_package::{
    ActionContract, ActivityDescriptor, BeamModule, BeamSet, CURRENT_FORMAT_VERSION,
    DeclaredActivity, Manifest, ManifestVersion, PackageBuilder, PackageContract, WorkerContract,
};
use aion_proto::generated;
use aion_store::{OutboxRow, OutboxStatus, ReadableEventStore};
use aion_store_haematite::HaematiteStore;
use serde_json::json;

pub const NAMESPACE: &str = "default";
/// The task queue the harness worker polls, and therefore the queue this
/// package's activities must be scoped onto.
pub const TASK_QUEUE: &str = "default";
/// The activity types `collect_four` schedules, one per fan-out ordinal — the
/// same four the harness worker registers for.
pub const FAN_ACTIVITY_TYPES: [&str; FAN_OUT] = ["fan:0", "fan:1", "fan:2", "fan:3"];
pub const OUTBOX_MODULE: &str = "aion_outbox_fixture";
pub const FAN_OUT: usize = 4;
/// Readiness/polling deadline for the spawned server. The fixture boots a
/// real 64-shard store: measured 18 s to "resident workflows recovered"
/// under a co-tenanted battery (2026-09-01, where 20 s timed out seconds
/// from ready, twice), single-digit seconds idle. Matches
/// `FANOUT_DECLARED_DEADLINE`: the bound is for boot under load, not a wait
/// the happy path ever spends.
pub const POLL_DEADLINE: Duration = Duration::from_secs(120);

const OUTBOX_BEAM: &[u8] = include_bytes!("../fixtures/aion_outbox_fixture.beam");
const OUTBOX_SOURCE: &[u8] = include_bytes!("../fixtures/aion_outbox_fixture.erl");
pub type TestError = Box<dyn std::error::Error + Send + Sync>;

pub fn test_error(message: impl Into<String>) -> TestError {
    std::io::Error::other(message.into()).into()
}

pub fn unique_temp_dir(name: &str) -> Result<tempfile::TempDir, TestError> {
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|elapsed| elapsed.as_nanos())
        .unwrap_or_default();
    let unique = COUNTER.fetch_add(1, Ordering::Relaxed);
    let pid = std::process::id();
    Ok(tempfile::Builder::new()
        .prefix(&format!(
            "aion-run-server-outbox-{name}-{pid}-{nanos}-{unique}-"
        ))
        .tempdir()?)
}

/// The typed wire descriptors for the four fan activities, derived from the SAME
/// concrete Rust types on both sides of admission: `collect_four` passes each
/// member the JSON string `"in"` and a handler answers with a [`String`].
///
/// One source, two consumers — the package's declaration ([`fixture_contract`])
/// and the harness worker's registration advertisement — so the two cannot drift
/// apart and `WORKER_CONTRACT_MISMATCH` compares two schemas with one origin.
///
/// # Errors
///
/// Returns the descriptor-generation error if a schema cannot be represented as
/// JSON.
pub fn fan_activity_descriptors() -> Result<Vec<ActivityDescriptor>, TestError> {
    let mut descriptors = Vec::with_capacity(FAN_ACTIVITY_TYPES.len());
    for activity_type in FAN_ACTIVITY_TYPES {
        descriptors.push(aion_worker::activity_descriptor::<String, String>(
            activity_type,
        )?);
    }
    Ok(descriptors)
}

/// The fixture's queue-scoped `.v4` contract.
///
/// `PackageContract::from_manifest` by design "never invents a queue", so a
/// manifest's bare activity names land in `unscoped_activities` — and this
/// server runs queue-routed, where an unscoped catalog is a terminal
/// `NO_QUEUE_DECLARATION` at start admission. That refusal is EARNED: an
/// unserved queue would otherwise wait silently forever. So the derived record
/// is amended rather than bypassed — the same four names move out of
/// `unscoped_activities` and onto the queue that serves them.
fn fixture_contract(manifest: &Manifest) -> Result<PackageContract, TestError> {
    let actions = fan_activity_descriptors()?
        .into_iter()
        .map(|descriptor| ActionContract {
            name: descriptor.name,
            input_schema: descriptor.input_schema,
            output_schema: descriptor.output_schema,
            node: None,
            timeout: None,
            retry: None,
            advisory: false,
            // A typed `String -> String` handler serves these, not an agent
            // harness.
            agent: false,
            // The harness worker serves this queue, so the declaration carries
            // no body of its own.
            body: None,
        })
        .collect();
    let mut contract = PackageContract::from_manifest(manifest);
    contract.workers = vec![WorkerContract {
        task_queue: TASK_QUEUE.to_owned(),
        actions,
    }];
    contract.unscoped_activities.clear();
    Ok(contract)
}

pub fn write_package_archive(dir: &Path) -> Result<PathBuf, TestError> {
    let beams = BeamSet::new(vec![BeamModule::new(OUTBOX_MODULE, OUTBOX_BEAM)])?;
    let manifest = Manifest {
        entry_module: OUTBOX_MODULE.to_owned(),
        entry_function: "collect_four".to_owned(),
        input_schema: json!({ "type": "object" }),
        output_schema: json!({}),
        timeout: Some(Duration::from_secs(30)),
        // The four ordinals `collect_four` actually fans out. This manifest used
        // to name one invented activity, `fixture_activity`, that the fixture
        // never schedules and no worker ever served.
        activities: FAN_ACTIVITY_TYPES
            .iter()
            .map(|activity_type| DeclaredActivity {
                activity_type: (*activity_type).to_owned(),
            })
            .collect(),
        version: ManifestVersion::new("stamped-by-builder"),
        format_version: CURRENT_FORMAT_VERSION,
        additional_workflows: Vec::new(),
    };
    let contract = fixture_contract(&manifest)?;
    let archive =
        PackageBuilder::with_source(manifest, beams, [(OUTBOX_MODULE, OUTBOX_SOURCE.to_vec())])
            .with_contract(contract)
            .write_to_bytes()?;
    let path = dir.join("collect_four.aion");
    std::fs::write(&path, archive)?;
    Ok(path)
}

pub struct ServerProcess {
    child: Child,
    /// The Aion home the child was given: its pid file, death note, run and
    /// logs directories all live under here and nowhere else.
    home: PathBuf,
}

impl ServerProcess {
    fn spawn(home: &Path, config_path: &Path) -> Result<Self, TestError> {
        let mut command = Command::new(std::env::current_exe()?);
        command
            .arg("--exact")
            .arg("run_server_child_process")
            .arg("--nocapture");
        for (name, _) in std::env::vars().filter(|(name, _)| name.starts_with("AION_")) {
            command.env_remove(name);
        }
        // The child is a REAL `run_server`: it claims a pid file at birth and
        // arms a death note, both under its Aion home. With `AION_HOME` unset
        // that home is `$HOME/.aion` — the operator's live server's home on the
        // machine that runs these tests — so every child wrote its death note
        // into the operator's logs and found the operator's pid file ("a live
        // server already holds this home's pid file … boots UNCLAIMED"). The
        // harness directory is the child's home; nothing of the operator's is
        // read or written.
        command
            .env("AION_HOME", home)
            .env("AION_RUN_SERVER_CHILD", "1")
            .env("AION_RUN_SERVER_CONFIG", config_path);
        Ok(Self {
            child: command.spawn()?,
            home: home.to_path_buf(),
        })
    }

    /// The child's Aion home — the harness directory it was spawned with.
    pub fn home(&self) -> &Path {
        &self.home
    }

    /// The child's process id, for the pin that reads its pid record back.
    pub fn pid(&self) -> u32 {
        self.child.id()
    }

    async fn wait_ready(&mut self, http: SocketAddr) -> Result<(), TestError> {
        let client = reqwest::Client::new();
        let url = format!("http://{http}/health/ready");
        let deadline = Instant::now() + POLL_DEADLINE;
        loop {
            if let Some(status) = self.child.try_wait()? {
                return Err(test_error(format!(
                    "run_server child exited before readiness: {status}"
                )));
            }
            if client
                .get(&url)
                .send()
                .await
                .is_ok_and(|response| response.status().is_success())
            {
                return Ok(());
            }
            if Instant::now() > deadline {
                return Err(test_error("timed out waiting for run_server readiness"));
            }
            tokio::time::sleep(Duration::from_millis(25)).await;
        }
    }

    /// Kill the server and WAIT for the process to exit.
    ///
    /// An abrupt death is the production way an in-flight dispatch is stranded,
    /// so this is the right stop for a test whose subject is recovery. The
    /// `wait` is also the barrier the caller needs before opening the durable
    /// store: haematite holds an exclusive writer lock on the data directory for
    /// the life of the process, and the OS releases it at exit — so a store
    /// opened after `wait` returns is opening an unlocked directory, with no
    /// sleep anywhere in the argument.
    pub fn stop(mut self) -> Result<(), TestError> {
        self.kill_and_wait()
    }

    /// Ask the server to drain and WAIT for it to exit.
    ///
    /// `run_server` treats the first termination signal as a graceful drain, so
    /// this exercises the production shutdown path and leaves nothing in flight.
    /// Like [`Self::stop`], the `wait` is the barrier that releases haematite's
    /// writer lock; unlike it, nothing is stranded by the stop itself.
    pub fn stop_gracefully(mut self) -> Result<(), TestError> {
        if self.child.try_wait()?.is_none() {
            let delivered = Command::new("kill")
                .args(["-TERM", &self.child.id().to_string()])
                .status()?;
            if !delivered.success() {
                return Err(test_error(format!(
                    "failed to deliver SIGTERM to the run_server child: {delivered}"
                )));
            }
        }
        let status = self.child.wait()?;
        if !status.success() {
            return Err(test_error(format!(
                "the run_server child did not drain cleanly: {status}"
            )));
        }
        Ok(())
    }

    fn kill_and_wait(&mut self) -> Result<(), TestError> {
        if self.child.try_wait()?.is_none() {
            self.child.kill()?;
        }
        let status = self.child.wait()?;
        std::hint::black_box(status);
        Ok(())
    }
}

impl Drop for ServerProcess {
    fn drop(&mut self) {
        if matches!(self.child.try_wait(), Ok(None)) {
            let _ = self.child.kill();
            let _ = self.child.wait();
        }
    }
}

pub async fn start_over_http(
    address: SocketAddr,
) -> Result<(aion_core::WorkflowId, aion_core::RunId), TestError> {
    let client = reqwest::Client::new();
    let url = format!("http://{address}/workflows/start");
    let mut last_error = String::new();
    for attempt in 0..5 {
        match client
            .post(&url)
            .header("content-type", "application/json")
            .header("x-aion-subject", "ci")
            .header("x-aion-namespaces", NAMESPACE)
            .json(&json!({
                "namespace": NAMESPACE,
                "workflow_type": OUTBOX_MODULE,
                "input": { "fixture": "input" },
            }))
            .send()
            .await
        {
            Ok(response) => {
                let status = response.status();
                let bytes = response.bytes().await?;
                if status.is_success() {
                    return workflow_ids_from_start_body(&bytes);
                }
                last_error = format!("HTTP {status}: {}", String::from_utf8_lossy(&bytes));
                if !status.is_server_error() {
                    break;
                }
            }
            Err(error) => last_error = error.to_string(),
        }
        tokio::time::sleep(Duration::from_millis(50 * (attempt + 1))).await;
    }
    Err(test_error(format!(
        "workflow start over HTTP failed: {last_error}"
    )))
}

fn workflow_ids_from_start_body(
    bytes: &[u8],
) -> Result<(aion_core::WorkflowId, aion_core::RunId), TestError> {
    // Clean wire contract: start response exposes plain UUID strings.
    let body: serde_json::Value = serde_json::from_slice(bytes)?;
    let workflow_id = body["workflow_id"]
        .as_str()
        .ok_or_else(|| test_error("start response missing workflow id"))?
        .parse::<uuid::Uuid>()?;
    let run_id = body["run_id"]
        .as_str()
        .ok_or_else(|| test_error("start response missing run id"))?
        .parse::<uuid::Uuid>()?;
    Ok((
        aion_core::WorkflowId::new(workflow_id),
        aion_core::RunId::new(run_id),
    ))
}

pub fn count_kind(history: &[Event], matcher: impl Fn(&Event) -> bool) -> usize {
    history.iter().filter(|event| matcher(event)).count()
}

pub fn count_completed(history: &[Event]) -> usize {
    count_kind(history, |event| {
        matches!(event, Event::ActivityCompleted { .. })
    })
}

pub fn count_completed_for(history: &[Event], ordinal: u64) -> usize {
    count_kind(history, |event| match event {
        Event::ActivityCompleted { activity_id, .. } => activity_id.sequence_position() == ordinal,
        _ => false,
    })
}

/// Read a workflow's history through the operator HTTP door.
///
/// The LIVE phases of these tests may not open the durable store: the server is
/// a separate process and haematite holds an exclusive writer lock on the data
/// directory for its whole life. Synchronisation before a stop therefore goes
/// through `POST /workflows/history` — the same endpoint the ops console reads.
///
/// `payload_limit_bytes = 0` disables payload elision, so each projected event
/// is the plain serde encoding of [`aion_core::Event`] and round-trips back into
/// one. The counting helpers in this module then apply unchanged to both doors,
/// and a history read over HTTP cannot silently mean something different from a
/// history read at the bytes.
pub async fn fetch_history_over_http(
    address: SocketAddr,
    workflow_id: &aion_core::WorkflowId,
) -> Result<Vec<Event>, TestError> {
    let client = reqwest::Client::new();
    let response = client
        .post(format!("http://{address}/workflows/history"))
        .header("content-type", "application/json")
        .header("x-aion-subject", "ci")
        .header("x-aion-namespaces", NAMESPACE)
        .json(&json!({
            "namespace": NAMESPACE,
            "workflow_id": workflow_id.to_string(),
            "limit": 1000,
            "payload_limit_bytes": 0,
        }))
        .send()
        .await?;
    let status = response.status();
    let body = response.text().await?;
    if !status.is_success() {
        return Err(test_error(format!(
            "history fetch failed with {status}: {body}"
        )));
    }
    let body: serde_json::Value = serde_json::from_str(&body)?;
    let events = body
        .get("events")
        .and_then(serde_json::Value::as_array)
        .ok_or_else(|| test_error(format!("history response carried no events array: {body}")))?;
    events
        .iter()
        .map(|event| {
            serde_json::from_value::<Event>(event.clone()).map_err(|error| {
                test_error(format!("history event did not decode as an Event: {error}"))
            })
        })
        .collect()
}

/// Poll the HTTP history door until `predicate` holds, or fail naming what it saw.
///
/// This is SYNCHRONISATION, not an assertion: it establishes that the state a
/// later byte-level assertion is about has been reached, so the server can be
/// stopped. The assertions themselves run against the durable store after the
/// process has exited.
pub async fn wait_for_history_over_http<F>(
    address: SocketAddr,
    workflow_id: &aion_core::WorkflowId,
    description: &str,
    predicate: F,
) -> Result<Vec<Event>, TestError>
where
    F: Fn(&[Event]) -> bool,
{
    let deadline = Instant::now() + POLL_DEADLINE;
    loop {
        let history = fetch_history_over_http(address, workflow_id).await?;
        if predicate(&history) {
            return Ok(history);
        }
        if Instant::now() > deadline {
            return Err(test_error(format!(
                "timed out waiting for {description}: {history:#?}"
            )));
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }
}

/// Every outbox row's durable lifecycle state, read at the bytes.
///
/// Only legitimate once the owning server process has exited — see
/// [`ServerProcess::stop`] for why the exit is the barrier.
pub async fn row_states(
    store: &HaematiteStore,
    workflow_id: &aion_core::WorkflowId,
    ordinals: &[u64],
) -> Result<Vec<aion_store_haematite::OutboxRowState>, TestError> {
    let mut states = Vec::with_capacity(ordinals.len());
    for ordinal in ordinals {
        let key = OutboxRow::dispatch_key_for(workflow_id, *ordinal);
        let state = store
            .outbox_row_state(&key)
            .await?
            .ok_or_else(|| test_error(format!("no outbox row for ordinal {ordinal}")))?;
        states.push(state);
    }
    Ok(states)
}

pub async fn wait_for_history<F>(
    store: &HaematiteStore,
    workflow_id: &aion_core::WorkflowId,
    description: &str,
    predicate: F,
) -> Result<Vec<Event>, TestError>
where
    F: Fn(&[Event]) -> bool,
{
    let deadline = Instant::now() + POLL_DEADLINE;
    loop {
        let history = store.read_history(workflow_id).await?;
        if predicate(&history) {
            return Ok(history);
        }
        if Instant::now() > deadline {
            return Err(test_error(format!(
                "timed out waiting for {description}: {history:#?}"
            )));
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }
}

pub async fn wait_for_rows<F>(
    store: &HaematiteStore,
    workflow_id: &aion_core::WorkflowId,
    ordinals: &[u64],
    description: &str,
    predicate: F,
) -> Result<Vec<OutboxStatus>, TestError>
where
    F: Fn(&[OutboxStatus]) -> bool,
{
    let deadline = Instant::now() + POLL_DEADLINE;
    loop {
        let statuses = row_statuses(store, workflow_id, ordinals).await?;
        if predicate(&statuses) {
            return Ok(statuses);
        }
        if Instant::now() > deadline {
            return Err(test_error(format!(
                "timed out waiting for {description}: {statuses:?}"
            )));
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }
}

async fn row_statuses(
    store: &HaematiteStore,
    workflow_id: &aion_core::WorkflowId,
    ordinals: &[u64],
) -> Result<Vec<OutboxStatus>, TestError> {
    let mut statuses = Vec::with_capacity(ordinals.len());
    for ordinal in ordinals {
        let key = OutboxRow::dispatch_key_for(workflow_id, *ordinal);
        let state = store
            .outbox_row_state(&key)
            .await?
            .ok_or_else(|| test_error(format!("no outbox row for ordinal {ordinal}")))?;
        statuses.push(state.status);
    }
    Ok(statuses)
}

pub fn worker_result(ordinal: u64) -> String {
    format!("\"worker-{ordinal}\"")
}

pub fn task_ordinal(task: &generated::ActivityTask) -> Result<u64, TestError> {
    task.activity_id
        .as_ref()
        .map(|id| id.sequence_position)
        .ok_or_else(|| test_error("pushed task missing activity id"))
}

pub fn assert_task_set(
    tasks: &[generated::ActivityTask],
    expected: &[u64],
) -> Result<(), TestError> {
    let mut ordinals = tasks
        .iter()
        .map(task_ordinal)
        .collect::<Result<Vec<_>, _>>()?;
    ordinals.sort_unstable();
    assert_eq!(ordinals, expected);
    Ok(())
}

pub async fn assert_fan_out_settled(
    reader: &HaematiteStore,
    workflow_id: &aion_core::WorkflowId,
) -> Result<Vec<Event>, TestError> {
    let history = wait_for_history(reader, workflow_id, "fan-out settled", |events| {
        count_completed(events) == FAN_OUT
            && count_kind(events, |event| {
                matches!(event, Event::WorkflowCompleted { .. })
            }) == 1
    })
    .await?;
    for ordinal in 0..FAN_OUT as u64 {
        assert_eq!(count_completed_for(&history, ordinal), 1);
    }
    wait_for_rows(
        reader,
        workflow_id,
        &[0, 1, 2, 3],
        "all rows done",
        |statuses| statuses.iter().all(|status| *status == OutboxStatus::Done),
    )
    .await?;
    assert_collect_result(&history)?;
    Ok(history)
}

fn assert_collect_result(history: &[Event]) -> Result<(), TestError> {
    let result = history
        .iter()
        .find_map(|event| match event {
            Event::WorkflowCompleted { result, .. } => Some(result.clone()),
            _ => None,
        })
        .ok_or_else(|| test_error("no WorkflowCompleted result payload"))?;
    let value: serde_json::Value = serde_json::from_slice(result.bytes())?;
    assert_eq!(
        value,
        json!([
            worker_result(0),
            worker_result(1),
            worker_result(2),
            worker_result(3),
        ])
    );
    Ok(())
}

pub async fn run_server_harness(
    dir: &Path,
    db_path: &Path,
    package_path: &Path,
) -> Result<(ServerProcess, SocketAddr, SocketAddr), TestError> {
    run_server_harness_with_reconciliation(dir, db_path, package_path, None).await
}

pub async fn run_server_harness_with_reconciliation(
    dir: &Path,
    db_path: &Path,
    package_path: &Path,
    reconciliation: Option<(u64, u64)>,
) -> Result<(ServerProcess, SocketAddr, SocketAddr), TestError> {
    let http = reserve_loopback_addr()?;
    let grpc = reserve_loopback_addr()?;
    let config = write_server_config(dir, db_path, package_path, http, grpc, reconciliation)?;
    let mut server = ServerProcess::spawn(dir, &config)?;
    server.wait_ready(http).await?;
    Ok((server, http, grpc))
}

fn reserve_loopback_addr() -> Result<SocketAddr, TestError> {
    let listener = StdTcpListener::bind("127.0.0.1:0")?;
    let address = listener.local_addr()?;
    drop(listener);
    Ok(address)
}

fn write_server_config(
    dir: &Path,
    db_path: &Path,
    package_path: &Path,
    http: SocketAddr,
    grpc: SocketAddr,
    reconciliation: Option<(u64, u64)>,
) -> Result<PathBuf, TestError> {
    let reconciliation = reconciliation.map_or_else(String::new, |(interval, stale_after)| {
        format!("reconcile_interval_ms = {interval}\nreconcile_stale_after_ms = {stale_after}\n")
    });
    let config = format!(
        r#"workflow_packages = [{package}]

[server]
listen_address = "{http}"
grpc_address = "{grpc}"

[store]
backend = "haematite"
data_dir = {db}
# Required, no default: the haematite boot path refuses a config that does not
# rule on the node cache's byte ceiling.
node_cache_budget = "unlimited"

[runtime]
scheduler_threads = 1
query_timeout_ms = 10000

[drain]
timeout_seconds = 30

[namespaces]
default = "{NAMESPACE}"

[metrics]
enabled = false

[websocket]
outbound_buffer_bound = 32
event_broadcast_capacity = 64
cluster_broadcast_capacity = 64

[observability]
# Required, no default: the transcript drain's flush policy.
max_batch_events = 64
max_batch_hold_ms = 0

[outbox]
enabled = true
# The default build compiles liminal-transport, which makes the outbox transport
# default to liminal; this test exercises the gRPC connected-worker path, so it
# selects grpc explicitly (otherwise an enabled outbox would require a liminal
# listen address).
transport = "grpc"
poll_interval_ms = 1000
batch_size = 16
max_attempts = 5
backoff_base_ms = 50
backoff_multiplier = 2
backoff_max_ms = 1000
{reconciliation}
"#,
        package = toml_string(&package_path.display().to_string())?,
        db = toml_string(&db_path.display().to_string())?,
    );
    let path = dir.join(format!("server-{}.toml", grpc.port()));
    std::fs::write(&path, config)?;
    Ok(path)
}

fn toml_string(value: &str) -> Result<String, TestError> {
    serde_json::to_string(value).map_err(Into::into)
}