clusterflux-public/crates/clusterflux-sdk/src/sdk_runtime.rs

use std::thread;
use std::time::{Duration, Instant};

use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use clusterflux_control::ControlSession;
#[cfg(target_arch = "wasm32")]
use clusterflux_core::TaskDefinitionId;
use clusterflux_core::{
    coordinator_wire_request, Digest, TaskBoundaryValue, TaskFailurePolicy, TaskJoinResult,
    TaskJoinState, TaskSpec,
};
use serde::de::DeserializeOwned;
use serde_json::{json, Value};

use crate::EnvRef;

const COORDINATOR_ENV: &str = "CLUSTERFLUX_SDK_COORDINATOR";
const TENANT_ENV: &str = "CLUSTERFLUX_SDK_TENANT";
const PROJECT_ENV: &str = "CLUSTERFLUX_SDK_PROJECT";
const PROCESS_ENV: &str = "CLUSTERFLUX_SDK_PROCESS";
const USER_ENV: &str = "CLUSTERFLUX_SDK_USER";
const SESSION_SECRET_ENV: &str = "CLUSTERFLUX_SDK_SESSION_SECRET";
const BUNDLE_DIGEST_ENV: &str = "CLUSTERFLUX_SDK_BUNDLE_DIGEST";
const WASM_MODULE_ENV: &str = "CLUSTERFLUX_SDK_WASM_MODULE_BASE64";
const WASM_MODULE_PATH_ENV: &str = "CLUSTERFLUX_SDK_WASM_MODULE_PATH";
const VFS_EPOCH_ENV: &str = "CLUSTERFLUX_SDK_VFS_EPOCH";
const JOIN_TIMEOUT_SECONDS_ENV: &str = clusterflux_core::limits::TASK_JOIN_TIMEOUT_SECONDS_ENV;
const DEFAULT_JOIN_TIMEOUT_SECONDS: u64 =
    clusterflux_core::limits::DEFAULT_TASK_JOIN_TIMEOUT_SECONDS;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ProductRuntimeConfig {
    pub coordinator: String,
    pub tenant: String,
    pub project: String,
    pub process: String,
    pub actor_user: String,
    pub session_secret: Option<String>,
    pub bundle_digest: Digest,
    pub wasm_module_base64: String,
    pub vfs_epoch: u64,
    pub join_poll_interval: Duration,
    pub join_timeout: Duration,
}

impl ProductRuntimeConfig {
    pub fn from_env() -> Result<Option<Self>, TaskRuntimeError> {
        let Some(coordinator) = nonempty_env(COORDINATOR_ENV) else {
            return Ok(None);
        };
        let tenant = required_env(TENANT_ENV)?;
        let project = required_env(PROJECT_ENV)?;
        let process = required_env(PROCESS_ENV)?;
        let actor_user = required_env(USER_ENV)?;
        let bundle_digest = parse_digest(&required_env(BUNDLE_DIGEST_ENV)?)?;
        let wasm_module_base64 = if let Some(module) = nonempty_env(WASM_MODULE_ENV) {
            module
        } else {
            let module_path = required_env(WASM_MODULE_PATH_ENV)?;
            let module = std::fs::read(&module_path).map_err(|error| {
                TaskRuntimeError::Configuration(format!(
                    "read {WASM_MODULE_PATH_ENV} `{module_path}`: {error}"
                ))
            })?;
            BASE64_STANDARD.encode(module)
        };
        let vfs_epoch = required_env(VFS_EPOCH_ENV)?.parse::<u64>().map_err(|_| {
            TaskRuntimeError::Configuration(format!("{VFS_EPOCH_ENV} must be an unsigned integer"))
        })?;
        Ok(Some(Self {
            coordinator,
            tenant,
            project,
            process,
            actor_user,
            session_secret: nonempty_env(SESSION_SECRET_ENV),
            bundle_digest,
            wasm_module_base64,
            vfs_epoch,
            join_poll_interval: Duration::from_millis(25),
            join_timeout: Duration::from_secs(optional_positive_u64_env(
                JOIN_TIMEOUT_SECONDS_ENV,
                DEFAULT_JOIN_TIMEOUT_SECONDS,
            )?),
        }))
    }
}

