nexus-sdk 2.0.0

//! Input schema generation for Move onchain tools.

use {
    super::types::{
        convert_move_signature_to_schema,
        is_agent_vertex_authorization_proof_param,
        is_hidden_internal_tool_param,
        is_onchain_tool_result_param,
        is_proof_of_uid_param,
        is_uid_requirements_param,
    },
    crate::{sui, types::OnchainToolMode},
    anyhow::{anyhow, bail, Result as AnyResult},
    serde_json::{Map, Value},
    std::sync::Arc,
};

/// Generate input schema by introspecting the execute function's parameters.
///
/// This function fetches the Move module from the chain and analyzes the
/// execute function's parameters to generate a JSON schema. It automatically
/// skips internal parameters such as `UIDRequirements`, `ProvenValue<AgentVertexAuthorization>`,
/// and `TxContext`.
pub async fn generate_input_schema(
    client: Arc<sui::grpc::Client>,
    package_address: sui::types::Address,
    module_name: &str,
    execute_function: &str,
) -> AnyResult<String> {
    let (schema, _) =
        generate_input_schema_with_mode(client, package_address, module_name, execute_function)
            .await?;
    Ok(schema)
}

/// Generate an input schema and derive its [`OnchainToolMode`].
pub async fn generate_input_schema_with_mode(
    client: Arc<sui::grpc::Client>,
    package_address: sui::types::Address,
    module_name: &str,
    execute_function: &str,
) -> AnyResult<(String, OnchainToolMode)> {
    let request = sui::grpc::GetPackageRequest::default().with_package_id(package_address);
    let mut client = client.as_ref().clone();

    // Fetch all normalized Move modules for the package.
    let Some(package) = client
        .package_client()
        .get_package(request)
        .await
        .map(|resp| resp.into_inner().package)?
    else {
        bail!("Package '{package_address}' not found")
    };
    let all_modules = package.modules();

    // Find the specific module.
    let normalized_module = all_modules
        .iter()
        .find(|m| m.name() == module_name)
        .ok_or_else(|| {
            anyhow!("Module '{module_name}' not found in package '{package_address}'")
        })?;

    // Find the execute function.
    let execute_func = normalized_module
        .functions()
        .iter()
        .find(|f| f.name() == execute_function)
        .ok_or_else(|| {
            anyhow!("Function '{execute_function}' not found in module '{module_name}'")
        })?;
    let mode = validate_execute_signature(execute_func, module_name, execute_function)?;

    // Parse function parameters.
    let mut schema_map = Map::new();
    let mut param_index = 0;

    for param_type in execute_func.parameters() {
        let body = param_type.body_opt().ok_or_else(|| {
            anyhow!("Parameter type missing body in function '{execute_function}' of module '{module_name}'")
        })?;

        if is_hidden_internal_tool_param(body) {
            continue;
        }

        // Convert the full signature (includes reference/mutability info).
        let param_schema = convert_move_signature_to_schema(param_type)?;

        // Store parameter information with index as the default name.
        let mut param_obj = match param_schema {
            Value::Object(obj) => obj,
            other => {
                let mut new_obj = Map::new();
                new_obj.insert("type".to_string(), other);
                new_obj
            }
        };

        let reference_kind = param_type
            .reference
            .and_then(|reference| sui::grpc::open_signature::Reference::try_from(reference).ok())
            .unwrap_or(sui::grpc::open_signature::Reference::Unknown);

        if param_obj.get("type").and_then(|v| v.as_str()) == Some("object")
            && reference_kind != sui::grpc::open_signature::Reference::Unknown
        {
            param_obj.insert(
                "mutable".to_string(),
                Value::Bool(reference_kind == sui::grpc::open_signature::Reference::Mutable),
            );
        }
        schema_map.insert(param_index.to_string(), Value::Object(param_obj));
        param_index += 1;
    }

    let schema_json = Value::Object(schema_map);
    let schema_string = serde_json::to_string(&schema_json)?;

    Ok((schema_string, mode))
}

