newton-chainio 0.5.2

use crate::{avs::errors::is_transient_rpc_error, provider::get_provider_with_retry};
use alloy::{
    primitives::{Address, Bytes},
    providers::Provider,
};
use eigensdk::{
    client_avsregistry::reader::{AvsRegistryChainReader, AvsRegistryReader},
    crypto_bls::{
        alloy_registry_g1_point_to_g1_affine, alloy_registry_g2_point_to_g2_affine, BlsG1Point, BlsG2Point, OperatorId,
    },
    types::operator::{operator_id_from_g1_pub_key, OperatorPubKeys},
    utils::slashing::middleware::{
        bls_apk_registry::BLSApkRegistry, operator_state_retriever::OperatorStateRetriever,
        socket_registry::SocketRegistry,
    },
};
use futures::stream::{StreamExt, TryStreamExt};
use newton_core::operator_registry::OperatorRegistry;
use std::{
    collections::HashMap,
    sync::{atomic::Ordering, Arc},
};
use tokio::sync::{mpsc::UnboundedSender, oneshot::Sender, RwLock};
use tracing::{debug, error, info, instrument, warn};

use crate::registry::operator;

/// Maximum number of per-operator RPC chains issued concurrently during
/// operator discovery.
///
/// Each registered operator triggers a short chain of contract reads (id +
/// pubkeys, id + socket, id + per-quorum stakes), previously launched all at
/// once via `join_all`. This cap is **preventative hygiene for larger operator
/// sets**: at the handful of operators in the deployments today it was never
/// the binding factor (a 2-operator set fanned out 2 chains, well under this
/// cap), so it was not what triggered the reported 429. The operative controls
/// for rate limiting live in the transport (`crate::provider`): a `ThrottleLayer`
/// paces requests to a per-gateway budget and a `RetryBackoffLayer` retries any
/// residual 429s. This cap simply keeps the burst flat as the set grows.
const MAX_CONCURRENT_OPERATOR_QUERIES: usize = 8;

/// Decide how a per-operator fetch failure should affect the whole-tick result.
///
/// The discovery fan-outs must distinguish two failure modes that look alike at
/// the call site but demand opposite handling:
///
/// * **Permanent contract reverts** — e.g. a whitelisted operator that has not
///   yet registered its BLS pubkey, so `getRegisteredPubkey` reverts. Re-running
///   against the same on-chain state always reverts the same way, so we drop
///   just that operator and keep the rest of the snapshot. (This is the original
///   `filter_map(.ok())` behavior, preserved for this case.)
/// * **Transient errors** — timeouts, connection drops, rate limits. These say
///   "we did not actually observe this operator", so a partial fetch must abort
///   the whole tick (`Err`) rather than masquerade as a smaller-but-complete
///   whitelist — otherwise the caller prunes healthy-but-unfetched operators and
///   collapses the pool toward 0.
///
/// Returns `Ok(true)` to skip the operator and continue, or `Err` to abort.
fn skip_or_abort_operator(operator: Address, context: &str, err: eyre::Error) -> eyre::Result<bool> {
    if is_transient_rpc_error(&err.to_string()) {
        Err(err)
    } else {
        debug!(operator = %operator, "[{context}] skipping operator on non-transient error: {err}");
        Ok(true)
    }
}

/// A registered operator captured from a single block-pinned `OperatorStateRetriever`
/// read: its address, `operatorId`, and per-quorum stake — all consistent at one block.
#[derive(Debug, Clone)]
pub(crate) struct RegisteredOperator {
    /// Operator address (cache key for pubkeys/sockets).
    pub address: Address,
    /// Operator id (cache key for stakes/sockets in the id-keyed maps).
    pub operator_id: OperatorId,
    /// Stake per quorum number, as reported inline by the retriever at the pinned block.
    /// A quorum the operator is registered for but holds zero stake in is present with `0`.
    pub stakes: HashMap<u8, u128>,
}

/// A block-consistent snapshot of the on-chain **registered** operator set.
#[derive(Debug)]
pub(crate) struct RegisteredSnapshot {
    /// Number of quorums (`quorumCount`) observed for this snapshot.
    pub quorum_count: u8,
    /// The deduplicated registered operators, each carrying its per-quorum stake.
    pub operators: Vec<RegisteredOperator>,
    /// Per-quorum total registered stake, summed from the retriever's inline stakes.
    pub total_stakes: HashMap<u8, u128>,
    /// The block this snapshot was pinned to. Threaded into the stake messages so the
    /// receiver can reconcile snapshot vs. event by block (last-writer-wins).
    pub block_number: u64,
}

