newton_cli/commands/

version.rs

/// Version compatibility checking and migration commands
use alloy::{
    network::EthereumWallet, primitives::Address, providers::ProviderBuilder, signers::local::PrivateKeySigner,
    transports::http::reqwest::Url,
};
use clap::Parser;
use eyre::{Context, Result};
use newton_prover_chainio::{
    policy_client::get_policy_address_for_client,
    version::{
        get_policy_data_factory_for_policy_data, get_policy_data_factory_version, get_policy_factory_for_policy,
        get_policy_factory_version,
    },
};
use newton_prover_core::{
    config::rpc::RpcProviderConfig,
    newton_policy::NewtonPolicy,
    newton_policy_client::NewtonPolicyClient,
    newton_policy_data::NewtonPolicyData,
    newton_policy_data_factory::NewtonPolicyDataFactory,
    newton_policy_factory::NewtonPolicyFactory,
    version::{is_compatible, MIN_COMPATIBLE_VERSION, PROTOCOL_VERSION},
};
use serde::Serialize;
use std::str::FromStr;
use tracing::{debug, error, info, warn};

#[derive(Parser, Debug)]
pub enum VersionCommand {
    /// Check compatibility of a policy client with current protocol version
    CheckCompatibility {
        /// Policy client contract address
        #[arg(long)]
        policy_client: Address,

        /// Chain ID
        #[arg(long)]
        chain_id: u64,

        /// RPC URL (optional, will use default from config)
        #[arg(long)]
        rpc_url: Option<String>,
    },

    /// Migrate a policy client to a compatible version
    Migrate {
        /// Policy client contract address to migrate
        #[arg(long)]
        policy_client: Address,

        /// Private key of the policy client owner
        #[arg(long, env = "PRIVATE_KEY")]
        private_key: String,

        /// Chain ID
        #[arg(long)]
        chain_id: u64,

        /// RPC URL (optional, will use default from config)
        #[arg(long)]
        rpc_url: Option<String>,

        /// Skip compatibility check before migration
        #[arg(long)]
        skip_check: bool,

        /// Dry run (show what would be done without executing)
        #[arg(long)]
        dry_run: bool,
    },

    /// Display protocol version information
    Info,
}

#[derive(Serialize, Debug)]
pub struct CompatibilityReport {
    pub protocol_version: String,
    pub policy_client: Address,
    pub policy_address: Address,
    pub policy_factory_version: String,
    pub policy_compatible: bool,
    pub policy_data_reports: Vec<PolicyDataReport>,
    pub all_compatible: bool,
    pub migration_required: bool,
}

#[derive(Serialize, Debug)]
pub struct PolicyDataReport {
    pub address: Address,
    pub factory_version: String,
    pub compatible: bool,
}

impl VersionCommand {
    pub async fn run(self) -> Result<()> {
        match self {
            VersionCommand::CheckCompatibility {
                policy_client,
                chain_id,
                rpc_url,
            } => check_compatibility(policy_client, chain_id, rpc_url).await,
            VersionCommand::Migrate {
                policy_client,
                private_key,
                chain_id,
                rpc_url,
                skip_check,
                dry_run,
            } => migrate_policy(policy_client, private_key, chain_id, rpc_url, skip_check, dry_run).await,
            VersionCommand::Info => {
                print_version_info();
                Ok(())
            }
        }
    }
}