fn validate_execute_signature(
    execute_func: &sui::grpc::FunctionDescriptor,
    module_name: &str,
    execute_function: &str,
) -> AnyResult<OnchainToolMode> {
    let visibility = execute_func
        .visibility
        .and_then(|visibility| {
            sui::grpc::function_descriptor::Visibility::try_from(visibility).ok()
        })
        .unwrap_or(sui::grpc::function_descriptor::Visibility::Unknown);
    if visibility != sui::grpc::function_descriptor::Visibility::Public {
        bail!("Onchain tool function '{module_name}::{execute_function}' must be public");
    }
    if !execute_func.returns().is_empty() {
        bail!(
            "Onchain tool function '{module_name}::{execute_function}' must not return values; finalize output through an owned OnchainToolResult argument"
        );
    }
    let has_owned_result_arg = execute_func.parameters().iter().any(|param| {
        param.reference.is_none()
            && param
                .body_opt()
                .map(is_onchain_tool_result_param)
                .unwrap_or(false)
    });
    if !has_owned_result_arg {
        bail!(
            "Onchain tool function '{module_name}::{execute_function}' must accept result: OnchainToolResult"
        );
    }

    let parameters = execute_func.parameters();
    let (mode, fixed_prefix_len) = match parameters {
        [authorization, worksheet, result, ..]
            if is_owned_param(authorization, is_agent_vertex_authorization_proof_param)
                && is_owned_param(worksheet, is_uid_requirements_param)
                && is_owned_param(result, is_onchain_tool_result_param) =>
        {
            (OnchainToolMode::WorkflowAuthorization, 3)
        }
        [worksheet, result, ..]
            if is_owned_param(worksheet, is_uid_requirements_param)
                && is_owned_param(result, is_onchain_tool_result_param) =>
        {
            (OnchainToolMode::Standard, 2)
        }
        _ => bail!(
            "Onchain tool function '{module_name}::{execute_function}' must begin with owned UIDRequirements and OnchainToolResult parameters, optionally preceded by an owned AgentVertexAuthorization proof"
        ),
    };
    if parameters[fixed_prefix_len..]
        .iter()
        .filter_map(sui::grpc::OpenSignature::body_opt)
        .any(|body| {
            is_agent_vertex_authorization_proof_param(body)
                || is_proof_of_uid_param(body)
                || is_uid_requirements_param(body)
                || is_onchain_tool_result_param(body)
        })
    {
        bail!(
            "Onchain tool function '{module_name}::{execute_function}' must contain internal parameters only in its fixed prefix"
        );
    }

    Ok(mode)
}

fn is_owned_param(
    parameter: &sui::grpc::OpenSignature,
    predicate: impl FnOnce(&sui::grpc::OpenSignatureBody) -> bool,
) -> bool {
    parameter.reference.is_none() && parameter.body_opt().is_some_and(predicate)
}

#[cfg(test)]
mod tests {
    use {
        super::*,
        sui::grpc::{function_descriptor::Visibility, open_signature_body::Type},
    };

    fn owned_onchain_tool_result_signature() -> sui::grpc::OpenSignature {
        sui::grpc::OpenSignature::default().with_body(
            sui::grpc::OpenSignatureBody::default()
                .with_type(Type::Datatype)
                .with_type_name("0x42::onchain_tool_result::OnchainToolResult"),
        )
    }

    fn referenced_onchain_tool_result_signature(
        reference: sui::grpc::open_signature::Reference,
    ) -> sui::grpc::OpenSignature {
        owned_onchain_tool_result_signature().with_reference(reference)
    }

    fn owned_proof_of_uid_signature() -> sui::grpc::OpenSignature {
        sui::grpc::OpenSignature::default().with_body(
            sui::grpc::OpenSignatureBody::default()
                .with_type(Type::Datatype)
                .with_type_name("0x42::proof_of_uid::ProofOfUID"),
        )
    }

    fn owned_uid_requirements_signature() -> sui::grpc::OpenSignature {
        sui::grpc::OpenSignature::default().with_body(
            sui::grpc::OpenSignatureBody::default()
                .with_type(Type::Datatype)
                .with_type_name("0x42::proof_of_uid::UIDRequirements"),
        )
    }

    fn owned_authorization_signature() -> sui::grpc::OpenSignature {
        sui::grpc::OpenSignature::default().with_body(
            sui::grpc::OpenSignatureBody::default()
                .with_type(Type::Datatype)
                .with_type_name(
                    "0x42::authorization::ProvenValue<0x43::authorization::AgentVertexAuthorization>",
                ),
        )
    }