/// Snapshot the **registered** operator set from on-chain AVS state at a single block.
///
/// NEWT-1715: discovery historically sourced membership from
/// `OperatorRegistry::getAllWhitelistedOperators()`, but the AVS state
/// (IndexRegistry/StakeRegistry, surfaced via `OperatorStateRetriever`) tracks
/// *registered* operators independently of whitelist membership. When an
/// operator is removed from the whitelist yet stays REGISTERED with non-zero
/// stake, whitelist-sourced pruning evicts its pubkeys/socket/stake from the
/// off-chain cache while `AvsRegistryServiceChainCaller` still enumerates it —
/// any resulting cache miss is mapped to `Err`, aborting aggregation / BN254
/// certificate witness building, and the gateway's quorum math counts its stake
/// in the denominator (`total_stake`) but never the numerator (`signed_stake`),
/// producing `QuorumNotReached` timeouts.
///
/// Sourcing membership from the registered set — the same view the AVS-state
/// caller uses — removes the mismatch by construction. `OperatorStateRetriever`
/// returns each operator's `operator` (address), `operatorId`, and `stake`
/// inline at a **single pinned block**, so this call yields membership,
/// `operatorId`, and per-quorum stake atomically — no `operatorId → address`
/// reverse lookup and no separate latest-block stake reads that would skew
/// membership-at-N against stakes-at-N+k.
///
/// The block-pinned read mirrors the mature cert path
/// (`AvsRegistryServiceChainCaller`), including the `len == quorum_count` guard.
///
/// Returns `Ok(None)` for a **no-signal tick**: the chain was reached successfully
/// but the result is not a trustworthy complete snapshot — either a
/// successful-but-short retriever result (fewer quorum vectors than `quorumCount`,
/// the invariant the cert path guards) or a `quorumCount == 0` read after a
/// non-zero count was previously observed. The caller leaves the cache untouched
/// and skips the prune rather than evicting healthy operators against a partial
/// view. This is distinct from a genuine transient RPC `Err` (which still
/// propagates); see NEWT-1715 re-review — returning `Err` for these no-signal
/// cases would exit the `start_service` select loop and tear the whole
/// operator-set service down rather than skipping the tick.
async fn query_registered_operator_snapshot(
    operator_registry_address: Address,
    operator_state_retriever_address: Address,
    http_rpc_url: String,
    observed_quorums: &std::sync::atomic::AtomicBool,
) -> eyre::Result<Option<RegisteredSnapshot>> {
    let provider = get_provider_with_retry(&http_rpc_url)?;
    let operator_registry = OperatorRegistry::new(operator_registry_address, provider.clone());

    // Pin the block FIRST so the whole snapshot — quorumCount, membership, and stake —
    // is genuinely read at the same block N. Reading quorumCount at latest and then
    // pinning afterward could disagree around quorum creation / head skew.
    let block_number = provider
        .get_block_number()
        .await
        .map_err(|e| eyre::eyre!("failed to read current block number: {e}"))?;
    let block_id = alloy::eips::BlockId::from(block_number);

    // `OperatorRegistry` is the slashing registry coordinator, so it doubles as the
    // `registry_coordinator` the AVS reader expects. `OperatorStateRetriever` is a
    // separate singleton whose address is not discoverable from the registry — it is
    // threaded in from deployment config by the caller.
    let quorum_count = operator_registry
        .quorumCount()
        .block(block_id)
        .call()
        .await
        .inspect_err(|e| error!("Failed to get quorum count: {e}"))?;

    if quorum_count == 0 {
        // A zero count is authoritative only during bootstrap. Once we have ever seen a
        // non-zero count, a later zero read is almost certainly a bad/transient response —
        // treating it as an empty snapshot would prune the entire cache. Return a
        // no-signal tick (`Ok(None)`) so the caller leaves the cache untouched. NOT `Err`:
        // a refresh `Err` propagates out of the `start_service` select loop and takes the
        // whole operator-set service down (NEWT-1715 re-review), which is not what a
        // transient zero-read should do.
        if observed_quorums.load(Ordering::Relaxed) {
            warn!(
                block_number,
                "quorumCount read as 0 after previously observing a non-zero count; treating as no-signal (cache untouched)"
            );
            return Ok(None);
        }
        debug!("[query_registered_operator_snapshot] quorum count is 0 (bootstrap); no registered operators");
        return Ok(Some(RegisteredSnapshot {
            quorum_count,
            operators: Vec::new(),
            total_stakes: HashMap::new(),
            block_number,
        }));
    }
    observed_quorums.store(true, Ordering::Relaxed);

    let reader = AvsRegistryChainReader::new(
        operator_registry_address,
        operator_state_retriever_address,
        http_rpc_url.clone(),
    )
    .await
    .map_err(|e| eyre::eyre!("failed to build AVS registry reader: {e}"))?;

    let quorum_numbers: Vec<u8> = (0..quorum_count).collect();
    let operators_by_quorum = reader
        .get_operators_stake_in_quorums_at_block(block_number, Bytes::from(quorum_numbers))
        .await
        .map_err(|e| eyre::eyre!("failed to enumerate registered operators: {e}"))?;

    let Some(snapshot) = assemble_registered_snapshot(quorum_count, &operators_by_quorum, block_number) else {
        // Short retriever result — a no-signal tick, not a smaller registered set.
        // `Ok(None)` so the caller skips the prune; see the fatal-channel note above.
        warn!(
            block_number,
            expected_quorums = quorum_count,
            got_quorums = operators_by_quorum.len(),
            "registered-operator snapshot is short; treating as no-signal (cache untouched)"
        );
        return Ok(None);
    };
    debug!(
        "[query_registered_operator_snapshot] {} registered operators across {} quorums at block {}",
        snapshot.operators.len(),
        quorum_count,
        block_number
    );
    Ok(Some(snapshot))
}

