refactor(e2e): observe on-chain outcomes past the RPC wait window

crates/e2e-tests/Cargo.toml

@@ -32,6 +32,7 @@ serde_json = { workspace = true }
 tempfile = { workspace = true }
 test-port-allocator = { workspace = true }
 test-utils = { workspace = true }
+thiserror = { workspace = true }
 tokio = { workspace = true }
 toml = { workspace = true }
 tracing = { workspace = true }

crates/e2e-tests/src/blockchain.rs

@@ -2,10 +2,13 @@ use std::time::Duration;
 
 use ed25519_dalek::SigningKey;
 use near_contract_transport::{CallContract, FunctionCallArgs};
-use near_kit::{AccountId, CryptoHash, Final, FinalExecutionOutcome};
+use near_kit::{
+    CallBuilder, CryptoHash, ExecutedOptimistic, Final, FinalExecutionOutcome, Included,
+};
 use near_mpc_contract_interface::client::MpcContractHandle;
 use near_mpc_contract_interface::types::ProtocolContractState;
 use serde::de::DeserializeOwned;
+use tokio::time::Instant;
 
 use crate::conversions::ToNearKey;
 
@@ -22,28 +25,147 @@ pub struct NearBlockchain {
     rpc_url: String,
 }
 
-/// A `near_kit::Near` client bound to a specific account: the e2e
-/// [`CallContract`] backend.
+/// A [`near_kit::Near`] client bound to a specific account: the e2e [`CallContract`]
+/// backend.
+///
+/// Without a timeout a call is bounded by nearcore's RPC polling window, so a request
+/// outliving it — a `sign` whose yield is still open — fails with
+/// [`CallError::RpcGaveUp`]. [`Self::with_timeout`] waits for the outcome instead.
 pub struct NearKitCaller {
     inner: near_kit::Near,
+    timeout: Option<Duration>,
+}
+
+impl NearKitCaller {
+    /// Observe the outcome for up to `timeout` instead of for as long as the RPC waits.
+    pub fn with_timeout(self, timeout: Duration) -> Self {
+        Self {
+            timeout: Some(timeout),
+            ..self
+        }
+    }
+
+    /// Bounded by nearcore's RPC polling window, across near-kit's retries.
+    async fn send(&self, call: CallBuilder) -> Result<FinalExecutionOutcome, CallError> {
+        let started = Instant::now();
+        call.send().await.map_err(|e| {
+            if timed_out_waiting(&e) {
+                CallError::RpcGaveUp {
+                    after: started.elapsed(),
+                    source: e,
+                }
+            } else {
+                CallError::Rpc(e)
+            }
+        })
+    }
+
+    /// Submitted once: retrying `send_tx` would re-broadcast, so only the polling repeats.
+    async fn send_and_observe(
+        &self,
+        call: CallBuilder,
+        timeout: Duration,
+    ) -> Result<FinalExecutionOutcome, CallError> {
+        // Waiting for inclusion first stops the poll below racing a hash the RPC has
+        // not seen yet.
+        let tx = call
+            .wait_until::<Included>()
+            .await
+            .map_err(CallError::Rpc)?
+            .transaction_hash;
+
+        // May not panic: submitting the call above already required a signer.
+        let sender = self.inner.account_id();
+
+        let poll = async {
+            loop {
+                match self
+                    .inner
+                    .tx_status(&tx, sender)
+                    .wait_until::<ExecutedOptimistic>()
+                    .await
+                {
+                    Ok(outcome) => return Ok(outcome),
+                    Err(e) if !is_retryable(&e) => return Err(CallError::Rpc(e)),
+                    Err(_) => tokio::time::sleep(TX_STATUS_POLL_INTERVAL).await,
+                }
+            }
+        };
+
+        tokio::time::timeout(timeout, poll)
+            .await
+            .unwrap_or_else(|_| Err(CallError::Deadline { tx, after: timeout }))
+    }
+}
+
+/// Wait a caller-chosen duration for a call's on-chain outcome.
+pub trait WithTimeout: Sized {
+    fn with_timeout(self, timeout: Duration) -> Self;
+}
+
+impl WithTimeout for MpcContractHandle<NearKitCaller> {
+    fn with_timeout(self, timeout: Duration) -> Self {
+        self.map_caller(|caller| caller.with_timeout(timeout))
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum CallError {
+    #[error(transparent)]
+    Rpc(near_kit::Error),
+
+    /// On chain and still unresolved — typically a `sign` whose yield has not been
+    /// answered. Widen the window with [`WithTimeout::with_timeout`].
+    #[error("tx {tx} still unresolved on chain after {after:?}")]
+    Deadline { tx: CryptoHash, after: Duration },
+
+    /// Inclusion was never confirmed, so whether the transaction reached the chain is
+    /// unknown. [`WithTimeout::with_timeout`] confirms it and waits for the outcome.
+    #[error("RPC stopped waiting after {after:?}; outcome unobserved")]
+    RpcGaveUp {
+        after: Duration,
+        #[source]
+        source: near_kit::Error,
+    },
 }
 
 impl CallContract for NearKitCaller {
     type Output = FinalExecutionOutcome;
-    type Error = near_kit::Error;
+    type Error = CallError;
 
     async fn call_contract(
         &self,
         contract_id: &near_kit::AccountId,
         call_args: FunctionCallArgs,
     ) -> Result<Self::Output, Self::Error> {
-        self.inner
+        let call = self
+            .inner
             .call(contract_id, &call_args.method_name)
             .args_raw(call_args.args)
             .gas(call_args.gas)
-            .deposit(call_args.deposit)
-            .send()
-            .await
+            .deposit(call_args.deposit);
+
+        match self.timeout {
+            None => self.send(call).await,
+            Some(timeout) => self.send_and_observe(call, timeout).await,
+        }
+    }
+}
+
+/// Polling through anything else would spin until the deadline, then report its cause as
+/// a timeout.
+fn is_retryable(error: &near_kit::Error) -> bool {
+    matches!(error, near_kit::Error::Rpc(rpc) if rpc.is_retryable())
+}
+
+/// Distinguishes "the RPC stopped waiting" from "the call failed".
+fn timed_out_waiting(error: &near_kit::Error) -> bool {
+    match error {
+        near_kit::Error::Rpc(rpc) => matches!(
+            **rpc,
+            near_kit::RpcError::Timeout(_) | near_kit::RpcError::RequestTimeout { .. }
+        ),
+        _ => false,
     }
 }
 
@@ -114,6 +236,7 @@ impl NearBlockchain {
     pub fn client_for(&self, account_id: &str, key: &SigningKey) -> anyhow::Result<NearKitCaller> {
         Ok(NearKitCaller {
             inner: self.make_client(account_id, key)?,
+            timeout: None,
         })
     }
 
@@ -227,54 +350,6 @@ impl DeployedContract {
             })
     }
 
-    /// Submit a call and return its tx hash once included. Pair with [`Self::wait_tx_final`]
-    /// for requests (e.g. a yielding `sign`) that resolve past the RPC timeout.
-    pub async fn call_from_with_deposit_included(
-        &self,
-        client: &NearKitCaller,
-        method: &str,
-        args: serde_json::Value,
-        gas: near_kit::Gas,
-        deposit: near_kit::NearToken,
-    ) -> anyhow::Result<CryptoHash> {
-        let response = client
-            .inner
-            .call(&self.contract_id, method)
-            .args(args)
-            .gas(gas)
-            .deposit(deposit)
-            .wait_until::<near_kit::Included>()
-            .await
-            .map_err(|e| anyhow::anyhow!("contract call `{method}` (included) failed: {e}"))?;
-        Ok(response.transaction_hash)
-    }
-
-    /// Poll `tx_status` for `tx_hash` (sent by `signer_id`) until `Final` or `timeout`.
-    pub async fn wait_tx_final(
-        &self,
-        tx_hash: CryptoHash,
-        signer_id: &AccountId,
-        timeout: Duration,
-    ) -> anyhow::Result<FinalExecutionOutcome> {
-        let deadline = tokio::time::Instant::now() + timeout;
-        loop {
-            match self
-                .client
-                .tx_status(&tx_hash, signer_id.as_str())
-                .wait_until::<Final>()
-                .await
-            {
-                Ok(outcome) => return Ok(outcome),
-                Err(e) if tokio::time::Instant::now() >= deadline => {
-                    return Err(anyhow::anyhow!(
-                        "tx {tx_hash} did not reach Final within {timeout:?}: {e}"
-                    ));
-                }
-                Err(_) => tokio::time::sleep(TX_STATUS_POLL_INTERVAL).await,
-            }
-        }
-    }
-
     pub async fn view<T: DeserializeOwned + Send + 'static>(
         &self,
         method: &str,

crates/e2e-tests/src/cluster.rs

@@ -9,7 +9,6 @@ use near_kit::AccountId;
 use near_mpc_bounded_collections::NonEmptyBTreeMap;
 use near_mpc_contract_interface::types::CKDRequestArgs;
 use near_mpc_contract_interface::{
-    call_args::SignArgs,
     client::MpcContractHandle,
     method_names,
     types::{
@@ -837,33 +836,6 @@ impl MpcCluster {
             .context("failed to send sign request")
     }
 
-    /// Like [`Self::send_sign_request`], but returns the tx hash once included instead
-    /// of awaiting it — for requests that resolve past the RPC timeout.
-    pub async fn send_sign_request_included(
-        &self,
-        domain_id: DomainId,
-        payload: Payload,
-        account_id: &AccountId,
-    ) -> anyhow::Result<near_kit::CryptoHash> {
-        let client = self.client_for(account_id)?;
-        let args = SignArgs::new(SignRequestArgs {
-            path: "test".to_string(),
-            payload,
-            domain_id,
-        });
-        self.contract
-            .call_from_with_deposit_included(
-                &client,
-                method_names::SIGN,
-                serde_json::to_value(&args)?,
-                near_mpc_contract_interface::client::SIGN_GAS,
-                near_kit::NearToken::from_yoctonear(
-                    near_mpc_contract_interface::deposits::SIGN_DEPOSIT_YOCTONEAR,
-                ),
-            )
-            .await
-    }
-
     /// Send a CKD (Confidential Key Derivation) request from the given user account.
     pub async fn send_ckd_request(
         &self,

crates/e2e-tests/src/lib.rs

@@ -6,7 +6,7 @@ pub mod metrics;
 pub mod mpc_node;
 pub mod near_sandbox;
 
-pub use blockchain::{DeployedContract, NearBlockchain, NearKitCaller};
+pub use blockchain::{CallError, DeployedContract, NearBlockchain, NearKitCaller, WithTimeout};
 pub use cluster::{
     CLUSTER_WAIT_TIMEOUT, DEFAULT_PRESIGNATURES_TO_BUFFER, DEFAULT_TRIPLES_TO_BUFFER, MpcCluster,
     MpcClusterConfig, MpcNodeState,

crates/e2e-tests/tests/distinct_reconstruction_thresholds.rs

@@ -4,9 +4,9 @@ use crate::common::{
     must_setup_cluster,
 };
 
-use e2e_tests::CLUSTER_WAIT_TIMEOUT;
+use e2e_tests::{CLUSTER_WAIT_TIMEOUT, WithTimeout};
 use near_mpc_contract_interface::types::{
-    DomainConfig, DomainId, DomainPurpose, Protocol, ReconstructionThreshold,
+    DomainConfig, DomainId, DomainPurpose, Protocol, ReconstructionThreshold, SignRequestArgs,
 };
 use rand::{SeedableRng, rngs::StdRng};
 
@@ -97,25 +97,25 @@ async fn distinct_reconstruction_thresholds__should_use_per_domain_threshold_whe
         outcome.failure_message()
     );
 
-    // Cait-Sith (needs all 6) is unanswerable: its yield outlives the JSON-RPC call, so
-    // we poll the tx to `Final` instead of awaiting it.
-    let user = cluster.default_user_account().clone();
-    let tx_hash = cluster
-        .send_sign_request_included(caitsith_domain.id, generate_ecdsa_payload(&mut rng), &user)
-        .await
-        .expect("failed to submit Cait-Sith sign request");
+    // Cait-Sith needs all 6, so its yield runs to the on-chain timeout — past the RPC's
+    // wait window, hence the deadline.
     let outcome = cluster
-        .contract
-        .wait_tx_final(tx_hash, &user, CLUSTER_WAIT_TIMEOUT)
+        .contract_handle(cluster.default_user_account())
+        .with_timeout(CLUSTER_WAIT_TIMEOUT)
+        .sign(SignRequestArgs {
+            path: "test".to_string(),
+            payload: generate_ecdsa_payload(&mut rng),
+            domain_id: caitsith_domain.id,
+        })
         .await
-        .expect("Cait-Sith sign request did not reach a final on-chain outcome");
+        .expect("Cait-Sith sign request did not reach an on-chain outcome");
     assert!(
         outcome.is_failure(),
         "Cait-Sith sign request succeeded with only 5 of its 6 required signers alive"
     );
     let message = outcome.failure_message().unwrap_or_default();
     assert!(
-        message.contains("timed out"),
+        message.contains("Request has timed out."),
         "Cait-Sith sign request failed for an unexpected reason: {message}"
     );
 }