simulator_api/

lib.rs

use std::{
    collections::{BTreeMap, BTreeSet},
    fmt,
};

use base64::{
    DecodeError as Base64DecodeError, Engine as _, engine::general_purpose::STANDARD as BASE64,
};
use serde::{Deserialize, Serialize};
use solana_address::Address;

/// Backtest RPC methods exposed to the client.
#[derive(Debug, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(tag = "method", content = "params", rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub enum BacktestRequest {
    CreateBacktestSession(CreateBacktestSessionRequest),
    Continue(ContinueParams),
    ContinueSessionV1(ContinueSessionRequestV1),
    CloseBacktestSession,
    CloseSessionV1(CloseSessionRequestV1),
    AttachBacktestSession {
        session_id: String,
        /// Last sequence number the client received. Responses after this sequence
        /// will be replayed from the session's buffer. None = replay entire buffer.
        last_sequence: Option<u64>,
    },
}

/// Versioned payload for `CreateBacktestSession`.
///
/// - `V0` keeps backwards-compatible shape by using `CreateSessionParams` directly.
/// - `V1` keeps the same shape and adds `parallel`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(untagged)]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub enum CreateBacktestSessionRequest {
    V1(CreateBacktestSessionRequestV1),
    V0(CreateSessionParams),
}

impl CreateBacktestSessionRequest {
    pub fn into_request_options(self) -> CreateBacktestSessionRequestOptions {
        match self {
            Self::V0(request) => CreateBacktestSessionRequestOptions {
                request,
                parallel: false,
            },
            Self::V1(CreateBacktestSessionRequestV1 { request, parallel }) => {
                CreateBacktestSessionRequestOptions { request, parallel }
            }
        }
    }

    pub fn into_request_and_parallel(self) -> (CreateSessionParams, bool) {
        let options = self.into_request_options();
        (options.request, options.parallel)
    }
}

impl From<CreateSessionParams> for CreateBacktestSessionRequest {
    fn from(value: CreateSessionParams) -> Self {
        Self::V0(value)
    }
}

impl From<CreateBacktestSessionRequestV1> for CreateBacktestSessionRequest {
    fn from(value: CreateBacktestSessionRequestV1) -> Self {
        Self::V1(value)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub struct CreateBacktestSessionRequestV1 {
    #[serde(flatten)]
    pub request: CreateSessionParams,
    pub parallel: bool,
}

#[derive(Debug, Clone)]
pub struct CreateBacktestSessionRequestOptions {
    pub request: CreateSessionParams,
    pub parallel: bool,
}

#[derive(Debug, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub struct ContinueSessionRequestV1 {
    pub session_id: String,
    pub request: ContinueParams,
}

#[derive(Debug, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub struct CloseSessionRequestV1 {
    pub session_id: String,
}

/// Parameters required to start a new backtest session.
#[serde_with::serde_as]
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
pub struct CreateSessionParams {
    /// First slot (inclusive) to replay.
    pub start_slot: u64,
    /// Last slot (inclusive) to replay.
    pub end_slot: u64,
    #[serde_as(as = "BTreeSet<serde_with::DisplayFromStr>")]
    #[serde(default)]
    #[cfg_attr(feature = "ts-rs", ts(as = "Vec<String>"))]
    /// Skip transactions signed by these addresses.
    pub signer_filter: BTreeSet<Address>,
    #[serde_as(as = "BTreeSet<serde_with::DisplayFromStr>")]
    #[serde(default)]
    #[cfg_attr(feature = "ts-rs", ts(as = "Vec<String>"))]
    /// Programs to preload before executing.
    pub preload_programs: BTreeSet<Address>,
    /// Account bundle IDs to preload before executing.
    #[serde(default)]
    pub preload_account_bundles: Vec<String>,
    /// When true, include a session summary with transaction statistics in client-facing
    /// `Completed` responses. Summary generation remains enabled internally for metrics.
    #[serde(default)]
    pub send_summary: bool,
    /// Maximum seconds to wait for ECS capacity-related startup retries before
    /// failing session creation. If not set (or 0), capacity errors fail immediately.
    #[serde(default)]
    #[cfg_attr(feature = "ts-rs", ts(optional))]
    pub capacity_wait_timeout_secs: Option<u16>,
    /// Maximum seconds to keep the session alive after the control websocket disconnects.
    /// If not set (or 0), the session tears down immediately on disconnect.
    /// Maximum value: 900 (15 minutes).
    #[serde(default)]
    pub disconnect_timeout_secs: Option<u16>,
    /// Extra compute units to add to each transaction's `SetComputeUnitLimit` budget.
    /// Useful when replaying with an account override whose program uses more CU than
    /// the original, causing otherwise-healthy transactions to run out of budget.
    /// Only applied when a `SetComputeUnitLimit` instruction is already present.
    #[serde(default)]
    pub extra_compute_units: Option<u32>,
    /// Agent configurations to run as sidecars alongside this session.
    #[serde(default)]
    pub agents: Vec<AgentParams>,
}