/// Build a [`RegisteredSnapshot`] from the retriever's per-quorum operator vectors.
///
/// Split out from the RPC so the short-snapshot guard and the dedup/stake-accumulation
/// logic are unit-testable without a chain. `operators_by_quorum[i]` is the operator
/// vector for quorum `i` (`0..quorum_count`).
///
/// Returns `None` for a **short** retriever result (fewer quorum vectors than
/// `quorum_count`) — a no-signal tick the caller must skip the prune on, distinct
/// from a valid empty-but-complete snapshot which is `Some` with empty operators.
fn assemble_registered_snapshot(
    quorum_count: u8,
    operators_by_quorum: &[Vec<OperatorStateRetriever::Operator>],
    block_number: u64,
) -> Option<RegisteredSnapshot> {
    // Guard the exact invariant the cert path guards (`avs/mod.rs` →
    // `InvalidQuorumNums`): the retriever must return one operator vector per
    // requested quorum. A successful-but-short result would otherwise flow through
    // as a complete snapshot and prune healthy operators against the missing
    // quorums. Return `None` so the caller records "no signal" and skips the prune
    // this tick — NOT `Err`, which would exit the `start_service` select loop and
    // tear the service down (NEWT-1715 re-review).
    if operators_by_quorum.len() != quorum_count as usize {
        return None;
    }

    // Union across quorums — an operator registered for multiple quorums appears
    // once per quorum, but the off-chain cache is keyed by address. Carry each
    // operator's inline `operatorId` and accumulate its per-quorum stake, and sum
    // per-quorum totals as we go so all three views come from this one pinned read.
    //
    // Build the `operators` Vec directly (preserving first-seen order) and use
    // `index_by_address` only to dedup and locate the existing entry — so there is
    // no separate drain step and no infallible-but-`expect`-able remove.
    let mut operators: Vec<RegisteredOperator> = Vec::new();
    let mut index_by_address: HashMap<Address, usize> = HashMap::new();
    let mut total_stakes: HashMap<u8, u128> = HashMap::new();
    for (quorum_index, quorum_operators) in operators_by_quorum.iter().enumerate() {
        let quorum_number = quorum_index as u8;
        let mut quorum_total: u128 = 0;
        for op in quorum_operators {
            let stake = op.stake.to::<u128>();
            quorum_total = quorum_total.saturating_add(stake);
            let index = *index_by_address.entry(op.operator).or_insert_with(|| {
                operators.push(RegisteredOperator {
                    address: op.operator,
                    operator_id: op.operatorId,
                    stakes: HashMap::new(),
                });
                operators.len() - 1
            });
            operators[index].stakes.insert(quorum_number, stake);
        }
        total_stakes.insert(quorum_number, quorum_total);
    }

    Some(RegisteredSnapshot {
        quorum_count,
        operators,
        total_stakes,
        block_number,
    })
}

/// Source of the operator info.
#[derive(strum_macros::Display, Debug, Clone, PartialEq)]
#[strum(serialize_all = "lowercase")]
pub enum StateSource {
    /// Historic state source.
    Historic,
    /// Event state source.
    Event,
}

/// A stake value stamped with the chain block it was observed at, used for
/// last-writer-wins reconciliation between the periodic block-pinned snapshot and
/// live `OperatorStakeUpdate` events.
///
/// NEWT-1715: a sticky per-entry `StateSource` label (Event-always-beats-Historic)
/// re-opened the zero-clear staleness — once any single-quorum event flipped an
/// operator to `Event`, every later Historic snapshot was refused and a quorum the
/// operator dropped out of never cleared. Stamping each per-quorum value with its
/// source block and letting the newer block win fixes this by construction: a later
/// snapshot supersedes an older event, and a later event supersedes an older
/// snapshot, per quorum.
pub type VersionedStake = (u128, u64);

/// Map of operator IDs to their per-quorum stake, each stamped with its source block.
pub type OperatorStakeMap = HashMap<OperatorId, HashMap<u8, VersionedStake>>;

/// State of the operator info service.
#[derive(Debug, Clone)]
pub struct OperatorStates {
    /// Operator info data.
    pub operator_info_data: Arc<RwLock<HashMap<Address, (StateSource, OperatorPubKeys)>>>,
    /// Operator address to ID mapping.
    pub operator_addr_to_id: Arc<RwLock<HashMap<Address, (StateSource, OperatorId)>>>,
    /// Socket dictionary.
    pub socket_dict: Arc<RwLock<HashMap<OperatorId, (StateSource, String)>>>,
    /// Operator stake amount, per-quorum, block-versioned (see [`VersionedStake`]).
    pub operator_stake: Arc<RwLock<OperatorStakeMap>>,
    /// Total stake amount, per-quorum, block-versioned (see [`VersionedStake`]).
    pub total_stake: Arc<RwLock<HashMap<u8, VersionedStake>>>,
}

/// OperatorSocket
#[derive(Debug)]
pub struct OperatorSocket {
    /// Operator ID.
    pub id: OperatorId,
    /// Socket address.
    pub socket: String,
}

/// Signal published when `OperatorRegistryEpochGovernance` advances its operator-set epoch.
///
/// Fired in two situations: once at `initializeEpochs(uint32)` when the governance contract
/// leaves the bootstrap phase (`epoch == 0`), and again every time `applyPendingChanges()`
/// crosses an epoch boundary on a chain where any queued whitelist / task-generator /
/// deregistration or `epochDurationBlocks` mutation took effect. Consumers should treat the
/// signal as a non-authoritative hint to refresh cached operator-set views; the contract
/// remains the source of truth, so a missed signal is not a correctness failure (the
/// periodic refresh already in `OperatorRegistryService::start_service` reconverges within
/// one epoch).
#[derive(Debug, Clone, Copy)]
pub struct EpochAdvancedSignal {
    /// Source-chain id this registry lives on (multichain disambiguation for fan-in consumers).
    pub chain_id: u64,
    /// Epoch number now active on the governance contract. `0` indicates the
    /// post-`initializeEpochs` boundary.
    pub epoch: u32,
    /// Block number on the registry's chain at which the new epoch started.
    pub start_block: u64,
    /// Epoch length in blocks for the newly-active epoch (post-`epochDurationBlocks` change, if any).
    pub duration_blocks: u32,
}

