newton_cli/commands/

task.rs

use alloy::primitives::{Address, Bytes, U256};
use clap::{Parser, Subcommand};
use eyre::Context;
use newton_prover_core::config::NewtonAvsConfig;
use newton_prover_rpc::types::{
    CreateTaskAndWaitRequest, ICreateTaskAndWaitRequest, ITaskIntent, SignedHashable, TaskIntent,
};
use serde_json::Value;
use std::{
    path::PathBuf,
    sync::atomic::{AtomicU64, Ordering},
};
use tracing::info;

use crate::config::NewtonCliConfig;

/// Task commands
#[derive(Debug, Parser)]
#[command(name = "task")]
pub struct TaskCommand {
    #[command(subcommand)]
    pub subcommand: TaskSubcommand,
}

#[derive(Debug, Subcommand)]
pub enum TaskSubcommand {
    /// Submit evaluation request to prover AVS
    #[command(name = "submit-evaluation-request")]
    SubmitEvaluationRequest(SubmitEvaluationRequestCommand),
}

/// Submit evaluation request command
#[derive(Debug, Parser)]
pub struct SubmitEvaluationRequestCommand {
    /// Path to task JSON file
    #[arg(long)]
    task_json: PathBuf,

    #[arg(long, env = "PRIVATE_KEY")]
    private_key: Option<String>,

    /// API key for authentication
    #[arg(long, env = "API_KEY")]
    api_key: Option<String>,
}

// Static counter for JSON-RPC request IDs
static NEXT_ID: AtomicU64 = AtomicU64::new(0);

fn get_next_id() -> u64 {
    NEXT_ID.fetch_add(1, Ordering::Relaxed) + 1
}

fn create_json_rpc_request_payload(method: &str, params: serde_json::Value) -> serde_json::Value {
    serde_json::json!({
        "jsonrpc": "2.0",
        "id": get_next_id(),
        "method": method,
        "params": params,
    })
}

// Helper function to convert hex string to U256
fn hex_to_u256(hex_str: &str) -> eyre::Result<U256> {
    let hex_str = hex_str.strip_prefix("0x").unwrap_or(hex_str);
    U256::from_str_radix(hex_str, 16).map_err(|e| eyre::eyre!("Failed to parse hex string '{}': {}", hex_str, e))
}

fn get_prover_avs_url(chain_id: u64, deployment_env: &str) -> eyre::Result<String> {
    match (deployment_env, chain_id) {
        ("stagef", 11155111) => Ok("https://prover-avs.stagef.sepolia.newt.foundation".to_string()),
        ("stagef", 1) => Ok("https://prover-avs.stagef.newt.foundation".to_string()),
        ("prod", 11155111) => Ok("https://prover-avs.sepolia.newt.foundation".to_string()),
        ("prod", 1) => Ok("https://prover-avs.newt.foundation".to_string()),
        _ => Err(eyre::eyre!(
            "Unsupported combination: DEPLOYMENT_ENV={}, CHAIN_ID={}",
            deployment_env,
            chain_id
        )),
    }
}

async fn http_post(url: &str, body: &serde_json::Value, api_key: Option<&str>) -> eyre::Result<serde_json::Value> {
    let client = reqwest::Client::new();

    let mut request = client.post(url).header("Content-Type", "application/json");

    if let Some(key) = api_key {
        request = request.header("x-newton-secret", key);
    }

    let response = request.json(body).send().await?;

    let status = response.status();

    if !status.is_success() {
        let error_text = response.text().await?;
        return Err(eyre::eyre!("HTTP error {}: {}", status, error_text));
    }

    let response_json: serde_json::Value = response.json().await?;
    Ok(response_json)
}

// Normalize value/chainId - handles bigint, number, or hex string
fn normalize_to_u256(value: &serde_json::Value) -> eyre::Result<U256> {
    match value {
        serde_json::Value::String(s) => hex_to_u256(s),
        serde_json::Value::Number(n) => {
            let num = n.as_u64().ok_or_else(|| eyre::eyre!("Number too large for u64"))?;
            Ok(U256::from(num))
        }
        _ => Err(eyre::eyre!("Invalid value type: expected string or number")),
    }
}

