khadhroony-solana-project/crates/ksp-onchain-transport-lib/src/ws_accounts.rs

// file: crates/ksp-onchain-transport-lib/src/ws_accounts.rs
// version: 3

const MAX_PROGRAM_SUBSCRIBE_FILTERS: usize = 4;
const MAX_PROGRAM_SUBSCRIBE_RAW_MEMCMP_BYTES: usize = 128;

/// Configuration accepted by the standard Solana `accountSubscribe` WebSocket method.
///
/// `minContextSlot` is deliberately absent: Agave `v4.2.1` carries that field in the shared account config but the PubSub handler ignores it, so KSP does
/// not expose it as an effective WebSocket option.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct SolanaAccountSubscribeConfig {
    encoding: std::option::Option<crate::SolanaAccountEncoding>,
    data_slice: std::option::Option<crate::SolanaDataSliceConfig>,
    commitment: std::option::Option<crate::SolanaCommitment>,
}

impl SolanaAccountSubscribeConfig {
    /// Creates an explicit `accountSubscribe` configuration.
    #[must_use]
    pub const fn new(
        encoding: std::option::Option<crate::SolanaAccountEncoding>,
        data_slice: std::option::Option<crate::SolanaDataSliceConfig>,
        commitment: std::option::Option<crate::SolanaCommitment>,
    ) -> Self {
        return Self { encoding, data_slice, commitment };
    }

    /// Returns the optional account-data encoding.
    #[must_use]
    pub const fn encoding(&self) -> std::option::Option<crate::SolanaAccountEncoding> {
        return self.encoding;
    }

    /// Returns the optional account-data slice.
    #[must_use]
    pub const fn data_slice(&self) -> std::option::Option<crate::SolanaDataSliceConfig> {
        return self.data_slice;
    }

    /// Returns the optional commitment level.
    #[must_use]
    pub const fn commitment(&self) -> std::option::Option<crate::SolanaCommitment> {
        return self.commitment;
    }

    fn is_empty(&self) -> bool {
        return self.encoding.is_none() && self.data_slice.is_none() && self.commitment.is_none();
    }

    fn to_json_value(&self) -> serde_json::Value {
        let mut object = serde_json::Map::new();
        if let std::option::Option::Some(encoding) = self.encoding {
            object.insert("encoding".to_owned(), serde_json::Value::String(encoding.as_str().to_owned()));
        }
        if let std::option::Option::Some(data_slice) = self.data_slice {
            object.insert("dataSlice".to_owned(), data_slice.to_json_value());
        }
        if let std::option::Option::Some(commitment) = self.commitment {
            object.insert("commitment".to_owned(), serde_json::Value::String(commitment.as_str().to_owned()));
        }
        return serde_json::Value::Object(object);
    }
}

/// Configuration accepted by the standard Solana `programSubscribe` WebSocket method.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct SolanaProgramSubscribeConfig {
    account: crate::SolanaAccountSubscribeConfig,
    filters: std::vec::Vec<crate::SolanaProgramAccountFilter>,
    with_context: std::option::Option<bool>,
}

impl SolanaProgramSubscribeConfig {
    /// Creates an explicit `programSubscribe` configuration.
    #[must_use]
    pub fn new(
        account: crate::SolanaAccountSubscribeConfig,
        filters: std::vec::Vec<crate::SolanaProgramAccountFilter>,
        with_context: std::option::Option<bool>,
    ) -> Self {
        return Self { account, filters, with_context };
    }

    /// Returns the shared WebSocket account configuration.
    #[must_use]
    pub const fn account(&self) -> &crate::SolanaAccountSubscribeConfig {
        return &self.account;
    }

    /// Returns the ordered program-account filters.
    #[must_use]
    pub fn filters(&self) -> &[crate::SolanaProgramAccountFilter] {
        return self.filters.as_slice();
    }

    /// Returns the optional `withContext` request; omission uses the upstream default `false`.
    #[must_use]
    pub const fn with_context(&self) -> std::option::Option<bool> {
        return self.with_context;
    }

    fn is_empty(&self) -> bool {
        return self.account.is_empty() && self.filters.is_empty() && self.with_context.is_none();
    }

