newton-aggregator 0.4.13

//! Aggregator crate

#[macro_use]
extern crate newton_core;

/// Config
pub mod config;
/// Median-based consensus logic for normalizing operator responses
pub mod consensus;
/// Core aggregator logic shared between implementations
pub mod core;
/// Aggregator error
pub mod error;
/// RPC server
pub mod rpc_server;
/// Shared simulation request/response types for operator delegation
pub mod simulation;
/// State-commit aggregation pipeline (Track 1 — NEWT-1115).
pub mod state_commit;
/// Local task supervisor (`spawn_monitored`) — avoids circular dep on gateway.
pub mod task_supervisor;
use alloy::{
    dyn_abi::SolType,
    providers::{Provider, ProviderBuilder, WsConnect},
    rpc::types::Filter,
    sol_types::{SolEvent, SolValue},
};
use ark_ec::AffineRepr;
pub use core::AggregatorCore;
use eigensdk::client_avsregistry::reader::AvsRegistryChainReader;
pub use error::{AggregatorCoreError, AggregatorError};
use futures_util::StreamExt;
use jsonrpsee::{
    server::{RpcModule, Server},
    types::ErrorObject,
};
pub use newton_chainio::bls::BlsAggregationServiceResponse;
use newton_chainio::bls::ServiceHandle;
use newton_core::{
    common::chain::get_block_duration_ms,
    config::{key::EcdsaKey, rpc::ChainRpcProviderConfig, NewtonAvsConfig},
    r#newton_prover_task_manager::{
        INewtonProverTaskManager::{Task, TaskResponse as BindingTaskResponse},
        NewtonProverTaskManager::NewTaskCreated,
    },
    TaskId,
};
pub use rpc_server::{
    ConsensusCommitRequest, ConsensusCommitResponse, ConsensusData, ConsensusPrepareRequest, ConsensusPrepareResponse,
    FieldAdjustment, OperatorErrorResponse, PartialDecryptionData, PreparePhaseChainContext, PreparePhaseFetchRequest,
    SignedTaskResponse, TaskResponse,
};
use std::{
    collections::HashMap,
    net::{SocketAddr, ToSocketAddrs},
    sync::Arc,
};
use tokio::{
    task::JoinHandle,
    time::{sleep, Duration},
};
use tokio_util::sync::CancellationToken;
use tracing::{error, info, instrument, warn, Instrument};

use crate::{config::AggregatorConfig, state_commit::StateCommitOrchestrator};
use newton_chainio::avs::writer::AvsWriter;

// Removed TASK_CHALLENGE_WINDOW_BLOCK constant - now fetched dynamically from contract

/// Aggregator
#[allow(missing_debug_implementations)]
pub struct Aggregator {
    /// Socket address
    port_address: String,
    /// AVS writer (Arc-shared with `AvsWriterCommitter` so all txs for the
    /// aggregator's ECDSA signer route through one submitter — see
    /// `.claude/rules/lessons.md` "single-signer routing")
    avs_writer: Arc<AvsWriter>,
    /// HashMap to store tasks
    tasks: HashMap<TaskId, Task>,
    /// Core aggregator logic shared with gateway aggregator
    core: Arc<AggregatorCore>,
    /// RPC provider configuration for all chains
    rpc: ChainRpcProviderConfig,
    /// Chain ID for RPC lookups (destination chain in multichain mode)
    chain_id: u64,
    /// Cached task response window block (fetched once at initialization)
    task_response_window_block: u64,
    /// Timeout for waiting for aggregation responses (in seconds)
    aggregation_timeout_secs: u64,
    /// Per-chain state-commit orchestrator (NEWT-1115). `Some` only when the
    /// `dev-stub` feature is enabled and a `StateCommitRegistry` address is
    /// configured; production builds will populate this in NEWT-1116 with the
    /// real `EnclavePcr0Provider`.
    state_commit_orchestrator: Option<Arc<StateCommitOrchestrator>>,
}