    fn valid_execute_descriptor() -> sui::grpc::FunctionDescriptor {
        sui::grpc::FunctionDescriptor::default()
            .with_visibility(Visibility::Public)
            .with_is_entry(true)
            .with_parameters(vec![
                owned_uid_requirements_signature(),
                owned_onchain_tool_result_signature(),
            ])
    }

    fn authorized_execute_descriptor() -> sui::grpc::FunctionDescriptor {
        sui::grpc::FunctionDescriptor::default()
            .with_visibility(Visibility::Public)
            .with_is_entry(true)
            .with_parameters(vec![
                owned_authorization_signature(),
                owned_uid_requirements_signature(),
                owned_onchain_tool_result_signature(),
            ])
    }

    #[test]
    fn validate_execute_signature_rejects_private_function() {
        let descriptor = valid_execute_descriptor().with_visibility(Visibility::Private);
        let err = validate_execute_signature(&descriptor, "tool", "execute").unwrap_err();
        assert!(err.to_string().contains("must be public"));
    }

    #[test]
    fn validate_execute_signature_accepts_public_nonentry_function() {
        let descriptor = valid_execute_descriptor().with_is_entry(false);
        validate_execute_signature(&descriptor, "tool", "execute").unwrap();
    }

    #[test]
    fn validate_execute_signature_rejects_return_values() {
        let descriptor =
            valid_execute_descriptor().with_returns(vec![sui::grpc::OpenSignature::default()
                .with_body(sui::grpc::OpenSignatureBody::default().with_type(Type::U64))]);
        let err = validate_execute_signature(&descriptor, "tool", "execute").unwrap_err();
        assert!(err.to_string().contains("must not return values"));
    }

    #[test]
    fn validate_execute_signature_rejects_referenced_onchain_tool_result() {
        let descriptor = sui::grpc::FunctionDescriptor::default()
            .with_visibility(Visibility::Public)
            .with_is_entry(true)
            .with_parameters(vec![
                owned_uid_requirements_signature(),
                referenced_onchain_tool_result_signature(
                    sui::grpc::open_signature::Reference::Immutable,
                ),
            ]);
        let err = validate_execute_signature(&descriptor, "tool", "execute").unwrap_err();
        assert!(err
            .to_string()
            .contains("must accept result: OnchainToolResult"));
    }

    #[test]
    fn validate_execute_signature_rejects_missing_requirements() {
        let descriptor = sui::grpc::FunctionDescriptor::default()
            .with_visibility(Visibility::Public)
            .with_is_entry(true)
            .with_parameters(vec![owned_onchain_tool_result_signature()]);
        let err = validate_execute_signature(&descriptor, "tool", "execute").unwrap_err();
        assert!(err.to_string().contains("must begin with"));
    }

    #[test]
    fn validate_execute_signature_rejects_proof_in_place_of_requirements() {
        let descriptor = sui::grpc::FunctionDescriptor::default()
            .with_visibility(Visibility::Public)
            .with_is_entry(true)
            .with_parameters(vec![
                owned_proof_of_uid_signature(),
                owned_onchain_tool_result_signature(),
            ]);
        let err = validate_execute_signature(&descriptor, "tool", "execute").unwrap_err();
        assert!(err.to_string().contains("UIDRequirements"));
    }

    #[test]
    fn validate_execute_signature_rejects_misplaced_authorization() {
        let descriptor = sui::grpc::FunctionDescriptor::default()
            .with_visibility(Visibility::Public)
            .with_is_entry(true)
            .with_parameters(vec![
                owned_uid_requirements_signature(),
                owned_authorization_signature(),
                owned_onchain_tool_result_signature(),
            ]);
        let err = validate_execute_signature(&descriptor, "tool", "execute").unwrap_err();
        assert!(err.to_string().contains("must begin with"));
    }

    #[test]
    fn validate_execute_signature_rejects_duplicate_internal_parameter() {
        let descriptor = sui::grpc::FunctionDescriptor::default()
            .with_visibility(Visibility::Public)
            .with_is_entry(true)
            .with_parameters(vec![
                owned_uid_requirements_signature(),
                owned_onchain_tool_result_signature(),
                owned_authorization_signature(),
            ]);
        let err = validate_execute_signature(&descriptor, "tool", "execute").unwrap_err();
        assert!(err.to_string().contains("fixed prefix"));
    }

    #[test]
    fn validate_execute_signature_accepts_entry_no_return_owned_result() {
        validate_execute_signature(&valid_execute_descriptor(), "tool", "execute").unwrap();
    }