/// Operator Registry info message
#[derive(Debug)]
pub(crate) enum OperatorsInfoMessage {
    InsertOperatorInfo(
        Option<Address>,
        Option<Box<OperatorPubKeys>>,
        Option<OperatorSocket>,
        StateSource,
    ),
    /// Apply a single-quorum operator stake observed at the given block
    /// (`operator_id`, `quorum`, `stake`, `block`). Block-versioned last-writer-wins:
    /// applied only if `block` is at least as new as the cached value for that quorum.
    /// Used by the per-quorum `OperatorStakeUpdate` event path.
    InsertOperatorStake(OperatorId, u8, u128, u64),
    /// Apply a single-quorum total stake observed at the given block
    /// (`quorum`, `stake`, `block`), block-versioned last-writer-wins.
    InsertTotalStake(u8, u128, u64),
    /// Reconcile an operator's entire per-quorum stake map against a block-pinned
    /// snapshot taken at `block` (`operator_id`, `stakes`, `block`). Per quorum the
    /// newer block wins; a quorum absent from `stakes` is cleared unless the cache
    /// holds a strictly-newer event for it — so a stake that dropped to zero no longer
    /// lingers as a stale weight feeding `signed_stake` (NEWT-1715), without clobbering
    /// a live event that arrived after the snapshot block.
    ReplaceOperatorStake(OperatorId, HashMap<u8, u128>, u64),
    /// Reconcile the entire per-quorum total-stake view against a block-pinned
    /// snapshot taken at `block` (`total_stakes`, `block`), same block-versioned rule.
    ReplaceTotalStakes(HashMap<u8, u128>, u64),
    #[allow(dead_code)]
    Remove(Address),
    GetPubKeys(Address, Sender<Option<OperatorPubKeys>>),
    GetSockets(Address, Sender<Option<String>>),
    GetOperatorStake(OperatorId, u8, Sender<Option<u128>>),
    GetTotalStake(u8, Sender<Option<u128>>),
}

/// Apply a single-quorum stake observed at `block` to a block-versioned per-quorum map,
/// keeping whichever observation is newer.
///
/// Last-writer-wins with `block` as the version: a strictly-older observation is
/// dropped, and a same-or-newer one wins. Ties (same block) take the incoming value —
/// harmless because two reads of the same block observe the same on-chain stake.
pub(crate) fn apply_versioned_stake(map: &mut HashMap<u8, VersionedStake>, quorum: u8, stake: u128, block: u64) {
    match map.get(&quorum) {
        Some((_, existing_block)) if *existing_block > block => {}
        _ => {
            map.insert(quorum, (stake, block));
        }
    }
}

/// Reconcile a whole per-quorum map against a block-pinned snapshot taken at
/// `snapshot_block`, in place.
///
/// Per quorum the newer block wins:
/// * a quorum in `snapshot` overwrites the cached value unless the cache holds a
///   strictly-newer entry (a live event that arrived after the snapshot block);
/// * a quorum absent from `snapshot` is cleared unless the cache holds a
///   strictly-newer entry — so a stake that dropped to zero is removed, while a
///   post-snapshot event is preserved.
///
/// This is the LWW generalization of the old "replace, but keep Event" logic that
/// fixes NEWT-1715: it no longer lets a stale event immunize a quorum against every
/// future snapshot, because the *block*, not the source label, decides.
pub(crate) fn reconcile_snapshot(
    map: &mut HashMap<u8, VersionedStake>,
    snapshot: HashMap<u8, u128>,
    snapshot_block: u64,
) {
    // Drop cached quorums that are absent from the snapshot and not newer than it.
    map.retain(|quorum, (_, block)| *block > snapshot_block || snapshot.contains_key(quorum));
    // Apply the snapshot's quorums under the same newer-wins rule.
    for (quorum, stake) in snapshot {
        apply_versioned_stake(map, quorum, stake, snapshot_block);
    }
}