async fn check_compatibility(policy_client: Address, chain_id: u64, rpc_url: Option<String>) -> Result<()> {
    let rpc_url = rpc_url.unwrap_or_else(|| {
        RpcProviderConfig::load(chain_id)
            .expect("Failed to load RPC config")
            .http
    });

    info!("Checking version compatibility for policy client: {}", policy_client);
    info!("Protocol version: {}", PROTOCOL_VERSION);

    // Get policy address
    let policy_address = get_policy_address_for_client(policy_client, rpc_url.clone()).await?;
    info!("Policy address: {}", policy_address);

    // Get policy factory and version
    let policy_factory = get_policy_factory_for_policy(policy_address, &rpc_url).await?;
    let policy_version = get_policy_factory_version(policy_factory, &rpc_url).await?;
    info!("Policy factory version: {}", policy_version);

    // Check policy compatibility
    let policy_compatible = is_compatible(&policy_version, MIN_COMPATIBLE_VERSION)?;
    if policy_compatible {
        info!("✓ Policy version is compatible");
    } else {
        warn!(
            "✗ Policy version is INCOMPATIBLE (minimum required: v{})",
            MIN_COMPATIBLE_VERSION
        );
    }

    // Check policy data compatibility
    info!("Checking policy data...");

    // Get policy data addresses from policy contract
    let provider = ProviderBuilder::new().connect_http(rpc_url.parse()?);
    let policy_contract = NewtonPolicy::new(policy_address, provider);

    let policy_data_addresses = policy_contract
        .getPolicyData()
        .call()
        .await
        .context("Failed to get policy data addresses")?;

    info!("Found {} policy data contracts", policy_data_addresses.len());

    let mut policy_data_reports = Vec::new();
    let mut all_compatible = policy_compatible;

    for policy_data_addr in policy_data_addresses {
        debug!("Checking policy data at {}...", policy_data_addr);

        match check_policy_data_compatibility(policy_data_addr, &rpc_url).await {
            Ok((version, compatible)) => {
                debug!("Version: {}, Compatible: {}", version, compatible);
                policy_data_reports.push(PolicyDataReport {
                    address: policy_data_addr,
                    factory_version: version,
                    compatible,
                });
                all_compatible = all_compatible && compatible;
            }
            Err(e) => {
                error!("Error checking compatibility: {}", e);
                // Consider policy data with errors as incompatible
                policy_data_reports.push(PolicyDataReport {
                    address: policy_data_addr,
                    factory_version: "unknown".to_string(),
                    compatible: false,
                });
                all_compatible = false;
            }
        }
    }

    let report = CompatibilityReport {
        protocol_version: PROTOCOL_VERSION.to_string(),
        policy_client,
        policy_address,
        policy_factory_version: policy_version.clone(),
        policy_compatible,
        policy_data_reports,
        all_compatible,
        migration_required: !all_compatible,
    };

    info!("=== Compatibility Report ===");
    info!("{}", serde_json::to_string_pretty(&report)?);

    if report.migration_required {
        warn!("Migration required!");
        info!("To migrate your policy to the latest version:");
        info!("  newton-cli policy migrate --policy-client {} \\", policy_client);
        info!("    --private-key $YOUR_PRIVATE_KEY \\");
        info!("    --chain-id {}", chain_id);
        info!("Guide: https://docs.newton.xyz/versioning/migration");
        std::process::exit(1);
    } else {
        info!("✓ All versions are compatible. No migration needed.");
        Ok(())
    }
}

fn print_version_info() {
    info!("Newton Protocol Version Information");
    info!("====================================");
    info!("Protocol Version: {}", PROTOCOL_VERSION);
    info!("Minimum Compatible Version: {}", MIN_COMPATIBLE_VERSION);
    info!("Version Compatibility:");
    info!("  - Major versions must match exactly");
    info!("  - Minor version must be >= minimum");
    info!("  - Patch version is ignored (backward-compatible bug fixes)");
    info!("For more information, see: https://docs.newton.xyz/versioning");
}

/// Check compatibility of a policy data contract
async fn check_policy_data_compatibility(policy_data_addr: Address, rpc_url: &str) -> Result<(String, bool)> {
    let factory = get_policy_data_factory_for_policy_data(policy_data_addr, rpc_url).await?;
    let version = get_policy_data_factory_version(factory, rpc_url).await?;
    let compatible = is_compatible(&version, MIN_COMPATIBLE_VERSION)?;
    Ok((version, compatible))
}

/// Get policy configuration from chain
async fn get_policy_info(policy_address: Address, rpc_url: &str) -> Result<PolicyInfo> {
    let provider = ProviderBuilder::new().connect_http(rpc_url.parse()?);
    let policy_contract = NewtonPolicy::new(policy_address, provider);

    let policy_cid = policy_contract
        .getPolicyCid()
        .call()
        .await
        .context("Failed to get policy CID")?;

    let schema_cid = policy_contract
        .getSchemaCid()
        .call()
        .await
        .context("Failed to get schema CID")?;

    let entrypoint = policy_contract
        .getEntrypoint()
        .call()
        .await
        .context("Failed to get entrypoint")?;

    let metadata_cid = policy_contract
        .getMetadataCid()
        .call()
        .await
        .context("Failed to get metadata CID")?;

    let policy_data = policy_contract
        .getPolicyData()
        .call()
        .await
        .context("Failed to get policy data addresses")?;

    Ok(PolicyInfo {
        policy_cid,
        schema_cid,
        entrypoint,
        metadata_cid,
        policy_data,
    })
}