impl Aggregator {
    /// Creates a new aggregator
    ///
    /// # Arguments
    ///
    /// * [`NewtonAvsConfig<AggregatorConfig>`] - The configuration for the aggregator
    ///
    /// # Returns
    ///
    /// * `Self` - The aggregator
    ///
    /// # Errors
    ///
    /// * `AggregatorError` - The error that occurred
    pub async fn new(config: NewtonAvsConfig<AggregatorConfig>) -> Result<Self, AggregatorError> {
        let chain_id = config.chain_id;
        let source_chain_id = config.get_source_chain_id();

        // Get RPC config for the current chain (destination chain in multichain mode)
        let rpc_config = config.rpc.get_or_err(chain_id)?.clone();

        // Source chain: where EigenLayer AVS contracts (OperatorRegistry, etc.) are deployed
        let source_rpc = config.rpc.get_or_err(source_chain_id)?.clone();

        // operator_state_retriever is source-only and validated by `load_src` →
        // `validate_for_role(ChainRole::Source)`. Aggregator only runs on source
        // chains, so this should be Some here; the typed error is defensive.
        let operator_state_retriever = config.contracts.avs.operator_state_retriever.ok_or_else(|| {
            std::io::Error::other(
                "aggregator requires source-chain operator_state_retriever but it was None — load_src must validate this",
            )
        })?;

        let avs_registry_chain_reader = AvsRegistryChainReader::new(
            config.contracts.avs.operator_registry,
            operator_state_retriever,
            source_rpc.http.clone(),
        )
        .await?;

        let state_commit_registry_addr = config.contracts.avs.state_commit_registry_or_zero();

        let avs_writer = Arc::new(
            AvsWriter::new(
                config.contracts.avs.newton_prover_task_manager,
                config.contracts.avs.identity_registry,
                state_commit_registry_addr,
                source_rpc.http.clone(),
                config.service.signer.clone(),
                config.chain_id,
            )
            .await?,
        );

        // Load the aggregator's ECDSA signer once. We keep both the address
        // (for metrics) and the signer itself (used by the state-commit HTTP
        // client to wrap each outbound `newt_*StateCommit*` call in an
        // EIP-712 `Authenticated<T>` envelope — operator-side dispatch
        // gates on `isTaskGenerator(signer)`).
        let aggregator_signer = newton_core::keys::load_ecdsa(&config.service.signer)
            .map_err(|e| AggregatorError::KeyError(e.to_string()))?;
        let aggregator_address = aggregator_signer.address();
        newton_metric::set_aggregator_address(&aggregator_address.to_string());

        let dest_operator_table_updater = config
            .contracts
            .destination_multichain()
            .ok()
            .map(|dm| dm.operator_table_updater);
        let core = AggregatorCore::new(
            avs_registry_chain_reader,
            config.contracts.avs.operator_registry,
            operator_state_retriever,
            config.is_destination_chain(),
            config.rpc.clone(),
            source_chain_id,
            chain_id,
            config.contracts.avs.newton_prover_task_manager,
            dest_operator_table_updater,
        )
        .await
        .map_err(|e| AggregatorError::KeyError(e.to_string()))?;

        // Fetch task response window block once at initialization (cached for reuse)
        let task_response_window_block =
            newton_core::common::chain::get_task_response_window_blocks(&rpc_config.http)
                .await
                .map_err(|e| AggregatorError::KeyError(format!("Failed to get task response window: {}", e)))?
                as u64;

        // Construct the per-chain orchestrator under `dev-stub`. Production
        // wiring lands in NEWT-1116: `EnclavePcr0Provider` replaces
        // `StubPcr0Provider`, the operator set is hydrated from
        // `AvsRegistryServiceCaller`, and the gate is lifted. We share
        // `Arc<AvsWriter>` with `process_aggregator_responses` so all txs
        // for the aggregator's ECDSA signer route through one submitter
        // (single-signer routing per `lessons.md`).
        #[cfg(any(test, feature = "dev-stub"))]
        let state_commit_orchestrator: Option<Arc<StateCommitOrchestrator>> = {
            use std::{collections::HashMap as StdHashMap, time::Duration};

            use alloy::providers::ProviderBuilder;
            use newton_core::state_commit_registry::StateCommitRegistry::StateCommitRegistryInstance;

            use crate::state_commit::{
                AvsWriterCommitter, HttpStateCommitOperatorClient, OnchainRegistryReader, StateCommitAggregator,
                StubOperatorSetSnapshotReader, StubPcr0Provider,
            };

            let http_provider = ProviderBuilder::new().connect_http(
                source_rpc
                    .http
                    .parse()
                    .map_err(|e| AggregatorError::KeyError(format!("invalid source RPC URL: {e}")))?,
            );
            let registry_instance = StateCommitRegistryInstance::new(state_commit_registry_addr, http_provider);
            let reader = Arc::new(OnchainRegistryReader::new(registry_instance, chain_id));
            let pcr0 = Arc::new(StubPcr0Provider::new());
            // Production wiring (NEWT-1116) will hydrate sockets from
            // `AvsRegistryServiceCaller`. Until then, the dev-stub uses an
            // empty map; the signer + chain_id + task_manager triple still
            // matter because the HTTP client embeds them in every outbound
            // envelope.
            let operator_client = Arc::new(
                HttpStateCommitOperatorClient::new(
                    StdHashMap::new(),
                    aggregator_signer.clone(),
                    chain_id,
                    config.contracts.avs.newton_prover_task_manager,
                )
                .map_err(|e| AggregatorError::KeyError(format!("state-commit HTTP client init: {e}")))?,
            );
            let aggregator = Arc::new(StateCommitAggregator::new(chain_id, 6_700));
            let committer = Arc::new(AvsWriterCommitter(Arc::clone(&avs_writer)));
            // Per-tick refresh: production impl (NEWT-1116) composes
            // `ViewBN254CertificateVerifier::latestReferenceTimestamp` with
            // `AvsRegistryServiceCaller`. Until then, the dev-stub returns an
            // empty snapshot so devnet smoke tests run without operators.
            let operator_set_reader = Arc::new(StubOperatorSetSnapshotReader::new(Vec::new(), 0));
            // Dev-stub: empty snapshot until the production wiring (NEWT-1116)
            // composes `BN254TableCalculator.getOperatorInfos` +
            // `BN254CertificateVerifier.latestReferenceTimestamp` against
            // configured contract addresses. The orchestrator's tick will
            // emit `empty_set` flavor against this empty provider — same
            // behavior as the operator-set reader stub above.
            let operator_table_provider =
                state_commit::operator_table::fake::FakeOperatorTableProvider::from_operators(0, Vec::new());
            // Dev-stub operator set: AVS = task manager address, id = 0.
            let operator_set =
                newton_core::bn254_certificate_verifier::ViewBN254CertificateVerifier::OperatorSet {
                    avs: config.contracts.avs.newton_prover_task_manager,
                    id: 0,
                };
            let orch = StateCommitOrchestrator::new(
                chain_id,
                Duration::from_secs(120),
                reader,
                pcr0,
                operator_client,
                aggregator,
                committer,
                operator_set_reader,
                operator_table_provider,
                operator_set,
            );
            Some(Arc::new(orch))
        };
        #[cfg(not(any(test, feature = "dev-stub")))]
        let state_commit_orchestrator: Option<Arc<StateCommitOrchestrator>> = None;

        Ok(Self {
            port_address: config.service.ip_address.clone(),
            avs_writer,
            tasks: HashMap::new(),
            core: Arc::new(core),
            rpc: config.rpc.clone(),
            chain_id,
            task_response_window_block,
            aggregation_timeout_secs: config.service.aggregation_timeout_secs,
            state_commit_orchestrator,
        })
    }