    fn to_json_value(&self) -> serde_json::Value {
        let account = self.account.to_json_value();
        let mut object = match account {
            serde_json::Value::Object(object) => object,
            _ => serde_json::Map::new(),
        };
        if !self.filters.is_empty() {
            let filters = self.filters.iter().map(crate::SolanaProgramAccountFilter::to_json_value).collect::<std::vec::Vec<_>>();
            object.insert("filters".to_owned(), serde_json::Value::Array(filters));
        }
        if let std::option::Option::Some(with_context) = self.with_context {
            object.insert("withContext".to_owned(), serde_json::Value::Bool(with_context));
        }
        return serde_json::Value::Object(object);
    }
}

/// One `programNotification` payload, preserving whether the upstream wire result was contextualized.
#[derive(Clone, Debug, PartialEq)]
pub enum SolanaProgramNotification {
    /// Program account payload without a surrounding RPC context.
    Account(crate::SolanaKeyedAccount),
    /// Program account payload wrapped in an RPC context.
    Context(crate::SolanaRpcResponse<crate::SolanaKeyedAccount>),
}

impl SolanaProgramNotification {
    /// Returns the program account regardless of the upstream context-wrapper form.
    #[must_use]
    pub const fn account(&self) -> &crate::SolanaKeyedAccount {
        return match self {
            Self::Account(account) => account,
            Self::Context(response) => response.value(),
        };
    }

    /// Returns the RPC context when the upstream notification included one.
    #[must_use]
    pub const fn context(&self) -> std::option::Option<&crate::SolanaRpcContext> {
        return match self {
            Self::Account(_) => std::option::Option::None,
            Self::Context(response) => std::option::Option::Some(response.context()),
        };
    }
}

impl crate::WsSession {
    /// Subscribes to changes for one Solana account through standard `accountSubscribe`.
    pub async fn account_subscribe(
        &self,
        account: &ksp_core_lib::Pubkey,
        config: std::option::Option<&crate::SolanaAccountSubscribeConfig>,
    ) -> ksp_core_lib::Result<crate::WsSubscription<crate::SolanaRpcResponse<crate::SolanaAccount>>> {
        let mut params = std::vec![serde_json::Value::String(account.to_string())];
        if let std::option::Option::Some(config) = config
            && !config.is_empty()
        {
            params.push(config.to_json_value());
        }
        return self
            .subscribe_typed(crate::WsSubscriptionKind::Account, params, |value| return decode_account_notification("accountSubscribe", value))
            .await;
    }

    /// Subscribes to account changes owned by one Solana program through standard `programSubscribe`.
    pub async fn program_subscribe(
        &self,
        program_id: &ksp_core_lib::Pubkey,
        config: std::option::Option<&crate::SolanaProgramSubscribeConfig>,
    ) -> ksp_core_lib::Result<crate::WsSubscription<crate::SolanaProgramNotification>> {
        if let std::option::Option::Some(config) = config {
            let validation = validate_program_subscribe_filters(config.filters());
            if let std::result::Result::Err(error) = validation {
                return std::result::Result::Err(error);
            }
        }
        let mut params = std::vec![serde_json::Value::String(program_id.to_string())];
        if let std::option::Option::Some(config) = config
            && !config.is_empty()
        {
            params.push(config.to_json_value());
        }
        return self
            .subscribe_typed(crate::WsSubscriptionKind::Program, params, |value| return decode_program_notification("programSubscribe", value))
            .await;
    }
}

#[derive(serde::Deserialize)]
struct WireRpcResponse {
    context: serde_json::Value,
    value: serde_json::Value,
}

#[derive(serde::Deserialize)]
#[serde(untagged)]
enum WireProgramNotification {
    Context(WireRpcResponse),
    Account(serde_json::Value),
}