// Main normalize function - takes a JSON intent and returns hex versions of chain id and value
fn normalize_intent(intent: &serde_json::Value) -> eyre::Result<serde_json::Value> {
    let mut normalized = intent.clone();

    // Normalize value
    if let Some(value) = normalized.get("value") {
        let normalized_value = normalize_to_u256(value)?;
        normalized["value"] = serde_json::Value::String(format!("0x{:x}", normalized_value));
    }

    // Normalize chainId
    if let Some(chain_id) = normalized.get("chainId") {
        let normalized_chain_id = normalize_to_u256(chain_id)?;
        normalized["chainId"] = serde_json::Value::String(format!("0x{:x}", normalized_chain_id));
    }

    Ok(normalized)
}

// Helper to convert hex string to Bytes
fn hex_to_bytes(hex_str: &str) -> eyre::Result<Bytes> {
    let hex_str = hex_str.strip_prefix("0x").unwrap_or(hex_str);
    let bytes = hex::decode(hex_str).map_err(|e| eyre::eyre!("Failed to decode hex: {}", e))?;
    Ok(Bytes::copy_from_slice(&bytes))
}

// Helper functions to get values from JSON
fn get_address(value: &Value) -> eyre::Result<Address> {
    let s = value
        .as_str()
        .ok_or_else(|| eyre::eyre!("Expected string for address"))?;
    s.parse::<Address>().map_err(|e| eyre::eyre!("Invalid address: {}", e))
}

fn get_bytes(value: &Value) -> eyre::Result<Bytes> {
    if value.is_null() {
        return Ok(Bytes::new());
    }
    if let Some(s) = value.as_str() {
        hex_to_bytes(s)
    } else {
        Err(eyre::eyre!("Expected string for bytes"))
    }
}

fn get_u256(value: &Value) -> eyre::Result<U256> {
    if let Some(s) = value.as_str() {
        hex_to_u256(s)
    } else if let Some(n) = value.as_u64() {
        Ok(U256::from(n))
    } else {
        Err(eyre::eyre!("Expected string or number for U256"))
    }
}

// Convert raw JSON intent to TaskIntent
fn json_intent_to_task_intent(intent: &serde_json::Value) -> eyre::Result<TaskIntent> {
    // Normalize value and chainId first
    let normalized_intent = normalize_intent(intent)?;

    let from = get_address(
        normalized_intent
            .get("from")
            .ok_or_else(|| eyre::eyre!("Missing from"))?,
    )?;
    let to = get_address(normalized_intent.get("to").ok_or_else(|| eyre::eyre!("Missing to"))?)?;
    let value = get_u256(
        normalized_intent
            .get("value")
            .ok_or_else(|| eyre::eyre!("Missing value"))?,
    )?;
    let chain_id = get_u256(
        normalized_intent
            .get("chainId")
            .ok_or_else(|| eyre::eyre!("Missing chainId"))?,
    )?;

    // Get data as hex string (with or without 0x prefix)
    let data_str = normalized_intent
        .get("data")
        .ok_or_else(|| eyre::eyre!("Missing data"))?
        .as_str()
        .ok_or_else(|| eyre::eyre!("data must be a string"))?;
    let data = if data_str.starts_with("0x") {
        data_str.to_string()
    } else {
        format!("0x{}", data_str)
    };

    // Get function signature as hex string (with or without 0x prefix)
    let function_signature = normalized_intent
        .get("functionSignature")
        .and_then(|v| v.as_str())
        .map(|s| {
            if s.starts_with("0x") {
                s.to_string()
            } else {
                format!("0x{}", s)
            }
        })
        .unwrap_or_default();

    Ok(TaskIntent {
        from,
        to,
        value,
        data,
        chain_id,
        function_signature,
    })
}