/// Query registered operator and fill database.
///
/// `operator_states` is read pre-insert to compute the diff between the
/// previously-cached operator set and the freshly-fetched registered set. Any
/// address present in the cache but absent from the new registered set gets a
/// `Remove(addr)` message — without this prune step, deregistered operators
/// would persist in `operator_addr_to_id` indefinitely (the chain-event
/// `Remove` path is only fan-out from log subscriptions, which can be missed
/// during socket reconnects). Downstream readers (gateway routing pool via
/// `get_operator_id_to_addr_mappings`, BLS-cert path via `get_operator_info`)
/// would then serve ghost entries until the next process restart.
///
/// NEWT-1715: membership is sourced from the on-chain **registered** set
/// (`OperatorStateRetriever`) rather than the whitelist, so an operator removed
/// from the whitelist but still REGISTERED with stake is retained — see
/// [`query_registered_operator_snapshot`].
pub(crate) async fn query_registered_operator_and_fill_db(
    operator_registry_address: Address,
    operator_state_retriever_address: Address,
    http_rpc_url: String,
    pub_keys: UnboundedSender<OperatorsInfoMessage>,
    operator_states: &OperatorStates,
    observed_quorums: &std::sync::atomic::AtomicBool,
) -> eyre::Result<()> {
    // Snapshot membership, per-operator stake, and per-quorum totals in one
    // block-pinned read. A transient RPC failure aborts the whole refresh here
    // (`?`) so the cache is left untouched rather than presenting an empty/partial
    // set that the prune step would mistake for "everyone deregistered". Because
    // stakes and totals ride along with this read, they are consistent with
    // membership at the same block — no separate latest-block stake reads that
    // could skew against membership.
    //
    // `Ok(None)` is a deliberate no-signal tick (short retriever result, or a
    // `quorumCount == 0` read after a non-zero count was observed): return early
    // with `Ok(())` so the cache is untouched and the prune is skipped. Crucially
    // this is NOT `Err` — a refresh `Err` propagates out of the `start_service`
    // select loop and tears the whole operator-set service down (NEWT-1715
    // re-review). Making the loop survive genuine transient errors is tracked in
    // the resilience follow-up (NEWT-1763).
    let Some(snapshot) = query_registered_operator_snapshot(
        operator_registry_address,
        operator_state_retriever_address,
        http_rpc_url.clone(),
        observed_quorums,
    )
    .await?
    else {
        return Ok(());
    };

    let snapshot_block = snapshot.block_number;
    let registered_addresses: Vec<Address> = snapshot.operators.iter().map(|o| o.address).collect();
    let registered_ids: Vec<(Address, OperatorId)> =
        snapshot.operators.iter().map(|o| (o.address, o.operator_id)).collect();
    // Per-operator stake keyed by address; carried straight from the pinned
    // snapshot so the receiver can reconcile the operator's whole stake map by block
    // (last-writer-wins), clearing quorums the operator no longer holds stake in
    // unless a strictly-newer event superseded the snapshot.
    let stakes_by_addr: HashMap<Address, HashMap<u8, u128>> = snapshot
        .operators
        .iter()
        .map(|o| (o.address, o.stakes.clone()))
        .collect();
    let total_stakes = snapshot.total_stakes;

    let pubkey_handle =
        query_registered_operator_pub_keys(operator_registry_address, http_rpc_url.clone(), registered_addresses);
    let socket_handle =
        query_registered_operator_sockets(operator_registry_address, http_rpc_url.clone(), registered_ids);

    let mut operator_sockets: HashMap<OperatorId, String> = HashMap::new();
    let mut operator_pub_keys: HashMap<Address, OperatorPubKeys> = HashMap::new();

    let (pub_keys_res, sockets_res) = futures::join!(pubkey_handle, socket_handle);

    // Critical: gate the prune step on a complete pubkey + socket snapshot. If
    // either RPC fails, the corresponding map stays empty — and a naive diff
    // against the cached set would mark every operator as stale, emitting
    // `Remove(addr)` for all of them and emptying the cache on any transient RPC
    // hiccup. The pre-NEWT-1175-followup behavior was "failed fetch ⇒ no updates,
    // cache stays as-is"; preserve that for the prune side too.
    //
    // Pruning removes stale operators *before* the insert loop below re-adds the
    // current registered set, and that loop hard-errors if any current operator
    // is missing a socket. So a partial fetch of either would publish `Remove`
    // messages and then bail without inserting replacements — an empty-pool
    // outcome. Stakes/totals no longer need a separate gate: they come from the
    // block-pinned snapshot already validated above (`?`), so if we reach here
    // they are complete and consistent with membership.
    //
    // Each query already distinguishes a transient failure (returns `Err`, drops
    // it here) from a permanent per-operator revert (drops just that operator
    // and still returns `Ok`), so this gate trips only on transient errors.
    let snapshot_complete = pub_keys_res.is_ok() && sockets_res.is_ok();
    if let Ok(pub_keys) = pub_keys_res {
        operator_pub_keys = pub_keys;
    }
    if let Ok(sockets) = sockets_res {
        operator_sockets = sockets;
    }

    debug!(
        "[query_registered_operator_and_fill_db] queried registered operator info. count of operators: {}",
        operator_pub_keys.len()
    );

    // Diff only when we actually observed the registered set. An empty registered
    // snapshot (validated above) is a valid signal that the set is empty (e.g.,
    // every operator deregistered) and we SHOULD prune all cached entries, but a
    // failed pubkey/socket fetch is a no-signal — leave the cache alone.
    if snapshot_complete {
        let cached_addrs: std::collections::HashSet<Address> = {
            let map = operator_states.operator_addr_to_id.read().await;
            map.keys().copied().collect()
        };
        let current_addrs: std::collections::HashSet<Address> = operator_pub_keys.keys().copied().collect();
        let stale: Vec<&Address> = cached_addrs.difference(&current_addrs).collect();
        if !stale.is_empty() {
            debug!(
                cached = cached_addrs.len(),
                current = current_addrs.len(),
                pruning = stale.len(),
                "[query_registered_operator_and_fill_db] pruning stale operators from cache"
            );
        }
        for stale_addr in stale {
            debug!(
                operator = %stale_addr,
                "[query_registered_operator_and_fill_db] pruning stale operator"
            );
            let _ = pub_keys.send(OperatorsInfoMessage::Remove(*stale_addr));
        }
    } else {
        warn!("[query_registered_operator_and_fill_db] incomplete operator snapshot (pubkey or socket RPC failed); skipping prune step (cache unchanged this tick)");
        return Ok(());
    }

    for (address, operator_pub_key) in operator_pub_keys.iter() {
        let operator_id = operator_id_from_g1_pub_key(operator_pub_key.g1_pub_key.clone())?;

        if let Some(socket) = operator_sockets.get(&operator_id) {
            let message = OperatorsInfoMessage::InsertOperatorInfo(
                Some(*address),
                Some(Box::new(operator_pub_key.clone())),
                Some(OperatorSocket {
                    id: operator_id,
                    socket: socket.clone(),
                }),
                StateSource::Historic,
            );
            debug!("Inserting operator info for operator {address} from historic state: operator id {operator_id}");
            let _ = pub_keys.send(message);
        } else {
            return Err(
                alloy::contract::Error::TransportError(alloy::transports::TransportError::ErrorResp(
                    alloy::rpc::json_rpc::ErrorPayload {
                        code: -1,
                        message: "Socket not found".into(),
                        data: None,
                    },
                ))
                .into(),
            );
        }

        // Reconcile the operator's entire per-quorum stake map against the block-pinned
        // snapshot. `ReplaceOperatorStake` clears any quorum absent from `stakes` — so an
        // operator whose stake in a quorum dropped to zero no longer leaves a stale
        // non-zero weight feeding `signed_stake` (NEWT-1715 / Octane denominator-numerator
        // mismatch) — unless the cache holds a strictly-newer event for that quorum, which
        // the block-versioned reconcile preserves. An operator with no stake anywhere is
        // sent an empty map, clearing its stale cached weights.
        let stakes = stakes_by_addr.get(address).cloned().unwrap_or_default();
        let _ = pub_keys.send(OperatorsInfoMessage::ReplaceOperatorStake(
            operator_id,
            stakes,
            snapshot_block,
        ));
    }

    // Reconcile the whole per-quorum total-stake view against the same pinned snapshot,
    // clearing any quorum absent from `total_stakes` (block-versioned), for the same reason.
    let _ = pub_keys.send(OperatorsInfoMessage::ReplaceTotalStakes(total_stakes, snapshot_block));

    Ok(())
}