/// Available agent types for sidecar participation in backtest sessions.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub enum AgentType {
    Arb,
}

/// Parameters for a circular arbitrage route.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub struct ArbRouteParams {
    pub base_mint: String,
    pub temp_mint: String,
    #[serde(default)]
    pub buy_dexes: Vec<String>,
    #[serde(default)]
    pub sell_dexes: Vec<String>,
    pub min_input: u64,
    pub max_input: u64,
    #[serde(default)]
    pub min_profit: u64,
}

/// Configuration for an agent to run alongside a backtest session.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub struct AgentParams {
    pub agent_type: AgentType,
    pub wallet: Option<String>,
    pub seed_sol_lamports: Option<u64>,
    #[serde(default)]
    pub arb_routes: Vec<ArbRouteParams>,
}

/// Account state modifications to apply.
#[serde_with::serde_as]
#[derive(Debug, Serialize, Deserialize, Default)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
pub struct AccountModifications(
    #[serde_as(as = "BTreeMap<serde_with::DisplayFromStr, _>")]
    #[serde(default)]
    #[cfg_attr(feature = "ts-rs", ts(as = "BTreeMap<String, AccountData>"))]
    pub BTreeMap<Address, AccountData>,
);

/// Arguments used to continue an existing session.
#[serde_with::serde_as]
#[derive(Debug, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
pub struct ContinueParams {
    #[serde(default = "ContinueParams::default_advance_count")]
    /// Number of blocks to advance before waiting.
    pub advance_count: u64,
    #[serde(default)]
    /// Base64-encoded transactions to execute before advancing.
    pub transactions: Vec<String>,
    #[serde(default)]
    /// Account state overrides to apply ahead of execution.
    pub modify_account_states: AccountModifications,
}

impl Default for ContinueParams {
    fn default() -> Self {
        Self {
            advance_count: Self::default_advance_count(),
            transactions: Vec::new(),
            modify_account_states: AccountModifications(BTreeMap::new()),
        }
    }
}

impl ContinueParams {
    pub fn default_advance_count() -> u64 {
        1
    }
}

/// Supported binary encodings for account/transaction payloads.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "lowercase")]
pub enum BinaryEncoding {
    Base64,
}

impl BinaryEncoding {
    pub fn encode(self, bytes: &[u8]) -> String {
        match self {
            Self::Base64 => BASE64.encode(bytes),
        }
    }

    pub fn decode(self, data: &str) -> Result<Vec<u8>, Base64DecodeError> {
        match self {
            Self::Base64 => BASE64.decode(data),
        }
    }
}

/// A blob paired with the encoding needed to decode it.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
pub struct EncodedBinary {
    /// Encoded payload.
    pub data: String,
    /// Encoding scheme used for the payload.
    pub encoding: BinaryEncoding,
}

impl EncodedBinary {
    pub fn new(data: String, encoding: BinaryEncoding) -> Self {
        Self { data, encoding }
    }

    pub fn from_bytes(bytes: &[u8], encoding: BinaryEncoding) -> Self {
        Self {
            data: encoding.encode(bytes),
            encoding,
        }
    }

    pub fn decode(&self) -> Result<Vec<u8>, Base64DecodeError> {
        self.encoding.decode(&self.data)
    }
}

/// Account snapshot used to seed or modify state in a session.
#[serde_with::serde_as]
#[derive(Debug, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
pub struct AccountData {
    /// Account data bytes and encoding.
    pub data: EncodedBinary,
    /// Whether the account is executable.
    pub executable: bool,
    /// Lamport balance.
    pub lamports: u64,
    #[serde_as(as = "serde_with::DisplayFromStr")]
    #[cfg_attr(feature = "ts-rs", ts(as = "String"))]
    /// Account owner pubkey.
    pub owner: Address,
    /// Allocated space in bytes.
    pub space: u64,
}

/// Responses returned over the backtest RPC channel.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(tag = "method", content = "params", rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub enum BacktestResponse {
    SessionCreated {
        session_id: String,
        rpc_endpoint: String,
    },
    SessionAttached {
        session_id: String,
        rpc_endpoint: String,
    },
    SessionsCreated {
        session_ids: Vec<String>,
    },
    ReadyForContinue,
    SlotNotification(u64),
    Error(BacktestError),
    Success,
    Completed {
        /// Session summary with transaction statistics.
        /// The session always computes this summary, but management may omit it from
        /// client-facing responses unless `send_summary` was requested at session creation.
        #[serde(skip_serializing_if = "Option::is_none")]
        summary: Option<SessionSummary>,
    },
    Status {
        status: BacktestStatus,
    },
    SessionEventV1 {
        session_id: String,
        event: SessionEventV1,
    },
}