    /// Starts the aggregator service with three tasks: one for the server,
    /// one for processing tasks and one for processing aggregator responses.
    ///
    /// # Arguments
    ///
    /// * `ws_rpc_url` - The websocket RPC URL
    ///
    /// # Returns
    ///
    /// * `eyre::Result<()>` - The result of the operation
    ///
    /// # Errors
    ///
    /// * The error that occurred
    pub async fn start(self, ws_rpc_url: String) -> eyre::Result<()> {
        info!("Starting aggregator");

        // Initialize Prometheus metrics exporter on port 9092
        let metrics_addr: SocketAddr = "0.0.0.0:9092".parse()?;
        let _metrics_handle = newton_metric::prometheus::init_prometheus_exporter(metrics_addr).await?;
        info!("Prometheus metrics server started on {}", metrics_addr);

        let Self {
            avs_writer,
            port_address,
            tasks,
            core,
            rpc: _,
            chain_id: _,
            task_response_window_block,
            aggregation_timeout_secs,
            state_commit_orchestrator,
        } = self;

        let tasks = Arc::new(tokio::sync::Mutex::new(tasks));
        let core_clone = Arc::clone(&core);

        // 4th task — per-chain state-commit orchestrator (120s tick). Runs
        // alongside the join-arm trio under a `CancellationToken`. The
        // orchestrator's `run()` loop is `biased` on the cancel arm, so
        // signalling `cancel()` after `try_join!` resolves drives a clean
        // shutdown without aborting an in-flight tick. In production builds
        // (no `dev-stub`), `state_commit_orchestrator` is `None` until
        // NEWT-1116 wires `EnclavePcr0Provider`.
        let orchestrator_token = CancellationToken::new();
        let orchestrator_handle: Option<JoinHandle<()>> = state_commit_orchestrator.map(|orch| {
            let token = orchestrator_token.clone();
            // `spawn_monitored` wraps the loop in `catch_unwind` so a panic
            // surfaces as a structured tracing event with `task = ...`. The
            // orchestrator's own `tick()` is already panic-isolated, but a
            // panic in the surrounding `select!` / shutdown plumbing would
            // otherwise vanish to stderr — see `.claude/rules/rust.md`
            // "Bare `tokio::spawn` for background loops" anti-pattern.
            crate::task_supervisor::spawn_monitored("state_commit_orchestrator", async move { orch.run(token).await })
        });

        // Spawn three tasks: one for the server, one for processing tasks, and one for processing aggregator responses
        // 1) Process signatures
        let server_handle = Self::start_server(port_address, Arc::clone(&core)).await?;
        // 2) Process tasks
        let process_handle = tokio::spawn(Self::process_tasks(
            ws_rpc_url.clone(),
            task_response_window_block,
            Arc::clone(&tasks),
            Arc::clone(&core),
        ));
        // 3) Process aggregator responses
        let responses_handle = tokio::spawn(Self::process_aggregator_responses(
            Arc::clone(&tasks),
            Arc::clone(&core),
            avs_writer,
            aggregation_timeout_secs,
        ));

        // Wait for the tasks to complete and handle potential errors
        let join_result = tokio::try_join!(server_handle, process_handle, responses_handle)
            .inspect_err(|e| error!("Error in task execution: {e:?}"));

        // Signal the orchestrator to stop and drain its handle before
        // surfacing any join error. `cancel()` is idempotent and the loop
        // exits at the next `select!` poll without aborting an in-flight
        // tick, so we never mid-tick on a `commit_state_root` submission.
        orchestrator_token.cancel();
        if let Some(handle) = orchestrator_handle {
            if let Err(e) = handle.await {
                error!("state-commit orchestrator join failed: {e:?}");
            }
        }

        let (_server_result, process_result, responses_result) =
            join_result.map_err(|e| eyre::eyre!("Task execution failed {e}"))?;

        process_result?;
        responses_result?;

        Ok(())
    }