impl TaskCommand {
    /// Execute the task command
    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        match self.subcommand {
            TaskSubcommand::SubmitEvaluationRequest(cmd) => {
                let private_key = cmd
                    .private_key
                    .ok_or_else(|| eyre::eyre!("PRIVATE_KEY is required (set via --private-key or env var)"))?;

                // Get API key from:
                // 1. Command line flag --api-key (highest priority)
                // 2. API_KEY environment variable (from .env file)
                // 3. Config file user_secret_key (from cli-login) (lowest priority)
                let api_key = cmd.api_key.or_else(|| config.service.user_secret_key.clone());

                info!("Reading task JSON from: {:?}", cmd.task_json);
                let contents = std::fs::read_to_string(&cmd.task_json)
                    .with_context(|| format!("Failed to read task JSON file: {:?}", cmd.task_json))?;

                let task: serde_json::Value = serde_json::from_str(&contents)
                    .with_context(|| format!("Failed to parse task JSON: {:?}", cmd.task_json))?;

                let intent_json = task
                    .get("intent")
                    .ok_or_else(|| eyre::eyre!("Missing 'intent' field in task"))?;

                // Convert raw JSON to TaskIntent (with normalization)
                let task_intent = json_intent_to_task_intent(intent_json)?;

                // Convert TaskIntent to ITaskIntent
                let i_task_intent = ITaskIntent::try_from(&task_intent)
                    .with_context(|| "Failed to convert TaskIntent to ITaskIntent")?;

                let intent_sig = task
                    .get("intentSignature")
                    .map(get_bytes)
                    .transpose()?
                    .unwrap_or_default();

                // Get policy client
                let policy_client = get_address(
                    task.get("policyClient")
                        .ok_or_else(|| eyre::eyre!("Missing policyClient"))?,
                )?;

                // Get optional fields
                let quorum_number = task.get("quorumNumber").map(get_bytes).transpose()?.unwrap_or_default();
                let quorum_threshold_percentage = task
                    .get("quorumThresholdPercentage")
                    .and_then(|v| v.as_u64())
                    .map(|v| v as u32)
                    .unwrap_or(0);
                let wasm_args = task.get("wasmArgs").map(get_bytes).transpose()?.unwrap_or_default();
                let timeout = task.get("timeout").and_then(|v| v.as_u64()).unwrap_or(60);

                // Build ICreateTaskAndWaitRequest
                let i_request = ICreateTaskAndWaitRequest {
                    policyClient: policy_client,
                    intent: i_task_intent,
                    intentSignature: intent_sig,
                    quorumNumber: quorum_number,
                    quorumThresholdPercentage: quorum_threshold_percentage,
                    wasmArgs: wasm_args,
                    timeout,
                };

                // Sign using the existing RPC type method
                let signed_request = CreateTaskAndWaitRequest::signed_from(&i_request, &private_key);

                // Serialize to JSON for JSON-RPC request
                let signed_request_json =
                    serde_json::to_value(&signed_request).with_context(|| "Failed to serialize signed request")?;

                let payload = create_json_rpc_request_payload(
                    "newton_createTaskAndWait",
                    serde_json::Value::Array(vec![signed_request_json]),
                );

                let chain_id = config.chain_id;
                let deployment_env = std::env::var("DEPLOYMENT_ENV").unwrap_or_else(|_| "prod".to_string());
                let prover_avs_url = get_prover_avs_url(chain_id, &deployment_env)?;

                info!("Submitting evaluation request to: {}", prover_avs_url);
                let response = http_post(&prover_avs_url, &payload, api_key.as_deref()).await?;

                info!("Response: {}", serde_json::to_string_pretty(&response)?);

                Ok(())
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::str::FromStr;

    #[test]
    fn test_get_next_id() {
        // Reset counter by creating new instances (note: this is a static, so we can't truly reset)
        // Just verify it returns incrementing values
        let id1 = get_next_id();
        let id2 = get_next_id();
        let id3 = get_next_id();
        assert!(id2 > id1);
        assert!(id3 > id2);
    }

    #[test]
    fn test_create_json_rpc_request_payload() {
        let method = "test_method";
        let params = serde_json::json!({"key": "value"});
        let payload = create_json_rpc_request_payload(method, params.clone());

        assert_eq!(payload["jsonrpc"], "2.0");
        assert_eq!(payload["method"], method);
        assert_eq!(payload["params"], params);
        assert!(payload["id"].is_number());
    }

    #[test]
    fn test_hex_to_u256_with_prefix() {
        let result = hex_to_u256("0x1a2b").unwrap();
        assert_eq!(result, U256::from(0x1a2b));
    }

    #[test]
    fn test_hex_to_u256_without_prefix() {
        let result = hex_to_u256("1a2b").unwrap();
        assert_eq!(result, U256::from(0x1a2b));
    }

    #[test]
    fn test_hex_to_u256_large_value() {
        let large_hex = "0xffffffffffffffffffffffffffffffff";
        let result = hex_to_u256(large_hex).unwrap();
        assert!(result > U256::from(u64::MAX));
    }

    #[test]
    fn test_hex_to_u256_invalid_hex() {
        let result = hex_to_u256("0xinvalid");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Failed to parse hex string"));
    }

    #[test]
    fn test_hex_to_u256_empty_string() {
        let result = hex_to_u256("");
        // Empty string should fail or return 0
        assert!(result.is_ok() || result.is_err());
    }

    #[test]
    fn test_get_prover_avs_url_stagef_sepolia() {
        let url = get_prover_avs_url(11155111, "stagef").unwrap();
        assert_eq!(url, "https://prover-avs.stagef.sepolia.newt.foundation");
    }

    #[test]
    fn test_get_prover_avs_url_stagef_mainnet() {
        let url = get_prover_avs_url(1, "stagef").unwrap();
        assert_eq!(url, "https://prover-avs.stagef.newt.foundation");
    }

    #[test]
    fn test_get_prover_avs_url_prod_sepolia() {
        let url = get_prover_avs_url(11155111, "prod").unwrap();
        assert_eq!(url, "https://prover-avs.sepolia.newt.foundation");
    }

    #[test]
    fn test_get_prover_avs_url_prod_mainnet() {
        let url = get_prover_avs_url(1, "prod").unwrap();
        assert_eq!(url, "https://prover-avs.newt.foundation");
    }

    #[test]
    fn test_get_prover_avs_url_unsupported() {
        let result = get_prover_avs_url(999, "stagef");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Unsupported combination"));
    }

    #[test]
    fn test_normalize_to_u256_from_hex_string() {
        let value = serde_json::json!("0x64");
        let result = normalize_to_u256(&value).unwrap();
        assert_eq!(result, U256::from(100));
    }

    #[test]
    fn test_normalize_to_u256_from_decimal_string() {
        // Note: normalize_to_u256 only handles hex strings (via hex_to_u256), not decimal strings
        // A decimal string like "100" will be interpreted as hex (which is 256 in decimal)
        let value = serde_json::json!("100");
        let result = normalize_to_u256(&value).unwrap();
        // "100" as hex = 0x100 = 256 in decimal
        assert_eq!(result, U256::from(256));
    }

    #[test]
    fn test_normalize_to_u256_from_number() {
        let value = serde_json::json!(42);
        let result = normalize_to_u256(&value).unwrap();
        assert_eq!(result, U256::from(42));
    }

    #[test]
    fn test_normalize_to_u256_invalid_type() {
        let value = serde_json::json!(true);
        let result = normalize_to_u256(&value);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Invalid value type"));
    }

    #[test]
    fn test_normalize_intent_with_value_and_chainid() {
        let intent = serde_json::json!({
            "value": "0x64",
            "chainId": 11155111,
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002"
        });
        let result = normalize_intent(&intent).unwrap();
        assert_eq!(result["value"], "0x64");
        assert_eq!(result["chainId"], "0xaa36a7"); // 11155111 in hex
    }

    #[test]
    fn test_normalize_intent_with_number_value() {
        let intent = serde_json::json!({
            "value": 100,
            "chainId": "0x1",
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002"
        });
        let result = normalize_intent(&intent).unwrap();
        assert_eq!(result["value"], "0x64"); // 100 in hex
        assert_eq!(result["chainId"], "0x1");
    }

    #[test]
    fn test_normalize_intent_without_value_or_chainid() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002"
        });
        let result = normalize_intent(&intent).unwrap();
        assert_eq!(result["from"], "0x0000000000000000000000000000000000000001");
        assert_eq!(result["to"], "0x0000000000000000000000000000000000000002");
    }

    #[test]
    fn test_normalize_intent_invalid_value() {
        let intent = serde_json::json!({
            "value": "invalid",
            "chainId": 1
        });
        let result = normalize_intent(&intent);
        assert!(result.is_err());
    }

    #[test]
    fn test_hex_to_bytes_with_prefix() {
        let result = hex_to_bytes("0x1234").unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0], 0x12);
        assert_eq!(result[1], 0x34);
    }

    #[test]
    fn test_hex_to_bytes_without_prefix() {
        let result = hex_to_bytes("1234").unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0], 0x12);
        assert_eq!(result[1], 0x34);
    }

    #[test]
    fn test_hex_to_bytes_empty() {
        let result = hex_to_bytes("0x").unwrap();
        assert_eq!(result.len(), 0);
    }

    #[test]
    fn test_hex_to_bytes_invalid_hex() {
        let result = hex_to_bytes("0xinvalid");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Failed to decode hex"));
    }

    #[test]
    fn test_get_address_valid() {
        let value = serde_json::json!("0x0000000000000000000000000000000000000001");
        let result = get_address(&value).unwrap();
        assert_eq!(
            result,
            Address::from_str("0x0000000000000000000000000000000000000001").unwrap()
        );
    }

    #[test]
    fn test_get_address_invalid_type() {
        let value = serde_json::json!(123);
        let result = get_address(&value);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Expected string"));
    }

    #[test]
    fn test_get_address_invalid_address() {
        let value = serde_json::json!("not_an_address");
        let result = get_address(&value);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Invalid address"));
    }

    #[test]
    fn test_get_bytes_null() {
        let value = serde_json::Value::Null;
        let result = get_bytes(&value).unwrap();
        assert_eq!(result.len(), 0);
    }

    #[test]
    fn test_get_bytes_hex_string() {
        let value = serde_json::json!("0x1234");
        let result = get_bytes(&value).unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0], 0x12);
        assert_eq!(result[1], 0x34);
    }

    #[test]
    fn test_get_bytes_invalid_type() {
        let value = serde_json::json!(123);
        let result = get_bytes(&value);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Expected string"));
    }

    #[test]
    fn test_get_u256_from_hex_string() {
        let value = serde_json::json!("0x64");
        let result = get_u256(&value).unwrap();
        assert_eq!(result, U256::from(100));
    }

    #[test]
    fn test_get_u256_from_number() {
        let value = serde_json::json!(42);
        let result = get_u256(&value).unwrap();
        assert_eq!(result, U256::from(42));
    }

    #[test]
    fn test_get_u256_invalid_type() {
        let value = serde_json::json!(true);
        let result = get_u256(&value);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Expected string or number"));
    }

    #[test]
    fn test_json_intent_to_task_intent_complete() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002",
            "value": "0x64",
            "chainId": 11155111,
            "data": "0x1234"
        });
        let result = json_intent_to_task_intent(&intent).unwrap();
        assert_eq!(
            result.from,
            Address::from_str("0x0000000000000000000000000000000000000001").unwrap()
        );
        assert_eq!(
            result.to,
            Address::from_str("0x0000000000000000000000000000000000000002").unwrap()
        );
        assert_eq!(result.value, U256::from(100));
        assert_eq!(result.chain_id, U256::from(11155111));
        assert_eq!(result.data, "0x1234");
    }

    #[test]
    fn test_json_intent_to_task_intent_with_function_signature() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002",
            "value": 100,
            "chainId": "0x1",
            "data": "1234",
            "functionSignature": "0xabcd"
        });
        let result = json_intent_to_task_intent(&intent).unwrap();
        assert_eq!(result.function_signature, "0xabcd");
        assert_eq!(result.data, "0x1234"); // Should add 0x prefix
    }

    #[test]
    fn test_json_intent_to_task_intent_without_function_signature() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002",
            "value": 100,
            "chainId": 1,
            "data": "0x1234"
        });
        let result = json_intent_to_task_intent(&intent).unwrap();
        assert_eq!(result.function_signature, "");
    }

    #[test]
    fn test_json_intent_to_task_intent_missing_from() {
        let intent = serde_json::json!({
            "to": "0x0000000000000000000000000000000000000002",
            "value": 100,
            "chainId": 1,
            "data": "0x1234"
        });
        let result = json_intent_to_task_intent(&intent);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Missing from"));
    }

    #[test]
    fn test_json_intent_to_task_intent_missing_to() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "value": 100,
            "chainId": 1,
            "data": "0x1234"
        });
        let result = json_intent_to_task_intent(&intent);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Missing to"));
    }

    #[test]
    fn test_json_intent_to_task_intent_missing_value() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002",
            "chainId": 1,
            "data": "0x1234"
        });
        let result = json_intent_to_task_intent(&intent);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Missing value"));
    }

    #[test]
    fn test_json_intent_to_task_intent_missing_chainid() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002",
            "value": 100,
            "data": "0x1234"
        });
        let result = json_intent_to_task_intent(&intent);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Missing chainId"));
    }

    #[test]
    fn test_json_intent_to_task_intent_missing_data() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002",
            "value": 100,
            "chainId": 1
        });
        let result = json_intent_to_task_intent(&intent);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Missing data"));
    }

    #[test]
    fn test_json_intent_to_task_intent_data_not_string() {
        let intent = serde_json::json!({
            "from": "0x0000000000000000000000000000000000000001",
            "to": "0x0000000000000000000000000000000000000002",
            "value": 100,
            "chainId": 1,
            "data": 123
        });
        let result = json_intent_to_task_intent(&intent);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("data must be a string"));
    }
}