#[derive(Debug, Clone)]
struct PolicyInfo {
    policy_cid: String,
    schema_cid: String,
    entrypoint: String,
    metadata_cid: String,
    policy_data: Vec<Address>,
}

/// Migrate a policy client to use the latest compatible policy version
async fn migrate_policy(
    policy_client: Address,
    private_key: String,
    chain_id: u64,
    rpc_url: Option<String>,
    skip_check: bool,
    dry_run: bool,
) -> Result<()> {
    let rpc_url = rpc_url.unwrap_or_else(|| {
        RpcProviderConfig::load(chain_id)
            .expect("Failed to load RPC config")
            .http
    });

    info!("=== Policy Migration Tool ===");
    info!("Policy Client: {}", policy_client);
    info!("Chain ID: {}", chain_id);
    info!("Protocol Version: {}", PROTOCOL_VERSION);

    // Step 1: Check compatibility (unless skipped)
    if !skip_check {
        info!("Step 1: Checking current compatibility...");

        let policy_address = get_policy_address_for_client(policy_client, rpc_url.clone()).await?;
        let policy_factory = get_policy_factory_for_policy(policy_address, &rpc_url).await?;
        let current_version = get_policy_factory_version(policy_factory, &rpc_url).await?;

        info!("Current policy version: {}", current_version);
        info!("Minimum required version: {}", MIN_COMPATIBLE_VERSION);

        let is_compatible_now = is_compatible(&current_version, MIN_COMPATIBLE_VERSION)?;

        if is_compatible_now {
            info!("✓ Policy is already compatible!");
            info!("No migration needed.");
            return Ok(());
        }

        warn!("✗ Policy is incompatible and needs migration");
    }

    // Step 2: Read current policy configuration
    info!("Step 2: Reading current policy configuration...");

    // Get current policy address and configuration
    let policy_address = get_policy_address_for_client(policy_client, rpc_url.clone()).await?;
    info!("Current policy address: {}", policy_address);

    // Get policy configuration from chain
    let policy_info = get_policy_info(policy_address, &rpc_url).await?;
    info!("- Policy CID: {}", policy_info.policy_cid);
    info!("- Schema CID: {}", policy_info.schema_cid);
    info!("- Entrypoint: {}", policy_info.entrypoint);
    info!("- Metadata CID: {}", policy_info.metadata_cid);
    info!("- Policy Data count: {}", policy_info.policy_data.len());

    // Step 3: Deploy new policy with latest factory
    info!("Step 3: Deploying new policy data (if needed) and new policy with latest factory version...");

    // First, handle policy data migration
    let mut new_policy_data_addresses = Vec::new();
    let mut incompatible_policy_data = Vec::new();

    for (i, policy_data_addr) in policy_info.policy_data.iter().enumerate() {
        debug!("Checking policy data {} at {}...", i + 1, policy_data_addr);

        match check_policy_data_compatibility(*policy_data_addr, &rpc_url).await {
            Ok((version, compatible)) => {
                debug!("Version: {}, Compatible: {}", version, compatible);
                if !compatible {
                    incompatible_policy_data.push(*policy_data_addr);
                } else {
                    // Compatible, reuse existing address
                    new_policy_data_addresses.push(*policy_data_addr);
                }
            }
            Err(e) => {
                error!("Error checking compatibility: {}", e);
                incompatible_policy_data.push(*policy_data_addr);
            }
        }
    }

    if !incompatible_policy_data.is_empty() {
        info!(
            "Step 3a: Migrating {} incompatible policy data contracts...",
            incompatible_policy_data.len()
        );

        if dry_run {
            warn!(
                "DRY RUN MODE - Would migrate {} incompatible policy data contracts:",
                incompatible_policy_data.len()
            );
            for addr in &incompatible_policy_data {
                info!("- {}", addr);
            }
            // In dry run, just add placeholders
            for _ in &incompatible_policy_data {
                new_policy_data_addresses.push(Address::ZERO);
            }
        } else {
            // Parse private key and create signer
            let private_key_str = private_key.strip_prefix("0x").unwrap_or(&private_key);
            let signer = PrivateKeySigner::from_str(private_key_str).context("Failed to parse private key")?;
            let signer_address = signer.address();
            let wallet = EthereumWallet::from(signer);

            // Create provider with signer
            let url = Url::parse(&rpc_url).context("Invalid RPC URL")?;
            let provider = ProviderBuilder::new().wallet(wallet).connect_http(url);

            for (i, policy_data_addr) in incompatible_policy_data.iter().enumerate() {
                info!(
                    "Migrating policy data {} of {}: {}",
                    i + 1,
                    incompatible_policy_data.len(),
                    policy_data_addr
                );

                // Get policy data factory address
                let policy_data_factory = get_policy_data_factory_for_policy_data(*policy_data_addr, &rpc_url).await?;
                info!("Using policy data factory at: {}", policy_data_factory);

                // Get current policy data configuration
                let policy_data_contract = NewtonPolicyData::new(*policy_data_addr, provider.clone());
                let wasm_cid = policy_data_contract
                    .getWasmCid()
                    .call()
                    .await
                    .context("Failed to get WASM CID")?;
                let secrets_schema_cid = policy_data_contract
                    .getSecretsSchemaCid()
                    .call()
                    .await
                    .context("Failed to get secrets schema")?;
                let expire_after = policy_data_contract
                    .getExpireAfter()
                    .call()
                    .await
                    .context("Failed to get expireAfter")?;
                let metadata_cid = policy_data_contract
                    .getMetadataCid()
                    .call()
                    .await
                    .context("Failed to get metadata CID")?;

                info!("- WASM CID: {}", wasm_cid);
                info!("- Secrets Schema CID: {}", secrets_schema_cid);
                info!("- Expire After: {}", expire_after);
                info!("- Metadata CID: {}", metadata_cid);

                // Deploy new policy data with same configuration
                let factory_contract = NewtonPolicyDataFactory::new(policy_data_factory, provider.clone());

                info!("Submitting deployPolicyData transaction...");
                let tx = factory_contract
                    .deployPolicyData(wasm_cid, secrets_schema_cid, expire_after, metadata_cid, signer_address)
                    .send()
                    .await
                    .context("Failed to send deployPolicyData transaction")?;

                info!("Transaction sent: {:?}", tx.tx_hash());
                info!("Waiting for transaction confirmation...");

                let receipt = tx.get_receipt().await.context("Failed to get transaction receipt")?;

                // Extract new policy data address from PolicyDataDeployed event
                let logs = receipt.inner.logs();
                let policy_data_deployed_event = logs
                    .iter()
                    .find_map(|log| log.log_decode::<NewtonPolicyDataFactory::PolicyDataDeployed>().ok())
                    .ok_or_else(|| eyre::eyre!("PolicyDataDeployed event not found in transaction logs"))?;

                let new_policy_data = policy_data_deployed_event.data().policyData;
                info!("✓ New policy data deployed at: {}", new_policy_data);
                info!(
                    "Factory version: {}",
                    policy_data_deployed_event.data().implementationVersion
                );

                new_policy_data_addresses.push(new_policy_data);
            }

            info!("✓ All incompatible policy data contracts migrated successfully");
        }
    } else {
        info!("✓ All policy data contracts are compatible, no migration needed");
    }

    // Now deploy the new policy
    info!("Step 3b: Deploying new policy with latest factory version...");

    let new_policy_address = if dry_run {
        info!("DRY RUN MODE - Would deploy new policy with:");
        info!("- Factory version: {} (latest)", PROTOCOL_VERSION);
        info!("- Entrypoint: {}", policy_info.entrypoint);
        info!("- Policy CID: {}", policy_info.policy_cid);
        info!("- Schema CID: {}", policy_info.schema_cid);
        info!("- Policy Data: {} contracts", new_policy_data_addresses.len());
        info!("- Metadata CID: {}", policy_info.metadata_cid);
        info!("Note: This would call NewtonPolicyFactory.deployPolicy() with the above parameters");
        Address::ZERO // Placeholder for dry run
    } else {
        info!("Deploying new policy using latest factory...");

        // Parse private key and create signer
        let private_key_str = private_key.strip_prefix("0x").unwrap_or(&private_key);
        let signer = PrivateKeySigner::from_str(private_key_str).context("Failed to parse private key")?;
        let signer_address = signer.address();
        let wallet = EthereumWallet::from(signer);

        // Create provider with signer
        let url = Url::parse(&rpc_url).context("Invalid RPC URL")?;
        let provider = ProviderBuilder::new().wallet(wallet).connect_http(url);

        // Get policy factory address from the existing policy
        let policy_factory = get_policy_factory_for_policy(policy_address, &rpc_url).await?;
        info!("Using policy factory at: {}", policy_factory);

        let factory_contract = NewtonPolicyFactory::new(policy_factory, provider.clone());

        // Deploy new policy with same configuration and new policy data addresses
        info!("Submitting deployPolicy transaction...");
        let tx = factory_contract
            .deployPolicy(
                policy_info.entrypoint.clone(),
                policy_info.policy_cid.clone(),
                policy_info.schema_cid.clone(),
                new_policy_data_addresses.clone(),
                policy_info.metadata_cid.clone(),
                signer_address,
            )
            .send()
            .await
            .context("Failed to send deployPolicy transaction")?;

        info!("Transaction sent: {:?}", tx.tx_hash());
        info!("Waiting for transaction confirmation...");

        let receipt = tx.get_receipt().await.context("Failed to get transaction receipt")?;

        // Extract new policy address from PolicyDeployed event
        let logs = receipt.inner.logs();
        let policy_deployed_event = logs
            .iter()
            .find_map(|log| log.log_decode::<NewtonPolicyFactory::PolicyDeployed>().ok())
            .ok_or_else(|| eyre::eyre!("PolicyDeployed event not found in transaction logs"))?;

        let new_policy = policy_deployed_event.data().policy;
        info!("✓ New policy deployed at: {}", new_policy);
        info!(
            "Factory version: {}",
            policy_deployed_event.data().implementationVersion
        );

        new_policy
    };

    // Step 4: Update policy client to point to new policy
    info!("Step 4: Updating policy client to use new policy...");

    if dry_run {
        info!(
            "DRY RUN MODE - Would call setPolicyAddress({}) on policy client {}",
            new_policy_address, policy_client
        );
        info!("Note: This calls the owner-only setPolicyAddress() function on the existing policy client");
        info!("No redeployment needed - the policy client address stays the same");
    } else {
        // Parse private key and create signer (same credentials as Step 3b)
        let private_key_str = private_key.strip_prefix("0x").unwrap_or(&private_key);
        let signer = PrivateKeySigner::from_str(private_key_str).context("Failed to parse private key")?;
        let wallet = EthereumWallet::from(signer);
        let url = Url::parse(&rpc_url).context("Invalid RPC URL")?;
        let provider = ProviderBuilder::new().wallet(wallet).connect_http(url);

        let policy_client_contract = NewtonPolicyClient::new(policy_client, provider);

        info!(
            "Calling setPolicyAddress({}) on policy client {}...",
            new_policy_address, policy_client
        );

        let tx = policy_client_contract
            .setPolicyAddress(new_policy_address)
            .send()
            .await
            .context("Failed to send setPolicyAddress transaction. Ensure the signer is the policy client owner.")?;

        info!("Transaction sent: {:?}", tx.tx_hash());
        info!("Waiting for confirmation...");

        let receipt = tx
            .get_receipt()
            .await
            .context("Failed to get setPolicyAddress receipt")?;

        if !receipt.status() {
            return Err(eyre::eyre!(
                "setPolicyAddress transaction reverted (tx: {:?}). Possible causes:\n\
                 - Signer is not the policy client owner\n\
                 - New policy factory version is incompatible with TaskManager minimum",
                receipt.transaction_hash
            ));
        }

        info!(
            "✓ Policy client {} now points to new policy {}",
            policy_client, new_policy_address
        );
    }

    // Step 5: Verify migration
    info!("Step 5: Verifying migration...");

    if dry_run {
        info!("DRY RUN MODE - Would verify:");
        info!("✓ New policy is deployed and accessible");
        info!("✓ Policy client correctly references new policy");
        info!("✓ New policy version is compatible with protocol v{}", PROTOCOL_VERSION);
        info!("✓ All policy data is migrated or compatible");
        info!("Verification would include:");
        info!("1. Calling getPolicyAddress() on policy client");
        info!("2. Checking factory version of new policy");
        info!("3. Verifying policy configuration matches original");
        info!("4. Testing policy functionality with a test transaction");
    } else {
        info!("Verifying migration...");

        // 1. Verify policy client points to new policy
        let current_policy = get_policy_address_for_client(policy_client, rpc_url.clone()).await?;
        if current_policy != new_policy_address {
            return Err(eyre::eyre!(
                "Verification failed: Policy client points to {}, expected {}",
                current_policy,
                new_policy_address
            ));
        }
        info!("✓ Policy client correctly references new policy");

        // 2. Verify new policy factory version is compatible
        let new_policy_factory = get_policy_factory_for_policy(new_policy_address, &rpc_url).await?;
        let new_policy_version = get_policy_factory_version(new_policy_factory, &rpc_url).await?;

        if !is_compatible(&new_policy_version, MIN_COMPATIBLE_VERSION)? {
            return Err(eyre::eyre!(
                "Verification failed: New policy version {} is not compatible with minimum required {}",
                new_policy_version,
                MIN_COMPATIBLE_VERSION
            ));
        }
        info!("✓ New policy version {} is compatible", new_policy_version);

        // 3. Verify policy configuration matches original
        let new_policy_info = get_policy_info(new_policy_address, &rpc_url).await?;
        if new_policy_info.policy_cid != policy_info.policy_cid {
            warn!(
                "Policy CID mismatch: old={}, new={}",
                policy_info.policy_cid, new_policy_info.policy_cid
            );
        }
        if new_policy_info.schema_cid != policy_info.schema_cid {
            warn!(
                "Schema CID mismatch: old={}, new={}",
                policy_info.schema_cid, new_policy_info.schema_cid
            );
        }
        if new_policy_info.entrypoint != policy_info.entrypoint {
            warn!(
                "Entrypoint mismatch: old={}, new={}",
                policy_info.entrypoint, new_policy_info.entrypoint
            );
        }
        if new_policy_info.metadata_cid != policy_info.metadata_cid {
            warn!(
                "Metadata CID mismatch: old={}, new={}",
                policy_info.metadata_cid, new_policy_info.metadata_cid
            );
        }
        info!("✓ Policy configuration verified");

        // 4. Verify all policy data
        if new_policy_info.policy_data.len() != new_policy_data_addresses.len() {
            return Err(eyre::eyre!(
                "Verification failed: Policy data count mismatch. Expected {}, got {}",
                new_policy_data_addresses.len(),
                new_policy_info.policy_data.len()
            ));
        }

        for (i, addr) in new_policy_info.policy_data.iter().enumerate() {
            if let Ok(factory) = get_policy_data_factory_for_policy_data(*addr, &rpc_url).await {
                if let Ok(version) = get_policy_data_factory_version(factory, &rpc_url).await {
                    if !is_compatible(&version, MIN_COMPATIBLE_VERSION)? {
                        return Err(eyre::eyre!(
                            "Verification failed: Policy data {} version {} is not compatible",
                            i + 1,
                            version
                        ));
                    }
                }
            }
        }
        info!("✓ All policy data verified as compatible");

        info!("✓ Migration verified successfully!");
    }

    info!("=== Migration Summary ===");

    if dry_run {
        info!("✓ Dry run completed successfully");
        info!("The following would be performed:");
        info!("1. Deploy new policy at factory version {}", PROTOCOL_VERSION);
        if !incompatible_policy_data.is_empty() {
            info!(
                "2. Migrate {} incompatible policy data contracts",
                incompatible_policy_data.len()
            );
        } else {
            info!("2. Reuse all existing policy data (all compatible)");
        }
        info!("3. Update policy client {} to reference new policy", policy_client);
        info!("4. Verify migration success");
        info!("");
        info!("Remove --dry-run to execute the actual migration");
    } else {
        info!("✓ Migration completed successfully!");
        info!("Summary:");
        info!("- Old policy: {}", policy_address);
        info!("- New policy: {}", new_policy_address);
        info!("- Policy client: {}", policy_client);
        if !incompatible_policy_data.is_empty() {
            info!("- Migrated {} policy data contracts", incompatible_policy_data.len());
        }
    }

    Ok(())
}