#[derive(Clone, Debug)]
pub(crate) struct RemoteTaskHandle {
    pub config: Option<ProductRuntimeConfig>,
    pub spec: TaskSpec,
    #[cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]
    pub host_handle_id: Option<u64>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TaskRuntimeError {
    Argument(String),
    Configuration(String),
    Transport(String),
    Protocol(String),
    RemoteTask(String),
    JoinTimeout {
        task: clusterflux_core::TaskInstanceId,
        waited: Duration,
    },
    ResultDecode(String),
    CommandFailed {
        program: String,
        status_code: Option<i32>,
        stdout: String,
        stderr: String,
        stdout_truncated: bool,
        stderr_truncated: bool,
    },
}

impl std::fmt::Display for TaskRuntimeError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let (kind, message) = match self {
            Self::Argument(message) => ("task argument", message),
            Self::Configuration(message) => ("product runtime configuration", message),
            Self::Transport(message) => ("coordinator transport", message),
            Self::Protocol(message) => ("coordinator protocol", message),
            Self::RemoteTask(message) => ("remote task", message),
            Self::JoinTimeout { task, waited } => {
                return write!(
                    formatter,
                    "task join timeout: {task} remained pending for {} seconds; the child task was left running",
                    waited.as_secs()
                );
            }
            Self::ResultDecode(message) => ("remote result", message),
            Self::CommandFailed {
                program,
                status_code,
                stdout,
                stderr,
                stdout_truncated,
                stderr_truncated,
            } => {
                return write!(
                    formatter,
                    "command `{program}` failed with status {status_code:?}; stdout{}: {stdout:?}; stderr{}: {stderr:?}",
                    if *stdout_truncated { " (truncated)" } else { "" },
                    if *stderr_truncated { " (truncated)" } else { "" },
                );
            }
        };
        write!(formatter, "{kind} error: {message}")
    }
}

impl std::error::Error for TaskRuntimeError {}

#[cfg(not(target_arch = "wasm32"))]
pub(crate) fn start_remote_task(
    config: ProductRuntimeConfig,
    task_definition: String,
    environment: Option<EnvRef>,
    args: Vec<TaskBoundaryValue>,
    failure_policy: TaskFailurePolicy,
) -> Result<RemoteTaskHandle, TaskRuntimeError> {
    let _ = (config, task_definition, environment, args, failure_policy);
    Err(TaskRuntimeError::Configuration(
        "native SDK task spawning requires coordinator EntrypointV1 execution; external TaskV1 launch is forbidden"
            .to_owned(),
    ))
}

pub(crate) fn join_remote_task<R>(handle: RemoteTaskHandle) -> Result<R, TaskRuntimeError>
where
    R: DeserializeOwned,
{
    #[cfg(target_arch = "wasm32")]
    if let Some(handle_id) = handle.host_handle_id {
        let response: Result<clusterflux_core::WasmHostTaskJoinResult, String> = guest_host_call(
            &clusterflux_core::WasmHostTaskJoinRequest {
                abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
                handle_id,
            },
            GuestHostCall::Join,
        )?;
        let response = response.map_err(TaskRuntimeError::RemoteTask)?;
        return decode_boundary_value(response.result);
    }
    let config = handle.config.as_ref().ok_or_else(|| {
        TaskRuntimeError::Protocol("remote task handle has no coordinator authority".to_owned())
    })?;
    let started = Instant::now();
    loop {
        let response = coordinator_request(
            config,
            json!({
                "type": "join_task",
                "tenant": config.tenant,
                "project": config.project,
                "actor_user": config.actor_user,
                "process": config.process,
                "task": handle.spec.task_instance,
            }),
        )?;
        if response.get("type").and_then(Value::as_str) == Some("error") {
            return Err(remote_error(&response));
        }
        if response.get("type").and_then(Value::as_str) != Some("task_joined") {
            return Err(TaskRuntimeError::Protocol(
                "join response was not task_joined".to_owned(),
            ));
        }
        let join: TaskJoinResult = serde_json::from_value(
            response
                .get("join")
                .cloned()
                .ok_or_else(|| TaskRuntimeError::Protocol("join result missing".to_owned()))?,
        )
        .map_err(|error| TaskRuntimeError::Protocol(error.to_string()))?;
        match join.state {
            TaskJoinState::Pending => {
                if started.elapsed() >= config.join_timeout {
                    return Err(TaskRuntimeError::JoinTimeout {
                        task: handle.spec.task_instance.clone(),
                        waited: config.join_timeout,
                    });
                }
                thread::sleep(config.join_poll_interval);
            }
            TaskJoinState::Completed => {
                if !join.remote_completion_observed {
                    return Err(TaskRuntimeError::Protocol(
                        "completed join did not prove a signed node completion".to_owned(),
                    ));
                }
                let boundary = join.result.ok_or_else(|| {
                    TaskRuntimeError::ResultDecode(
                        "completed remote task did not return a boundary value".to_owned(),
                    )
                })?;
                return decode_boundary_value(boundary);
            }
            TaskJoinState::Failed | TaskJoinState::Cancelled => {
                return Err(TaskRuntimeError::RemoteTask(join.message));
            }
        }
    }
}