    /// Starts the RPC server and returns a future that ends once the server is stopped
    async fn start_server(
        port_address: String,
        core: Arc<AggregatorCore>,
    ) -> eyre::Result<JoinHandle<()>, AggregatorError> {
        // See https://github.com/paritytech/jsonrpsee/blob/42461391fee47c94d42c4a7303355525291df9f6/examples/examples/cors_server.rs
        let mut module = RpcModule::new(core);
        module
            .register_async_method("process_signed_task_response", |params, ctx, _| async move {
                let core = ctx.as_ref();
                let signed_task_response = params
                    .parse::<SignedTaskResponse>()
                    .map_err(|err| ErrorObject::owned(0, err.to_string(), None::<()>))?;

                // Call the process_signed_response function from AggregatorCore
                let result = core.process_signed_response(signed_task_response).await;

                match result {
                    Ok(()) => Ok(true),
                    Err(err) => Err(ErrorObject::owned(0, err.to_string(), None::<()>)),
                }
            })
            .expect("method name is unique");
        let socket: SocketAddr = port_address
            .to_socket_addrs()
            .map_err(AggregatorError::IOError)?
            .next()
            .ok_or_else(|| {
                AggregatorError::IOError(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "No addresses found for hostname",
                ))
            })?;
        let middleware = tower::ServiceBuilder::new();
        let server = Server::builder().set_http_middleware(middleware).build(socket).await?;

