vaea_flash_sdk/

retry.rs

/// VAEA Flash — Smart Retry (Rust)
///
/// Automatic transaction retry with:
/// - Blockhash refresh on expiry
/// - Priority fee escalation on congestion
/// - Never retries program errors (Custom(...), InstructionError)
///
/// Uses standard Solana RPC + optional Jito bundles.

use solana_sdk::{
    compute_budget::ComputeBudgetInstruction,
    instruction::Instruction,
    message::{v0, VersionedMessage},
    signature::Keypair,
    signer::Signer,
    transaction::VersionedTransaction,
    address_lookup_table::AddressLookupTableAccount,
    commitment_config::CommitmentConfig,
};
use solana_client::nonblocking::rpc_client::RpcClient;
use crate::jito::{JitoConfig, resolve_block_engine_url, calculate_tip, build_tip_instruction, send_jito_bundle, poll_bundle_status};

// ═══════════════════════════════════════════════════════════
//  Types
// ═══════════════════════════════════════════════════════════

/// Reason for retry.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum RetryReason {
    Expired,
    Congestion,
    ProgramError,
}

/// Retry configuration.
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of attempts (default: 3)
    pub max_attempts: u32,
    /// Retry strategy
    pub strategy: RetryStrategy,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum RetryStrategy {
    /// No retry — fail immediately
    None,
    /// Smart retry: blockhash refresh + fee escalation
    Adaptive,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            strategy: RetryStrategy::Adaptive,
        }
    }
}

/// Send mode: standard RPC or Jito bundle.
#[derive(Debug, Clone)]
pub enum SendMode {
    Rpc,
    Jito(JitoConfig),
}

// ═══════════════════════════════════════════════════════════
//  Smart Retry
// ═══════════════════════════════════════════════════════════

/// Send a transaction with smart retry logic.
///
/// - Refreshes blockhash on expiry
/// - Escalates priority fee on congestion (×1.5 per attempt)
/// - Never retries program errors
///
/// # Example
/// ```rust,no_run
/// use vaea_flash_sdk::retry::{send_with_retry, RetryConfig, SendMode};
///
/// let sig = send_with_retry(
///     &rpc, &wallet, instructions, &lookup_tables,
///     RetryConfig::default(), Some(1000), SendMode::Rpc,
/// ).await?;
/// ```
pub async fn send_with_retry(
    rpc: &RpcClient,
    wallet: &Keypair,
    instructions: Vec<Instruction>,
    lookup_tables: &[AddressLookupTableAccount],
    config: RetryConfig,
    initial_priority_micro_lamports: Option<u64>,
    send_mode: SendMode,
) -> Result<String, String> {
    if config.strategy == RetryStrategy::None {
        return send_once(rpc, wallet, &instructions, lookup_tables,
            initial_priority_micro_lamports.unwrap_or(0)).await;
    }

    let mut last_error = String::new();
    let mut priority = initial_priority_micro_lamports.unwrap_or(1_000);

    for attempt in 1..=config.max_attempts {
        let result = match &send_mode {
            SendMode::Rpc => {
                send_once(rpc, wallet, &instructions, lookup_tables, priority).await
            }
            SendMode::Jito(jito_config) => {
                send_once_via_jito(rpc, wallet, &instructions, lookup_tables, priority, jito_config).await
            }
        };

        match result {
            Ok(sig) => return Ok(sig),
            Err(err) => {
                let reason = classify_error(&err);
                last_error = err;

                // Never retry program errors
                if reason == RetryReason::ProgramError {
                    return Err(last_error);
                }
                if attempt >= config.max_attempts {
                    return Err(last_error);
                }

                // Escalate on congestion
                if reason == RetryReason::Congestion {
                    priority = (priority as f64 * 1.5) as u64;
                    let backoff = 400 * 2u64.pow(attempt - 1);
                    tokio::time::sleep(std::time::Duration::from_millis(backoff)).await;
                }
                // Expired: just rebuild with new blockhash (next iteration)
            }
        }
    }

    Err(last_error)
}

async fn send_once(
    rpc: &RpcClient,
    wallet: &Keypair,
    instructions: &[Instruction],
    lookup_tables: &[AddressLookupTableAccount],
    priority_micro_lamports: u64,
) -> Result<String, String> {
    let (blockhash, _last_valid_block_height) = rpc
        .get_latest_blockhash_with_commitment(CommitmentConfig::confirmed())
        .await
        .map_err(|e| e.to_string())?;

    let mut all_ixs: Vec<Instruction> = Vec::new();
    if priority_micro_lamports > 0 {
        all_ixs.push(ComputeBudgetInstruction::set_compute_unit_price(priority_micro_lamports));
    }
    all_ixs.extend_from_slice(instructions);

    let msg = v0::Message::try_compile(
        &wallet.pubkey(), &all_ixs, lookup_tables, blockhash,
    ).map_err(|e| e.to_string())?;

    let tx = VersionedTransaction::try_new(VersionedMessage::V0(msg), &[wallet])
        .map_err(|e| e.to_string())?;

    let sig = rpc.send_transaction(&tx)
        .await
        .map_err(|e| e.to_string())?;

    rpc.confirm_transaction_with_spinner(&sig, &blockhash, CommitmentConfig::confirmed())
        .await
        .map_err(|e| e.to_string())?;

    Ok(sig.to_string())
}

async fn send_once_via_jito(
    rpc: &RpcClient,
    wallet: &Keypair,
    instructions: &[Instruction],
    lookup_tables: &[AddressLookupTableAccount],
    priority_micro_lamports: u64,
    jito_config: &JitoConfig,
) -> Result<String, String> {
    let block_engine_url = resolve_block_engine_url(&jito_config.region);
    let tip_lamports = calculate_tip(&jito_config.tip, 10_000);

    let mut all_ixs: Vec<Instruction> = Vec::new();
    if priority_micro_lamports > 0 {
        all_ixs.push(ComputeBudgetInstruction::set_compute_unit_price(priority_micro_lamports));
    }
    all_ixs.extend_from_slice(instructions);
    all_ixs.push(build_tip_instruction(&wallet.pubkey(), tip_lamports));

    let (blockhash, _) = rpc
        .get_latest_blockhash_with_commitment(CommitmentConfig::confirmed())
        .await
        .map_err(|e| e.to_string())?;

    let msg = v0::Message::try_compile(
        &wallet.pubkey(), &all_ixs, lookup_tables, blockhash,
    ).map_err(|e| e.to_string())?;

    let tx = VersionedTransaction::try_new(VersionedMessage::V0(msg), &[wallet])
        .map_err(|e| e.to_string())?;

    let bundle_id = send_jito_bundle(&block_engine_url, &[tx]).await?;
    poll_bundle_status(&block_engine_url, &bundle_id, 30_000).await
}

fn classify_error(err: &str) -> RetryReason {
    if err.contains("Blockhash") || err.contains("expired") || err.contains("block height exceeded") {
        RetryReason::Expired
    } else if err.contains("InstructionError") || err.contains("Custom(") || err.contains("custom program error") {
        RetryReason::ProgramError
    } else {
        RetryReason::Congestion
    }
}