#[cfg(target_arch = "wasm32")]
pub(crate) fn start_guest_host_task(
    task_definition: TaskDefinitionId,
    environment: Option<EnvRef>,
    args: Vec<TaskBoundaryValue>,
    failure_policy: TaskFailurePolicy,
) -> Result<RemoteTaskHandle, TaskRuntimeError> {
    let request = clusterflux_core::WasmHostTaskStartRequest {
        abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
        task_definition,
        environment_id: environment.map(|environment| environment.name.to_owned()),
        args,
        failure_policy,
    };
    request.validate().map_err(TaskRuntimeError::Argument)?;
    let response: Result<clusterflux_core::WasmHostTaskHandle, String> =
        guest_host_call(&request, GuestHostCall::Start)?;
    let response = response.map_err(TaskRuntimeError::RemoteTask)?;
    Ok(RemoteTaskHandle {
        config: None,
        spec: response.task_spec,
        host_handle_id: Some(response.handle_id),
    })
}

#[cfg(target_arch = "wasm32")]
enum GuestHostCall {
    Start,
    Join,
    Command,
    TaskControl,
    DebugProbe,
    SourceSnapshot,
    Vfs,
}

#[cfg(target_arch = "wasm32")]
pub(crate) fn run_guest_host_command(
    program: String,
    args: Vec<String>,
    working_directory: String,
    environment_variables: std::collections::BTreeMap<String, String>,
    timeout_ms: u64,
    network: clusterflux_core::CommandNetworkPolicy,
) -> Result<clusterflux_core::WasmHostCommandResult, TaskRuntimeError> {
    let request = clusterflux_core::WasmHostCommandRequest {
        abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
        program,
        args,
        working_directory,
        environment_variables,
        timeout_ms,
        network,
    };
    request.validate().map_err(TaskRuntimeError::Argument)?;
    let response: Result<clusterflux_core::WasmHostCommandResult, String> =
        guest_host_call(&request, GuestHostCall::Command)?;
    response.map_err(TaskRuntimeError::RemoteTask)
}

#[cfg(target_arch = "wasm32")]
pub(crate) fn guest_cancellation_requested() -> Result<bool, TaskRuntimeError> {
    let request = clusterflux_core::WasmHostTaskControlRequest {
        abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
    };
    request.validate().map_err(TaskRuntimeError::Argument)?;
    let response: Result<clusterflux_core::WasmHostTaskControlResult, String> =
        guest_host_call(&request, GuestHostCall::TaskControl)?;
    response
        .map(|result| result.cancellation_requested)
        .map_err(TaskRuntimeError::RemoteTask)
}

#[cfg(target_arch = "wasm32")]
pub(crate) fn guest_debug_probe(
    symbol: String,
) -> Result<clusterflux_core::WasmHostDebugProbeResult, TaskRuntimeError> {
    let request = clusterflux_core::WasmHostDebugProbeRequest {
        abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
        symbol,
    };
    request.validate().map_err(TaskRuntimeError::Argument)?;
    let response: Result<clusterflux_core::WasmHostDebugProbeResult, String> =
        guest_host_call(&request, GuestHostCall::DebugProbe)?;
    response.map_err(TaskRuntimeError::RemoteTask)
}