fn decode_account_notification(method: &str, value: serde_json::Value) -> ksp_core_lib::Result<crate::SolanaRpcResponse<crate::SolanaAccount>> {
    let decoded = crate::decode_wire_json::<WireRpcResponse>(method, value);
    let wire = match decoded {
        std::result::Result::Ok(wire) => wire,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let context = crate::SolanaRpcContext::decode_wire(method, wire.context);
    let context = match context {
        std::result::Result::Ok(context) => context,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let account = crate::SolanaAccount::decode_wire(method, wire.value);
    let account = match account {
        std::result::Result::Ok(account) => account,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    return std::result::Result::Ok(crate::SolanaRpcResponse::new(context, account));
}

fn decode_program_notification(method: &str, value: serde_json::Value) -> ksp_core_lib::Result<crate::SolanaProgramNotification> {
    let decoded = crate::decode_wire_json::<WireProgramNotification>(method, value);
    let wire = match decoded {
        std::result::Result::Ok(wire) => wire,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    return match wire {
        WireProgramNotification::Account(value) => {
            let account = crate::SolanaKeyedAccount::decode_wire(method, value);
            match account {
                std::result::Result::Ok(account) => std::result::Result::Ok(crate::SolanaProgramNotification::Account(account)),
                std::result::Result::Err(error) => std::result::Result::Err(error),
            }
        },
        WireProgramNotification::Context(wire) => {
            let context = crate::SolanaRpcContext::decode_wire(method, wire.context);
            let context = match context {
                std::result::Result::Ok(context) => context,
                std::result::Result::Err(error) => return std::result::Result::Err(error),
            };
            let account = crate::SolanaKeyedAccount::decode_wire(method, wire.value);
            let account = match account {
                std::result::Result::Ok(account) => account,
                std::result::Result::Err(error) => return std::result::Result::Err(error),
            };
            std::result::Result::Ok(crate::SolanaProgramNotification::Context(crate::SolanaRpcResponse::new(context, account)))
        },
    };
}

fn validate_program_subscribe_filters(filters: &[crate::SolanaProgramAccountFilter]) -> ksp_core_lib::Result<()> {
    if filters.len() > MAX_PROGRAM_SUBSCRIBE_FILTERS {
        return std::result::Result::Err(
            ksp_core_lib::Error::new(crate::ERROR_CODE_INVALID_RPC_PARAMETERS, "programSubscribe accepts at most 4 filters on the targeted Agave runtime")
                .with_context("rpc_method", "programSubscribe")
                .with_context("filter_count", filters.len().to_string()),
        );
    }
    for filter in filters {
        if let crate::SolanaProgramAccountFilter::Memcmp(memcmp) = filter
            && let crate::SolanaMemcmpBytes::Bytes(bytes) = memcmp.bytes()
            && bytes.len() > MAX_PROGRAM_SUBSCRIBE_RAW_MEMCMP_BYTES
        {
            return std::result::Result::Err(
                ksp_core_lib::Error::new(crate::ERROR_CODE_INVALID_RPC_PARAMETERS, "raw programSubscribe memcmp data accepts at most 128 bytes")
                    .with_context("rpc_method", "programSubscribe")
                    .with_context("memcmp_byte_count", bytes.len().to_string()),
            );
        }
    }
    return std::result::Result::Ok(());
}

impl crate::SolanaStandardWsSession {
    /// Subscribes to changes for one Solana account through standard `accountSubscribe`.
    pub async fn account_subscribe(
        &self,
        account: &ksp_core_lib::Pubkey,
        config: std::option::Option<&crate::SolanaAccountSubscribeConfig>,
    ) -> ksp_core_lib::Result<crate::WsSubscription<crate::SolanaRpcResponse<crate::SolanaAccount>>> {
        return self.physical_session().account_subscribe(account, config).await;
    }

    /// Subscribes to account changes owned by one Solana program through standard `programSubscribe`.
    pub async fn program_subscribe(
        &self,
        program_id: &ksp_core_lib::Pubkey,
        config: std::option::Option<&crate::SolanaProgramSubscribeConfig>,
    ) -> ksp_core_lib::Result<crate::WsSubscription<crate::SolanaProgramNotification>> {
        return self.physical_session().program_subscribe(program_id, config).await;
    }
}

impl crate::HeliusLaserStreamWsSession {
    /// Subscribes to account changes through the standard `accountSubscribe` wire supported by Helius LaserStream WebSocket.
    pub async fn account_subscribe(
        &self,
        account: &ksp_core_lib::Pubkey,
        config: std::option::Option<&crate::SolanaAccountSubscribeConfig>,
    ) -> ksp_core_lib::Result<crate::WsSubscription<crate::SolanaRpcResponse<crate::SolanaAccount>>> {
        return self.physical_session().account_subscribe(account, config).await;
    }

    /// Subscribes to program-owned account changes through the standard `programSubscribe` wire supported by Helius LaserStream WebSocket.
    pub async fn program_subscribe(
        &self,
        program_id: &ksp_core_lib::Pubkey,
        config: std::option::Option<&crate::SolanaProgramSubscribeConfig>,
    ) -> ksp_core_lib::Result<crate::WsSubscription<crate::SolanaProgramNotification>> {
        return self.physical_session().program_subscribe(program_id, config).await;
    }
}

#[cfg(test)]
#[path = "../unit_tests/ws_accounts.rs"]
mod tests;