    #[test]
    fn validate_execute_signature_detects_standard_mode() {
        assert_eq!(
            validate_execute_signature(&valid_execute_descriptor(), "tool", "execute").unwrap(),
            OnchainToolMode::Standard
        );
    }

    #[test]
    fn validate_execute_signature_detects_workflow_authorization() {
        assert_eq!(
            validate_execute_signature(&authorized_execute_descriptor(), "tool", "execute")
                .unwrap(),
            OnchainToolMode::WorkflowAuthorization
        );
    }

    #[tokio::test]
    #[cfg(feature = "test_utils")]
    async fn test_generate_input_schema_from_published_package() {
        use crate::test_utils;

        // Spin up the Sui instance.
        let test_utils::containers::SuiInstance {
            rpc_port,
            faucet_port,
            container: _container,
            ..
        } = test_utils::containers::setup_sui_instance().await;

        let rpc_url = format!("http://127.0.0.1:{rpc_port}");
        let faucet_url = format!("http://127.0.0.1:{faucet_port}/gas");

        let mut rng = rand::thread_rng();

        // Create a wallet and request some gas tokens.
        let pk = sui::crypto::Ed25519PrivateKey::generate(&mut rng);
        let addr = pk.public_key().derive_address();

        test_utils::faucet::request_tokens(&faucet_url, addr)
            .await
            .expect("Failed to request tokens from faucet.");

        let (gas_coin, _) = test_utils::gas::fetch_gas_coins(&rpc_url, addr)
            .await
            .expect("Failed to fetch gas coin.")
            .into_iter()
            .next()
            .unwrap();

        let package = crate::integration_test_support::compile_move_package(
            &std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/move/onchain_tool_test"),
        )
        .expect("failed to compile Move package");
        let response =
            test_utils::contracts::publish_move_package(&pk, &rpc_url, package, gas_coin).await;

        let pkg_id = response
            .objects
            .iter()
            .find_map(|c| match c.data() {
                sui::types::ObjectData::Package(m) => Some(m.id),
                _ => None,
            })
            .expect("Move package must be published");

        let client = Arc::new(
            sui::grpc::client(format!("http://127.0.0.1:{rpc_port}"))
                .expect("Could not create gRPC client"),
        );

        // Generate input schema for the onchain_tool::execute function.
        let schema_str = generate_input_schema(
            client,
            pkg_id.to_string().parse().unwrap(),
            "onchain_tool",
            "execute",
        )
        .await
        .expect("Failed to generate input schema");

        // Parse the schema.
        let schema: serde_json::Value =
            serde_json::from_str(&schema_str).expect("Failed to parse schema JSON");

        // Verify schema structure.
        // The execute function has hidden internal parameters:
        // execute(requirements: UIDRequirements, result: OnchainToolResult, current_execution: &mut DAGExecution, counter: &mut RandomCounter, increase_with: u64, _ctx: &mut TxContext)
        // After skipping hidden internal parameters, we should have:
        // - Parameter 0: current_execution (&mut DAGExecution) - object type, mutable
        // - Parameter 1: counter (&mut RandomCounter) - object type, mutable
        // - Parameter 2: increase_with (u64).

        // Check parameter 0 (current_execution).
        let param0 = schema
            .get("0")
            .expect("Schema should have parameter 0 (current_execution)");
        assert_eq!(param0["type"], "object");
        assert_eq!(param0["mutable"], true);
        assert!(param0["description"]
            .as_str()
            .unwrap()
            .contains("DAGExecution"));

        // Check parameter 1 (counter).
        let param1 = schema
            .get("1")
            .expect("Schema should have parameter 1 (counter)");
        assert_eq!(param1["type"], "object");
        assert_eq!(param1["mutable"], true);
        assert!(param1["description"]
            .as_str()
            .unwrap()
            .contains("RandomCounter"));

        // Check parameter 2 (increase_with).
        let param2 = schema
            .get("2")
            .expect("Schema should have parameter 2 (increase_with)");
        assert_eq!(param2["type"], "u64");
        assert_eq!(param2["description"], "64-bit unsigned integer");
        assert!(param2.get("mutable").is_none());

        // Verify only 3 user inputs remain after internal params are skipped.
        assert_eq!(schema.as_object().unwrap().len(), 3);
    }
}