/// Queries existing operators from for a particular block range.
/// # Arguments
///
/// * `operator_registry_address` - The address of the operator registry contract.
/// * `http_rpc_url` - The HTTP RPC url to use for querying.
///
/// # Returns
///
/// * (`HashMap<Address, OperatorPubKeys>`) - A hashmap of operator addresses and its
///   corresponding operator pub keys.
#[instrument(skip_all)]
pub async fn query_registered_operator_pub_keys(
    operator_registry_address: Address,
    http_rpc_url: String,
    registered_operators: Vec<Address>,
) -> eyre::Result<(HashMap<Address, OperatorPubKeys>)> {
    debug!("[query_registered_operator_pub_keys] operator_registry_address {operator_registry_address}");
    let provider = get_provider_with_retry(&http_rpc_url)?;
    let operator_registry = OperatorRegistry::new(operator_registry_address, provider.clone());

    let bls_apk_registry_address = operator_registry
        .blsApkRegistry()
        .call()
        .await
        .inspect_err(|e| error!("Failed to get BLS APK registry address: {e}"))?;

    let bls_apk_registry = BLSApkRegistry::new(bls_apk_registry_address, provider.clone());

    debug!(
        "[query_registered_operator_pub_keys] registered_operators count: {}",
        registered_operators.len()
    );

    // This map is the authoritative "did we observe the registered set this tick"
    // signal that `query_registered_operator_and_fill_db` uses to gate cache
    // pruning. A per-operator failure is handled by `skip_or_abort_operator`:
    // a permanent contract revert (e.g. registered but BLS key not yet set)
    // drops just that operator and the snapshot stays valid, while a transient
    // error aborts the whole tick (`Err`) so the caller skips pruning rather
    // than mistaking a partial fetch for a smaller registered set and collapsing
    // the pool toward 0. Each future yields `Some` to keep, `None` to skip.
    futures::stream::iter(registered_operators)
        .map(|operator| {
            let bls_apk_registry = bls_apk_registry.clone();
            async move {
                let fetch = async {
                    let operator_id = bls_apk_registry.getOperatorId(operator).call().await?;
                    let g1 = bls_apk_registry.getRegisteredPubkey(operator).call().await?._0;
                    let g2 = bls_apk_registry.getOperatorPubkeyG2(operator).call().await?;
                    Ok::<_, eyre::Error>((operator_id, g1, g2))
                };
                match fetch.await {
                    Ok((operator_id, g1, g2)) => {
                        let operator_pub_key = OperatorPubKeys {
                            g1_pub_key: BlsG1Point::new(alloy_registry_g1_point_to_g1_affine(g1)),
                            g2_pub_key: BlsG2Point::new(alloy_registry_g2_point_to_g2_affine(g2)),
                        };
                        debug!(
                            "[query_registered_operator_pub_keys] operator {}, operator_id {}",
                            operator, operator_id
                        );
                        debug!("[query_registered_operator_pub_keys] operator_pub_key {operator_pub_key:?}");
                        Ok(Some((operator, operator_pub_key)))
                    }
                    Err(e) => {
                        skip_or_abort_operator(operator, "query_registered_operator_pub_keys", e).map(|_skipped| None)
                    }
                }
            }
        })
        .buffer_unordered(MAX_CONCURRENT_OPERATOR_QUERIES)
        .try_filter_map(|opt| async move { Ok(opt) })
        .try_collect()
        .await
        .inspect_err(|e| {
            error!("Transient failure fetching operator pubkeys; treating registered set as unobserved this tick: {e}")
        })
}