#[cfg(target_arch = "wasm32")]
pub(crate) fn guest_vfs_operation(
    operation: clusterflux_core::WasmHostVfsOperation,
) -> Result<clusterflux_core::WasmHostVfsResult, TaskRuntimeError> {
    let request = clusterflux_core::WasmHostVfsRequest {
        abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
        operation,
    };
    request.validate().map_err(TaskRuntimeError::Argument)?;
    let response: Result<clusterflux_core::WasmHostVfsResult, String> =
        guest_host_call(&request, GuestHostCall::Vfs)?;
    response.map_err(TaskRuntimeError::RemoteTask)
}

#[cfg(target_arch = "wasm32")]
pub(crate) fn guest_source_snapshot(
) -> Result<clusterflux_core::WasmHostSourceSnapshotResult, TaskRuntimeError> {
    let request = clusterflux_core::WasmHostSourceSnapshotRequest {
        abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
    };
    request.validate().map_err(TaskRuntimeError::Argument)?;
    let response: Result<clusterflux_core::WasmHostSourceSnapshotResult, String> =
        guest_host_call(&request, GuestHostCall::SourceSnapshot)?;
    response.map_err(TaskRuntimeError::RemoteTask)
}

#[cfg(target_arch = "wasm32")]
fn guest_host_call<I, O>(input: &I, call: GuestHostCall) -> Result<O, TaskRuntimeError>
where
    I: serde::Serialize,
    O: DeserializeOwned,
{
    #[link(wasm_import_module = "clusterflux")]
    unsafe extern "C" {
        #[link_name = "task_start_v1"]
        fn task_start_v1(
            input_pointer: u32,
            input_length: u32,
            output_pointer: u32,
            output_capacity: u32,
        ) -> i32;
        #[link_name = "task_join_v1"]
        fn task_join_v1(
            input_pointer: u32,
            input_length: u32,
            output_pointer: u32,
            output_capacity: u32,
        ) -> i32;
        #[link_name = "command_run_v1"]
        fn command_run_v1(
            input_pointer: u32,
            input_length: u32,
            output_pointer: u32,
            output_capacity: u32,
        ) -> i32;
        #[link_name = "vfs_operation_v1"]
        fn vfs_operation_v1(
            input_pointer: u32,
            input_length: u32,
            output_pointer: u32,
            output_capacity: u32,
        ) -> i32;
        #[link_name = "source_snapshot_v1"]
        fn source_snapshot_v1(
            input_pointer: u32,
            input_length: u32,
            output_pointer: u32,
            output_capacity: u32,
        ) -> i32;
        #[link_name = "task_control_v1"]
        fn task_control_v1(
            input_pointer: u32,
            input_length: u32,
            output_pointer: u32,
            output_capacity: u32,
        ) -> i32;
        #[link_name = "debug_probe_v1"]
        fn debug_probe_v1(
            input_pointer: u32,
            input_length: u32,
            output_pointer: u32,
            output_capacity: u32,
        ) -> i32;
    }

    let input =
        serde_json::to_vec(input).map_err(|error| TaskRuntimeError::Protocol(error.to_string()))?;
    if input.len() > clusterflux_core::MAX_WASM_TASK_ENVELOPE_BYTES {
        return Err(TaskRuntimeError::Argument(
            "Wasm host-call request exceeds the task ABI limit".to_owned(),
        ));
    }
    let mut output = vec![0_u8; clusterflux_core::MAX_WASM_TASK_ENVELOPE_BYTES];
    let written = unsafe {
        match call {
            GuestHostCall::Start => task_start_v1(
                input.as_ptr() as u32,
                input.len() as u32,
                output.as_mut_ptr() as u32,
                output.len() as u32,
            ),
            GuestHostCall::Join => task_join_v1(
                input.as_ptr() as u32,
                input.len() as u32,
                output.as_mut_ptr() as u32,
                output.len() as u32,
            ),
            GuestHostCall::Command => command_run_v1(
                input.as_ptr() as u32,
                input.len() as u32,
                output.as_mut_ptr() as u32,
                output.len() as u32,
            ),
            GuestHostCall::TaskControl => task_control_v1(
                input.as_ptr() as u32,
                input.len() as u32,
                output.as_mut_ptr() as u32,
                output.len() as u32,
            ),
            GuestHostCall::DebugProbe => debug_probe_v1(
                input.as_ptr() as u32,
                input.len() as u32,
                output.as_mut_ptr() as u32,
                output.len() as u32,
            ),
            GuestHostCall::Vfs => vfs_operation_v1(
                input.as_ptr() as u32,
                input.len() as u32,
                output.as_mut_ptr() as u32,
                output.len() as u32,
            ),
            GuestHostCall::SourceSnapshot => source_snapshot_v1(
                input.as_ptr() as u32,
                input.len() as u32,
                output.as_mut_ptr() as u32,
                output.len() as u32,
            ),
        }
    };
    if written <= 0 || written as usize > output.len() {
        return Err(TaskRuntimeError::Protocol(format!(
            "Wasm host-call returned invalid response length {written}"
        )));
    }
    serde_json::from_slice(&output[..written as usize])
        .map_err(|error| TaskRuntimeError::Protocol(error.to_string()))
}