        let handle = server.start(module);

        info!("Server running at {socket}");

        Ok(tokio::spawn(handle.stopped()))
    }

    /// Processes new tasks created in the task manager contract
    #[instrument(skip_all)]
    async fn process_tasks(
        ws_rpc_url: String,
        task_response_window_block: u64,
        tasks: Arc<tokio::sync::Mutex<HashMap<TaskId, Task>>>,
        core: Arc<AggregatorCore>,
    ) -> eyre::Result<()> {
        let ws = WsConnect::new(ws_rpc_url.clone());
        let provider = ProviderBuilder::new()
            .disable_recommended_fillers()
            .connect_ws(ws)
            .await?;
        let chain_id = provider.get_chain_id().await?;

        let filter = Filter::new().event_signature(NewTaskCreated::SIGNATURE_HASH);
        let sub = provider.subscribe_logs(&filter).await?;
        let mut stream = sub.into_stream();

        while let Some(log) = stream.next().await {
            let NewTaskCreated { taskId, task, state: _ } = log.log_decode()?.inner.data;

            let span = tracing::info_span!(
                "process_task",
                task_id = %taskId,
                policy_client = %task.policyClient
            );

            async {
                info!("aggregator: new task created");

                tasks.lock().await.insert(taskId, task.clone());

                let mut quorum_nums: Vec<u8> = vec![];
                for val in task.quorumNumbers.iter() {
                    quorum_nums.push(*val);
                }
                info!(
                    quorum_nums = %hex!(&quorum_nums),
                    "aggregator: initialize_task quorum_nums"
                );

                let block_time_ms = get_block_duration_ms(chain_id).unwrap_or(12000);
                let time_to_expiry = tokio::time::Duration::from_millis(task_response_window_block * block_time_ms);

                info!("aggregator: initializing task");
                core.initialize_task(
                    taskId,
                    task.taskCreatedBlock.into(),
                    quorum_nums,
                    task.quorumThresholdPercentage.try_into()?,
                    time_to_expiry,
                    0, // Expected operator count unknown for standalone aggregator (disables early exit)
                )
                .await
                .map_err(|e| {
                    error!("Failed to initialize task {}: {:?}", taskId, e);
                    eyre::eyre!(e)
                })?;

                Ok::<(), eyre::Error>(())
            }
            .instrument(span)
            .await?;
        }

        Ok(())
    }

    /// Processes BLS Aggregator Service responses
    ///
    /// With per-task channels, we iterate over active tasks and wait for their responses.
    /// This eliminates mutex contention and response stealing.
    async fn process_aggregator_responses(
        tasks: Arc<tokio::sync::Mutex<HashMap<TaskId, Task>>>,
        core: Arc<AggregatorCore>,
        avs_writer: Arc<AvsWriter>,
        aggregation_timeout_secs: u64,
    ) -> eyre::Result<()> {
        loop {
            // Get list of active task IDs
            let active_task_ids: Vec<TaskId> = {
                let tasks_guard = tasks.lock().await;
                tasks_guard.keys().cloned().collect()
            };

            if active_task_ids.is_empty() {
                // No active tasks, wait a bit before checking again
                tokio::time::sleep(Duration::from_millis(100)).await;
                continue;
            }

            // Filter to only tasks that have receivers (not already consumed)
            // Check receivers efficiently to avoid unnecessary wait_for_aggregation calls
            let tasks_with_receivers: Vec<TaskId> = {
                let mut tasks_with_receivers = Vec::new();
                for task_id in active_task_ids {
                    if core.has_task_receiver(task_id).await {
                        tasks_with_receivers.push(task_id);
                    }
                }
                tasks_with_receivers
            };

            if tasks_with_receivers.is_empty() {
                // No tasks with receivers, wait a bit before checking again
                tokio::time::sleep(Duration::from_millis(200)).await;
                continue;
            }

            // Use a longer timeout (1-2 seconds) to reduce polling overhead
            // This is more efficient than 100ms polling and reduces latency
            // We process tasks sequentially but with a reasonable timeout per task
            let mut found_response = false;
            for task_id in tasks_with_receivers {
                // Double-check receiver still exists (may have been consumed by another call)
                if !core.has_task_receiver(task_id).await {
                    continue;
                }

                // Use configurable timeout to accommodate destination chain operations
                // (certificate verifier queries add ~500ms-1s of latency)
                let timeout = Duration::from_secs(aggregation_timeout_secs);
                match core.wait_for_aggregation(task_id, timeout).await {
                    Ok(service_response) => {
                        found_response = true;
                        info!("Received aggregated response for task_id: {}", service_response.task_id);

                        let task = match tasks.lock().await.get(&service_response.task_id).cloned() {
                            Some(task) => task,
                            None => {
                                info!("Task not found for task_id: {}", service_response.task_id);
                                continue;
                            }
                        };

                        let task_response = core.task_states.get(&task.taskId).and_then(|state| {
                            state
                                .task_responses
                                .get(&service_response.task_response_digest)
                                .cloned()
                        });

                        if let Some(task_response) = task_response {
                            if let Err(err) = Self::submit_aggregated_response_with_retry(
                                &core,
                                &avs_writer,
                                &task,
                                &task_response,
                                &service_response,
                            )
                            .await
                            {
                                error!(
                                    task_id = hex!(task.taskId),
                                    error = ?err,
                                    "Failed to submit aggregated response to contract"
                                );
                            }
                        } else {
                            info!(
                                "Not found task_response for task_id: {}",
                                hex!(service_response.task_id)
                            );
                        }
                        break; // Process one response per loop iteration
                    }
                    Err(crate::error::AggregatorCoreError::Timeout { .. }) => {
                        // Timeout is expected, continue to next task
                        continue;
                    }
                    Err(crate::error::AggregatorCoreError::TaskNotInitialized { .. }) => {
                        // Receiver already consumed (by gateway or previous call) - skip this task
                        continue;
                    }
                    Err(crate::error::AggregatorCoreError::InvalidTimeoutDuration(_)) => {
                        // Invalid timeout - this shouldn't happen, but handle gracefully
                        warn!(task_id = %task_id, "Invalid timeout duration - skipping task");
                        continue;
                    }
                    Err(e) => {
                        error!(task_id = %task_id, "Error receiving aggregated response: {:?}", e);
                        continue;
                    }
                }
            }

            if !found_response {
                // No responses found, wait a bit before checking again
                // Reduced from 50ms to avoid excessive CPU usage when no tasks are active
                tokio::time::sleep(Duration::from_millis(100)).await;
            }
        }
    }

    async fn submit_aggregated_response_with_retry(
        core: &AggregatorCore,
        avs_writer: &AvsWriter,
        task: &Task,
        task_response: &BindingTaskResponse,
        service_response: &BlsAggregationServiceResponse,
    ) -> Result<(), AggregatorError> {
        const MAX_RETRIES: usize = 3;
        let mut attempt: usize = 0;
        let mut delay = Duration::from_millis(500);

        loop {
            attempt += 1;

            let result = core
                .submit_aggregated_response(
                    avs_writer,
                    task.clone(),
                    task_response.clone(),
                    service_response.clone(),
                )
                .await
                .map(|_| ())
                .map_err(|_| AggregatorError::SendAggregatedResponseError);

            match result {
                Ok(()) => return Ok(()),
                Err(err) => {
                    let err_str = err.to_string();
                    let should_retry =
                        err_str.contains("SendAggregatedResponseError") || err_str.contains("Failed to submit");
                    if !should_retry || attempt >= MAX_RETRIES {
                        return Err(AggregatorError::SendAggregatedResponseError);
                    }

                    warn!(
                        task_id = hex!(service_response.task_id),
                        attempt,
                        retry_in_secs = delay.as_secs(),
                        "Failed to submit aggregated response; retrying"
                    );

                    sleep(delay).await;
                    delay = (delay.saturating_mul(2)).min(Duration::from_millis(2000));
                }
            }
        }
    }

    fn should_retry_contract_submission(err: &AggregatorError) -> bool {
        matches!(err, AggregatorError::SendAggregatedResponseError)
    }
}