/// Queries operator sockets for the given registered operators.
///
/// # Arguments
///
/// * `operator_registry_address` - The address of the operator registry contract.
/// * `http_rpc_url` - The HTTP RPC url to use for querying.
/// * `registered_operators` - The registered `(address, operatorId)` pairs to fetch sockets for.
///   The `operatorId` is carried from the block-pinned snapshot, so no per-operator
///   `getOperatorId` lookup is needed here.
///
/// # Returns
///
/// * `HashMap<OperatorId, String>` - Operator Id to socket mapping containing all the operator
///   sockets registered in the given block range
#[instrument(skip_all)]
pub async fn query_registered_operator_sockets(
    operator_registry_address: Address,
    http_rpc_url: String,
    registered_operators: Vec<(Address, OperatorId)>,
) -> eyre::Result<HashMap<OperatorId, String>> {
    debug!("[query_registered_operator_sockets] operator_registry_address {operator_registry_address}");
    let provider = get_provider_with_retry(&http_rpc_url)?;
    let operator_registry = OperatorRegistry::new(operator_registry_address, provider.clone());

    let socket_registry_address = operator_registry
        .socketRegistry()
        .call()
        .await
        .inspect_err(|e| error!("Failed to get socket_registry_address: {e}"))?;
    debug!("query_registered_operator_sockets socket_registry_address {socket_registry_address}");
    let socket_registry = SocketRegistry::new(socket_registry_address, provider);

    debug!(
        "[query_registered_operator_sockets] registered_operators count: {}",
        registered_operators.len()
    );

    // The socket map is the other half of the prune snapshot in
    // `query_registered_operator_and_fill_db`: removals are published before the
    // insert loop runs, so a partial socket fetch that silently dropped failures
    // would let stale operators get pruned while their replacements (missing a
    // socket) never get inserted — the same empty-pool outcome the pubkey path
    // guards against, via a different trigger. As with pubkeys, a permanent
    // contract revert drops just that operator while a transient error aborts
    // the tick (`Err`) so the caller can gate the prune on a complete snapshot.
    futures::stream::iter(registered_operators)
        .map(|(operator, operator_id)| {
            let socket_registry = socket_registry.clone();
            async move {
                match socket_registry.getOperatorSocket(operator_id).call().await {
                    Ok(socket) => {
                        debug!(
                            "[query_registered_operator_sockets] operator {}, operator_id {}, socket {}",
                            operator, operator_id, socket
                        );
                        Ok(Some((operator_id, socket)))
                    }
                    Err(e) => skip_or_abort_operator(operator, "query_registered_operator_sockets", e.into())
                        .map(|_skipped| None),
                }
            }
        })
        .buffer_unordered(MAX_CONCURRENT_OPERATOR_QUERIES)
        .try_filter_map(|opt| async move { Ok(opt) })
        .try_collect()
        .await
        .inspect_err(|e| {
            error!("Transient failure fetching operator sockets; treating snapshot as incomplete this tick: {e}")
        })
}