fn coordinator_request(
    config: &ProductRuntimeConfig,
    payload: Value,
) -> Result<Value, TaskRuntimeError> {
    let payload = if let Some(session_secret) = &config.session_secret {
        let request = payload
            .as_object()
            .cloned()
            .ok_or_else(|| TaskRuntimeError::Protocol("request must be an object".to_owned()))?;
        let request = request
            .into_iter()
            .filter(|(key, _)| !matches!(key.as_str(), "tenant" | "project" | "actor_user"))
            .collect::<serde_json::Map<_, _>>();
        json!({
            "type": "authenticated",
            "session_secret": session_secret,
            "request": request,
        })
    } else {
        payload
    };
    let wire = coordinator_wire_request("sdk-product-runtime", payload);
    let mut session = ControlSession::connect(&config.coordinator)
        .map_err(|error| TaskRuntimeError::Transport(error.to_string()))?;
    session
        .request(&wire)
        .map_err(|error| TaskRuntimeError::Transport(error.to_string()))
}

fn decode_boundary_value<R>(value: TaskBoundaryValue) -> Result<R, TaskRuntimeError>
where
    R: DeserializeOwned,
{
    let value = value
        .materialize()
        .map_err(TaskRuntimeError::ResultDecode)?;
    serde_json::from_value(value).map_err(|error| TaskRuntimeError::ResultDecode(error.to_string()))
}

fn optional_positive_u64_env(name: &str, default: u64) -> Result<u64, TaskRuntimeError> {
    let Some(value) = nonempty_env(name) else {
        return Ok(default);
    };
    value
        .parse::<u64>()
        .ok()
        .filter(|value| *value > 0)
        .ok_or_else(|| {
            TaskRuntimeError::Configuration(format!("{name} must be a positive integer"))
        })
}

fn remote_error(response: &Value) -> TaskRuntimeError {
    TaskRuntimeError::RemoteTask(
        response
            .get("message")
            .and_then(Value::as_str)
            .unwrap_or("coordinator rejected task request")
            .to_owned(),
    )
}

fn parse_digest(value: &str) -> Result<Digest, TaskRuntimeError> {
    let digest: Digest = serde_json::from_value(Value::String(value.to_owned()))
        .map_err(|error| TaskRuntimeError::Configuration(error.to_string()))?;
    if !digest.is_valid_sha256() {
        return Err(TaskRuntimeError::Configuration(format!(
            "{BUNDLE_DIGEST_ENV} must be a lowercase sha256 digest"
        )));
    }
    Ok(digest)
}

fn required_env(name: &str) -> Result<String, TaskRuntimeError> {
    nonempty_env(name).ok_or_else(|| {
        TaskRuntimeError::Configuration(format!(
            "{name} is required when {COORDINATOR_ENV} enables product mode"
        ))
    })
}

fn nonempty_env(name: &str) -> Option<String> {
    std::env::var(name)
        .ok()
        .map(|value| value.trim().to_owned())
        .filter(|value| !value.is_empty())
}