impl BacktestResponse {
    pub fn is_completed(&self) -> bool {
        matches!(self, BacktestResponse::Completed { .. })
    }

    pub fn is_terminal(&self) -> bool {
        match self {
            BacktestResponse::Completed { .. } => true,
            BacktestResponse::Error(e) => matches!(
                e,
                BacktestError::NoMoreBlocks
                    | BacktestError::AdvanceSlotFailed { .. }
                    | BacktestError::Internal { .. }
            ),
            _ => false,
        }
    }
}

impl From<BacktestStatus> for BacktestResponse {
    fn from(status: BacktestStatus) -> Self {
        Self::Status { status }
    }
}

impl From<String> for BacktestResponse {
    fn from(message: String) -> Self {
        BacktestError::Internal { error: message }.into()
    }
}

impl From<&str> for BacktestResponse {
    fn from(message: &str) -> Self {
        BacktestError::Internal {
            error: message.to_string(),
        }
        .into()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(tag = "method", content = "params", rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub enum SessionEventV1 {
    ReadyForContinue,
    SlotNotification(u64),
    Error(BacktestError),
    Success,
    Completed {
        #[serde(skip_serializing_if = "Option::is_none")]
        summary: Option<SessionSummary>,
    },
    Status {
        status: BacktestStatus,
    },
}

/// Wire format wrapper for responses sent over the control websocket.
/// Sessions with `disconnect_timeout_secs > 0` use this to track client position.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "ts-rs", ts(export))]
pub struct SequencedResponse {
    pub seq_id: u64,
    #[serde(flatten)]
    pub response: BacktestResponse,
}

/// High-level progress states during a `Continue` call.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
pub enum BacktestStatus {
    /// Runtime startup is in progress.
    StartingRuntime,
    DecodedTransactions,
    AppliedAccountModifications,
    ReadyToExecuteUserTransactions,
    ExecutedUserTransactions,
    ExecutingBlockTransactions,
    ExecutedBlockTransactions,
    ProgramAccountsLoaded,
}

/// Summary of transaction execution statistics for a completed backtest session.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
pub struct SessionSummary {
    /// Number of simulations where simulator outcome matched on-chain outcome
    /// (`true_success + true_failure`).
    pub correct_simulation: usize,
    /// Number of simulations where simulator outcome did not match on-chain outcome
    /// (`false_success + false_failure`).
    pub incorrect_simulation: usize,
    /// Number of transactions that had execution errors in simulation.
    pub execution_errors: usize,
    /// Number of transactions with different balance diffs.
    pub balance_diff: usize,
    /// Number of transactions with different log diffs.
    pub log_diff: usize,
}

impl SessionSummary {
    /// Returns true if there were any execution deviations (errors or mismatched results).
    pub fn has_deviations(&self) -> bool {
        self.incorrect_simulation > 0 || self.execution_errors > 0 || self.balance_diff > 0
    }

    /// Total number of transactions processed.
    pub fn total_transactions(&self) -> usize {
        self.correct_simulation
            + self.incorrect_simulation
            + self.execution_errors
            + self.balance_diff
            + self.log_diff
    }
}

impl std::fmt::Display for SessionSummary {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let total = self.total_transactions();
        write!(
            f,
            "Session summary: {total} transactions\n\
             \x20  - {} correct simulation\n\
             \x20  - {} incorrect simulation\n\
             \x20  - {} execution errors\n\
             \x20  - {} balance diffs\n\
             \x20  - {} log diffs",
            self.correct_simulation,
            self.incorrect_simulation,
            self.execution_errors,
            self.balance_diff,
            self.log_diff,
        )
    }
}

/// Error variants surfaced to backtest RPC clients.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "ts-rs", derive(ts_rs::TS))]
#[serde(rename_all = "camelCase")]
pub enum BacktestError {
    InvalidTransactionEncoding {
        index: usize,
        error: String,
    },
    InvalidTransactionFormat {
        index: usize,
        error: String,
    },
    InvalidAccountEncoding {
        address: String,
        encoding: BinaryEncoding,
        error: String,
    },
    InvalidAccountOwner {
        address: String,
        error: String,
    },
    InvalidAccountPubkey {
        address: String,
        error: String,
    },
    NoMoreBlocks,
    AdvanceSlotFailed {
        slot: u64,
        error: String,
    },
    InvalidRequest {
        error: String,
    },
    Internal {
        error: String,
    },
    InvalidBlockhashFormat {
        slot: u64,
        error: String,
    },
    InitializingSysvarsFailed {
        slot: u64,
        error: String,
    },
    ClerkError {
        error: String,
    },
    SimulationError {
        error: String,
    },
    SessionNotFound {
        session_id: String,
    },
    SessionOwnerMismatch,
}

/// One contiguous block range available on the history clerk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AvailableRange {
    pub bundle_start_slot: u64,
    pub bundle_start_slot_utc: Option<String>,
    pub max_bundle_end_slot: Option<u64>,
    pub max_bundle_end_slot_utc: Option<String>,
}

/// Split a user-requested `[start_slot, end_slot]` range across the available
/// bundle ranges, returning a list of non-overlapping `(start, end)` pairs — one
/// per bundle that intersects the requested range.
///
/// Returns an error if the requested start slot is not covered by any range.
pub fn split_range(
    ranges: &[AvailableRange],
    requested_start: u64,
    requested_end: u64,
) -> Result<Vec<(u64, u64)>, String> {
    if requested_end < requested_start {
        return Err(format!(
            "invalid range: start_slot {requested_start} > end_slot {requested_end}"
        ));
    }

    let mut candidates: Vec<(u64, u64)> = ranges
        .iter()
        .filter_map(|r| Some((r.bundle_start_slot, r.max_bundle_end_slot?)))
        .filter(|(start, end)| end >= start)
        .collect();

    candidates.sort_by(|a, b| a.0.cmp(&b.0).then(b.1.cmp(&a.1)));
    candidates.dedup_by_key(|(start, _)| *start);

    if candidates.is_empty() {
        return Err("no available bundle ranges found on server".to_string());
    }

    let mut non_overlapping: Vec<(u64, u64)> = Vec::with_capacity(candidates.len());
    for (i, (start, mut end)) in candidates.iter().copied().enumerate() {
        if let Some((next_start, _)) = candidates.get(i + 1).copied()
            && next_start <= end
        {
            end = next_start.saturating_sub(1);
        }
        if end >= start {
            non_overlapping.push((start, end));
        }
    }

    let anchor = non_overlapping
        .iter()
        .enumerate()
        .rev()
        .find(|(_, (s, e))| *s <= requested_start && *e >= requested_start)
        .map(|(i, _)| i)
        .ok_or_else(|| {
            format!("start_slot {requested_start} is not covered by any available bundle range")
        })?;

    let mut result = Vec::new();
    for (start, range_end) in non_overlapping.into_iter().skip(anchor) {
        if start > requested_end {
            break;
        }
        let end = range_end.min(requested_end);
        if end >= start {
            result.push((start, end));
        }
        if end == requested_end {
            break;
        }
    }

    if result.is_empty() {
        return Err(format!(
            "no available bundle ranges intersect requested range [{requested_start}-{requested_end}]"
        ));
    }

    Ok(result)
}

impl From<BacktestError> for BacktestResponse {
    fn from(error: BacktestError) -> Self {
        Self::Error(error)
    }
}

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

impl fmt::Display for BacktestError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BacktestError::InvalidTransactionEncoding { index, error } => {
                write!(f, "invalid transaction encoding at index {index}: {error}")
            }
            BacktestError::InvalidTransactionFormat { index, error } => {
                write!(f, "invalid transaction format at index {index}: {error}")
            }
            BacktestError::InvalidAccountEncoding {
                address,
                encoding,
                error,
            } => write!(
                f,
                "invalid encoding for account {address} ({encoding:?}): {error}"
            ),
            BacktestError::InvalidAccountOwner { address, error } => {
                write!(f, "invalid owner for account {address}: {error}")
            }
            BacktestError::InvalidAccountPubkey { address, error } => {
                write!(f, "invalid account pubkey {address}: {error}")
            }
            BacktestError::NoMoreBlocks => write!(f, "no more blocks available"),
            BacktestError::AdvanceSlotFailed { slot, error } => {
                write!(f, "failed to advance to slot {slot}: {error}")
            }
            BacktestError::InvalidRequest { error } => write!(f, "invalid request: {error}"),
            BacktestError::Internal { error } => write!(f, "internal error: {error}"),
            BacktestError::InvalidBlockhashFormat { slot, error } => {
                write!(f, "invalid blockhash at slot {slot}: {error}")
            }
            BacktestError::InitializingSysvarsFailed { slot, error } => {
                write!(f, "failed to initialize sysvars at slot {slot}: {error}")
            }
            BacktestError::ClerkError { error } => write!(f, "clerk error: {error}"),
            BacktestError::SimulationError { error } => {
                write!(f, "simulation error: {error}")
            }
            BacktestError::SessionNotFound { session_id } => {
                write!(f, "session not found: {session_id}")
            }
            BacktestError::SessionOwnerMismatch => {
                write!(f, "session owner mismatch")
            }
        }
    }
}