// NEWT-1715: per-operator stake (`query_registered_operator_stakes`) and per-quorum
// total-stake (`query_total_stakes`) fan-outs were removed. Both are now sourced from
// the single block-pinned `query_registered_operator_snapshot` read: the retriever
// returns each operator's inline per-quorum stake atomically with membership, and the
// per-quorum totals are summed from those same values. This eliminates the redundant
// latest-block RPCs and the membership-at-N / stakes-at-N+k consistency skew.

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

    #[test]
    fn contract_revert_skips_operator() {
        // A registered operator that has not yet set its BLS pubkey reverts on
        // getRegisteredPubkey; re-running always reverts the same way, so we drop
        // just that operator and keep the rest of the snapshot (the original
        // filter_map behavior).
        let err = eyre::eyre!("execution reverted: custom error 0x25ec6c1f, data: \"0x25ec6c1f\"");
        let skipped = skip_or_abort_operator(Address::ZERO, "test", err).expect("revert should skip, not abort");
        assert!(skipped);
    }

    #[test]
    fn transient_error_aborts_tick() {
        // A transient error means we did not observe this operator; aborting the
        // whole tick keeps the caller from mistaking a partial fetch for a
        // smaller registered set and pruning healthy operators.
        for msg in ["connection refused", "request timeout", "429 Too Many Requests"] {
            let res = skip_or_abort_operator(Address::ZERO, "test", eyre::eyre!("{msg}"));
            assert!(res.is_err(), "transient error {msg:?} should abort the tick");
        }
    }

    use alloy::primitives::{aliases::U96, FixedBytes};

    /// Build a retriever `Operator` for tests.
    fn op(addr: u8, id: u8, stake: u128) -> OperatorStateRetriever::Operator {
        OperatorStateRetriever::Operator {
            operator: Address::repeat_byte(addr),
            operatorId: FixedBytes::<32>::repeat_byte(id),
            stake: U96::from(stake),
        }
    }

    #[test]
    fn snapshot_retains_registered_operator() {
        // NEWT-1715 core: an operator present in the registered set (regardless of
        // whitelist membership) with non-zero stake must survive into the snapshot,
        // carrying its address, id, and per-quorum stake.
        let dewhitelisted = op(0xAA, 0x01, 500);
        let snapshot = assemble_registered_snapshot(1, &[vec![dewhitelisted]], 100).expect("complete snapshot");

        assert_eq!(snapshot.operators.len(), 1);
        let kept = &snapshot.operators[0];
        assert_eq!(kept.address, Address::repeat_byte(0xAA));
        assert_eq!(kept.operator_id, FixedBytes::<32>::repeat_byte(0x01));
        assert_eq!(kept.stakes.get(&0), Some(&500));
        assert_eq!(snapshot.total_stakes.get(&0), Some(&500));
    }

    #[test]
    fn snapshot_dedups_operator_across_quorums() {
        // An operator registered for multiple quorums appears once per quorum in the
        // retriever result, but the cache is address-keyed — it must dedup to a single
        // entry whose per-quorum stakes and the per-quorum totals are both preserved.
        let q0 = vec![op(0xAA, 0x01, 300), op(0xBB, 0x02, 700)];
        let q1 = vec![op(0xAA, 0x01, 400)];
        let snapshot = assemble_registered_snapshot(2, &[q0, q1], 100).expect("complete snapshot");

        assert_eq!(snapshot.operators.len(), 2, "0xAA must not be double-counted");
        let a = snapshot
            .operators
            .iter()
            .find(|o| o.address == Address::repeat_byte(0xAA))
            .expect("0xAA present");
        assert_eq!(a.stakes.get(&0), Some(&300));
        assert_eq!(a.stakes.get(&1), Some(&400));
        // Per-quorum totals are summed independently across the two quorums.
        assert_eq!(snapshot.total_stakes.get(&0), Some(&1000));
        assert_eq!(snapshot.total_stakes.get(&1), Some(&400));
    }

    #[test]
    fn short_snapshot_is_rejected() {
        // C1: a successful-but-short retriever result (fewer quorum vectors than
        // quorumCount) must be a no-signal tick (`None`) so the caller skips the
        // prune, rather than pruning healthy operators against the missing quorums.
        // `None` — not `Err` — because the caller maps this to `Ok(())` (skip tick)
        // instead of a fatal refresh error that would tear down the service.
        let only_one_quorum = vec![vec![op(0xAA, 0x01, 500)]];
        assert!(
            assemble_registered_snapshot(3, &only_one_quorum, 100).is_none(),
            "short snapshot must be a no-signal None"
        );
    }

    #[test]
    fn empty_registered_set_is_a_valid_signal() {
        // An `Ok` snapshot with the right number of (empty) quorum vectors is a valid
        // "everyone deregistered" signal — distinct from the short-snapshot Err above.
        let snapshot = assemble_registered_snapshot(2, &[vec![], vec![]], 100).expect("empty but complete");
        assert!(snapshot.operators.is_empty());
        assert_eq!(snapshot.total_stakes.get(&0), Some(&0));
        assert_eq!(snapshot.total_stakes.get(&1), Some(&0));
    }

    // ---- Block-versioned last-writer-wins reconciliation (NEWT-1715 re-review) ----

    #[test]
    fn apply_versioned_stake_newer_block_wins() {
        let mut map: HashMap<u8, VersionedStake> = HashMap::new();
        apply_versioned_stake(&mut map, 0, 100, 10);
        assert_eq!(map.get(&0), Some(&(100, 10)));
        // Newer block overwrites.
        apply_versioned_stake(&mut map, 0, 200, 20);
        assert_eq!(map.get(&0), Some(&(200, 20)));
        // Older block is ignored.
        apply_versioned_stake(&mut map, 0, 999, 15);
        assert_eq!(map.get(&0), Some(&(200, 20)));
        // Equal block takes the incoming value (same-block reads agree on-chain).
        apply_versioned_stake(&mut map, 0, 250, 20);
        assert_eq!(map.get(&0), Some(&(250, 20)));
    }

    #[test]
    fn reconcile_snapshot_clears_dropped_quorum() {
        // Regression for the sticky-source-label bug: an operator that received a stake
        // event (any quorum) must still have a quorum it dropped out of cleared by a
        // later snapshot. With block-versioning this Just Works — the newer snapshot
        // block supersedes the older event.
        let mut map: HashMap<u8, VersionedStake> = HashMap::new();
        // An event landed for quorum 0 and quorum 1 at block 10.
        apply_versioned_stake(&mut map, 0, 100, 10);
        apply_versioned_stake(&mut map, 1, 50, 10);
        // A later snapshot at block 20 shows the operator only in quorum 0.
        reconcile_snapshot(&mut map, HashMap::from([(0u8, 120u128)]), 20);
        assert_eq!(map.get(&0), Some(&(120, 20)), "quorum 0 updated to snapshot value");
        assert_eq!(
            map.get(&1),
            None,
            "quorum 1 (dropped) must be cleared despite the earlier event"
        );
    }

    #[test]
    fn reconcile_snapshot_preserves_newer_event() {
        // A live event that arrived AFTER the snapshot block must survive reconciliation
        // (the snapshot is stale for that quorum).
        let mut map: HashMap<u8, VersionedStake> = HashMap::new();
        // Event at block 30 for quorum 0.
        apply_versioned_stake(&mut map, 0, 500, 30);
        // Older snapshot (block 20) that omits quorum 0.
        reconcile_snapshot(&mut map, HashMap::new(), 20);
        assert_eq!(
            map.get(&0),
            Some(&(500, 30)),
            "post-snapshot event must not be cleared by an older snapshot"
        );
    }

    #[test]
    fn reconcile_snapshot_older_event_loses_to_snapshot() {
        // An event older than the snapshot must be superseded by the snapshot value.
        let mut map: HashMap<u8, VersionedStake> = HashMap::new();
        apply_versioned_stake(&mut map, 0, 500, 10); // stale event
        reconcile_snapshot(&mut map, HashMap::from([(0u8, 700u128)]), 25);
        assert_eq!(map.get(&0), Some(&(700, 25)), "newer snapshot supersedes older event");
    }
}