miden-multisig-client 0.17.0

//! Recovery primitives for restoring note state after device loss.
//!
//! After recovery on a fresh device the local store has no note-transport
//! cursor — and in a shared dirty store another account's sync may have
//! advanced the cursor past notes belonging to the newly recovered account.
//! The primitives here rescan sources that normal forward sync would skip.

use std::collections::BTreeSet;

use miden_client::ClientError;
use miden_client::note_transport::{NOTE_TRANSPORT_COVERED_TAGS_KEY, NoteTransportError};
use miden_client::sync::NoteTagSource;
use miden_protocol::note::NoteTag;

use super::MultisigClient;
use crate::MidenSdkClient;
use crate::error::{Result, error_chain};
use crate::rpc::{is_transient_note_transport_error, is_transient_rpc_error};

/// Outcome class of a private-note transport backlog drain.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransportRecoveryStatus {
    /// The full transport backlog was scanned; every note the transport still
    /// holds for the tracked tags is now in the local store.
    Completed,
    /// The transport could not be consulted at all: it is disabled for this
    /// client (no endpoint configured) or unreachable before anything was
    /// imported (`imported` is always 0 — a connection lost mid-drain after
    /// partial progress reports `Failed` instead). The rest of a recovery
    /// flow should proceed without transport notes.
    Unavailable,
    /// The drain started but did not finish; the backlog may be partially
    /// imported. `retryable` distinguishes transient failures (rerun the
    /// drain) from permanent ones.
    Failed,
}

/// Result of [`MultisigClient::drain_private_note_backlog`].
///
/// Transport problems are reported here rather than returned as errors so a
/// transport failure never aborts the rest of a recovery flow.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TransportRecoveryReport {
    /// Outcome class of the drain.
    pub status: TransportRecoveryStatus,
    /// Number of note records newly imported into the local store by this
    /// drain. Can be non-zero even on `Failed`: batches imported before the
    /// failure stay imported.
    pub imported: usize,
    /// Whether rerunning the drain can plausibly succeed (transient
    /// connectivity failures, the upstream pagination convergence guard).
    /// Always `false` on `Completed`.
    pub retryable: bool,
    /// Human-readable cause when the drain did not complete.
    pub reason: Option<String>,
}

impl MultisigClient {
    /// Rescans the full private-note transport backlog for every tracked note
    /// tag and imports what it finds, regardless of the stored transport
    /// cursor.
    ///
    /// Use after account recovery: a fresh store has no transport cursor, and
    /// in a shared store another account's sync may have advanced the cursor
    /// past this account's notes. The drain is idempotent, tag-scoped (the
    /// recovered account must already be in the store so its note tag is
    /// tracked — [`MultisigClient::pull_account`] does this), and never
    /// regresses an already-advanced cursor.
    ///
    /// Transport recovery is bounded by the transport service's retention:
    /// senders may bypass the transport entirely and relayed blobs are pruned
    /// after the retention window, so this is a best-effort rescan, **not** a
    /// backup. Transport-disabled and transport-unreachable outcomes are
    /// reported in the [`TransportRecoveryReport`] rather than returned as
    /// errors; an `Err` from this method means the local store itself failed.
    ///
    /// The rescan runs as many transport syncs as the upstream per-sync tag
    /// backfill cap requires to cover every tracked candidate tag, and a
    /// failed drain restores the pre-drain covered-tags bookkeeping so
    /// normal sync keeps working exactly as it did before the attempt.
    pub(crate) async fn drain_private_note_backlog(&mut self) -> Result<TransportRecoveryReport> {
        if !self.miden_client.is_note_transport_enabled() {
            return Ok(TransportRecoveryReport {
                status: TransportRecoveryStatus::Unavailable,
                imported: 0,
                retryable: false,
                reason: Some(
                    "note transport is not configured for this client; \
                     set a transport endpoint to relay private notes"
                        .to_string(),
                ),
            });
        }

        let before = self.input_note_count().await?;
        // Snapshot the covered-tags set before clearing it: the clear is
        // durable, and upstream re-marks a tag covered only after its
        // backfill succeeds — so without a restore, a drain that fails on a
        // tag with a permanently bad relay blob would leave every tag
        // uncovered and make every subsequent normal sync re-attempt (and
        // fail) the same backfill. Restoring on failure returns the client
        // to its working pre-drain state.
        let covered_snapshot = self.covered_tags_snapshot().await?;
        // miden-client 0.16 replaced the explicit full drain with covered-tag
        // bookkeeping inside `sync_note_transport`: each tag not yet marked
        // covered is drained from the start with a local cursor (the global
        // cursor is never regressed), then the steady-state fetch runs.
        // Clearing the covered-tags marker first forces that full per-tag
        // re-drain for every tracked tag, which is exactly the recovery
        // semantic this primitive promises. Imports dedupe, so re-draining
        // already-seen history is harmless. Upstream backfills at most
        // `MAX_BACKFILL_TAGS_PER_SYNC` uncovered tags per call, so run
        // enough passes to cover every candidate tag before reporting the
        // backlog fully scanned.
        let passes = self
            .backfill_candidate_count()
            .await?
            .div_ceil(MidenSdkClient::MAX_BACKFILL_TAGS_PER_SYNC)
            .max(1);
        let drain_result = match self
            .miden_client
            .remove_setting(NOTE_TRANSPORT_COVERED_TAGS_KEY.to_string())
            .await
        {
            Ok(_) => {
                let mut result = Ok(());
                for _ in 0..passes {
                    if let Err(err) = self.miden_client.sync_note_transport().await {
                        result = Err(err);
                        break;
                    }
                }
                result
            }
            Err(err) => Err(err),
        };
        if let Err(drain_err) = &drain_result
            && let Err(restore_err) = self.restore_covered_tags(&covered_snapshot).await
        {
            // The drain failed AND the store write restoring the
            // covered-tags bookkeeping failed: the store is left altered and
            // subsequent normal syncs may re-drain (and re-fail on) old
            // transport history. That is a local-store environment failure
            // the whole recovery flow must see, not a transport outcome.
            return Err(crate::error::MultisigError::miden_client_with_context(
                format!(
                    "the transport drain failed ({}) and restoring the covered-tags \
                     bookkeeping also failed; subsequent syncs may re-drain old transport \
                     history",
                    error_chain(drain_err)
                ),
                restore_err,
            ));
        }
        // Count even when the drain failed: each fetched batch is imported as
        // it arrives, so notes recovered before the failure stay in the store.
        // A drain never removes records (and `&mut self` excludes concurrent
        // client operations), so the length delta is the count of newly
        // imported records.
        let after = self.input_note_count().await?;
        let imported = after.saturating_sub(before);

        match drain_result {
            Ok(()) => Ok(TransportRecoveryReport {
                status: TransportRecoveryStatus::Completed,
                imported,
                retryable: false,
                reason: None,
            }),
            // A broken local store is an environment failure, not a transport
            // outcome: the whole recovery flow needs to know, so it
            // propagates instead of being folded into the report.
            Err(err @ ClientError::StoreError(_)) => {
                Err(crate::error::MultisigError::miden_client_with_context(
                    "local store failed during the transport drain",
                    err,
                ))
            }
            Err(err) => {
                let (status, retryable) = classify_drain_failure(&err);
                // `Unavailable` promises "nothing was imported"; a connection
                // lost mid-drain after partial progress is an interrupted
                // drain, so report it as a retryable failure instead.
                let status = if imported > 0 && status == TransportRecoveryStatus::Unavailable {
                    TransportRecoveryStatus::Failed
                } else {
                    status
                };
                Ok(TransportRecoveryReport {
                    status,
                    imported,
                    retryable,
                    reason: Some(error_chain(&err)),
                })
            }
        }
    }

    /// The covered-tags set as stored, for restore-on-failure. Unreadable
    /// bytes decode to an empty set — upstream resets the entry to empty on
    /// load, so that is the state the drain actually starts from.
    async fn covered_tags_snapshot(&mut self) -> Result<BTreeSet<NoteTag>> {
        match self
            .miden_client
            .get_setting(NOTE_TRANSPORT_COVERED_TAGS_KEY.to_string())
            .await
        {
            Ok(Some(tags)) => Ok(tags),
            Ok(None) => Ok(BTreeSet::new()),
            Err(err @ ClientError::StoreError(_)) => {
                Err(crate::error::MultisigError::miden_client_with_context(
                    "failed to read the covered-tags setting before the transport drain",
                    err,
                ))
            }
            Err(_) => Ok(BTreeSet::new()),
        }
    }

    /// Restores the covered-tags set after a failed drain: the pre-drain
    /// covered set united with whatever the interrupted backfill already
    /// re-covered. A tag restored here was covered before the drain, so
    /// restoring it only returns the client to its working pre-drain state;
    /// a later successful drain re-covers everything from scratch anyway.
    /// The read of the partial progress is best-effort, but a failed
    /// restoring write is returned — the caller must not swallow a store
    /// left with its bookkeeping cleared.
    // See `existing_records_by_commitment`: the Err variant is upstream's `ClientError`.
    #[allow(clippy::result_large_err)]
    async fn restore_covered_tags(
        &mut self,
        snapshot: &BTreeSet<NoteTag>,
    ) -> std::result::Result<(), ClientError> {
        if snapshot.is_empty() {
            return Ok(());
        }
        let current: BTreeSet<NoteTag> = self
            .miden_client
            .get_setting(NOTE_TRANSPORT_COVERED_TAGS_KEY.to_string())
            .await
            .ok()
            .flatten()
            .unwrap_or_default();
        let merged: BTreeSet<NoteTag> = snapshot.union(&current).copied().collect();
        self.miden_client
            .set_setting(NOTE_TRANSPORT_COVERED_TAGS_KEY.to_string(), merged)
            .await
    }

    /// Number of tags the transport backfill considers candidates (`User`-
    /// and `Account`-source, matching upstream `backfill_candidate_tags`).
    async fn backfill_candidate_count(&mut self) -> Result<usize> {
        let tags = self.miden_client.get_note_tags().await.map_err(|e| {
            crate::error::MultisigError::miden_client_with_context("failed to list note tags", e)
        })?;
        Ok(tags
            .iter()
            .filter(|record| {
                matches!(
                    record.source,
                    NoteTagSource::User | NoteTagSource::Account(_)
                )
            })
            .count())
    }

    /// Number of input note records currently in the local store.
    async fn input_note_count(&mut self) -> Result<usize> {
        let records = self
            .miden_client
            .get_input_notes(miden_client::store::NoteFilter::All)
            .await
            .map_err(|e| {
                crate::error::MultisigError::miden_client_with_context(
                    "failed to list input notes",
                    e,
                )
            })?;
        Ok(records.len())
    }
}

/// Maps a transport-drain (`sync_note_transport`) failure onto a report
/// class: `(status, retryable)`.
fn classify_drain_failure(err: &ClientError) -> (TransportRecoveryStatus, bool) {
    match err {
        // The match is deliberately exhaustive so a new upstream variant is a
        // compile error here rather than a silently wrong classification.
        ClientError::NoteTransportError(transport_err) => match transport_err {
            // No transport configured — retrying cannot help until the client
            // is rebuilt with an endpoint.
            NoteTransportError::Disabled => (TransportRecoveryStatus::Unavailable, false),
            // `Connection` wraps endpoint parsing, TLS configuration, and
            // actual connect failures indiscriminately; the shared classifier
            // inspects the cause chain to tell a retry-worthy dropped
            // connection from a permanently misconfigured endpoint.
            NoteTransportError::Connection(_) => (
                TransportRecoveryStatus::Unavailable,
                is_transient_note_transport_error(transport_err),
            ),
            // The transport answered with an error — worth retrying once the
            // service recovers.
            NoteTransportError::Network(_) => (TransportRecoveryStatus::Unavailable, true),
            // The upstream convergence guard tripped (the server cursor kept
            // advancing for 1000 iterations without an empty batch) — a
            // server-side bug, not an honest backlog. Retryable in the sense
            // that a rerun is safe (imports are idempotent) and succeeds once
            // the server recovers; while the server misbehaves, each rerun
            // repeats the same full scan and trips the same guard.
            NoteTransportError::PaginationDidNotTerminate(_) => {
                (TransportRecoveryStatus::Failed, true)
            }
            // Undecodable payloads: a rerun would hit the same bytes again.
            NoteTransportError::Deserialization(_) => (TransportRecoveryStatus::Failed, false),
        },
        // Each fetched batch is imported through the node (inclusion-proof
        // lookup), so a node RPC failure interrupts the drain mid-way; the
        // shared gRPC classifier decides whether a rerun can help.
        ClientError::RpcError(rpc_err) => (
            TransportRecoveryStatus::Failed,
            is_transient_rpc_error(rpc_err),
        ),
        _ => (TransportRecoveryStatus::Failed, false),
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;
    use std::sync::Arc;

    use miden_client::ClientError;
    use miden_client::note_transport::{NoteTransportClient, NoteTransportError};
    use miden_client::rpc::Endpoint;
    use miden_client::testing::note_transport::{MockNoteTransportApi, MockNoteTransportNode};
    use miden_protocol::note::{Note, NoteTag, NoteType};
    use miden_tx::utils::sync::RwLock;

    use super::*;
    use crate::client::test_support::{
        add_to_transport, mock_transport, offline_client, p2id_note_for, test_wallet,
    };

    // ---------------------------------------------------------------------
    // classification
    // ---------------------------------------------------------------------

    fn transport_error(err: NoteTransportError) -> ClientError {
        ClientError::NoteTransportError(err)
    }

    #[test]
    fn disabled_transport_is_unavailable_and_not_retryable() {
        let (status, retryable) =
            classify_drain_failure(&transport_error(NoteTransportError::Disabled));
        assert_eq!(status, TransportRecoveryStatus::Unavailable);
        assert!(!retryable);
    }

    #[test]
    fn unreachable_transport_is_unavailable_and_retryable() {
        let connection = NoteTransportError::Connection(Box::new(std::io::Error::new(
            std::io::ErrorKind::ConnectionRefused,
            "connection refused",
        )));
        let (status, retryable) = classify_drain_failure(&transport_error(connection));
        assert_eq!(status, TransportRecoveryStatus::Unavailable);
        assert!(retryable);

        let network = NoteTransportError::Network("transport 503".to_string());
        let (status, retryable) = classify_drain_failure(&transport_error(network));
        assert_eq!(status, TransportRecoveryStatus::Unavailable);
        assert!(retryable);
    }

    /// `Connection` also wraps endpoint-parse/TLS misconfiguration; the
    /// shared cause-chain classifier marks those permanent so a recovery
    /// flow does not loop retrying a client that can never connect.
    #[test]
    fn misconfigured_transport_endpoint_is_unavailable_and_not_retryable() {
        let connection = NoteTransportError::Connection(Box::new(std::io::Error::other(
            "invalid uri: missing scheme",
        )));
        let (status, retryable) = classify_drain_failure(&transport_error(connection));
        assert_eq!(status, TransportRecoveryStatus::Unavailable);
        assert!(!retryable);
    }

    #[test]
    fn pagination_convergence_guard_is_a_retryable_failure() {
        let (status, retryable) = classify_drain_failure(&transport_error(
            NoteTransportError::PaginationDidNotTerminate(1_000),
        ));
        assert_eq!(status, TransportRecoveryStatus::Failed);
        assert!(retryable);
    }

    #[test]
    fn node_rpc_failure_mid_drain_is_a_retryable_failure() {
        let err: ClientError = miden_client::rpc::RpcError::RequestError {
            endpoint: miden_client::rpc::RpcEndpoint::SyncChainMmr,
            error_kind: miden_client::rpc::GrpcError::Unavailable,
            endpoint_error: None,
            source: None,
        }
        .into();
        let (status, retryable) = classify_drain_failure(&err);
        assert_eq!(status, TransportRecoveryStatus::Failed);
        assert!(retryable);
    }

    /// The shared gRPC classifier decides retryability for node failures:
    /// a permanent status (bad request, wrong node version) must not tell
    /// callers to rerun the drain.
    #[test]
    fn permanent_node_rpc_failure_mid_drain_is_not_retryable() {
        let err: ClientError = miden_client::rpc::RpcError::RequestError {
            endpoint: miden_client::rpc::RpcEndpoint::SyncChainMmr,
            error_kind: miden_client::rpc::GrpcError::InvalidArgument,
            endpoint_error: None,
            source: None,
        }
        .into();
        let (status, retryable) = classify_drain_failure(&err);
        assert_eq!(status, TransportRecoveryStatus::Failed);
        assert!(!retryable);
    }

    #[test]
    fn other_client_errors_are_permanent_failures() {
        let (status, retryable) = classify_drain_failure(&ClientError::AddNewAccountWithoutSeed);
        assert_eq!(status, TransportRecoveryStatus::Failed);
        assert!(!retryable);
    }

    // ---------------------------------------------------------------------
    // offline behavioral tests (mock chain + mock transport)
    // ---------------------------------------------------------------------

    /// A distinct (per `seed`) private P2ID note addressed at `target`.
    fn private_note_for(target: &miden_protocol::account::Account, seed: u32) -> Note {
        p2id_note_for(target, seed, NoteType::Private)
    }

    #[tokio::test]
    async fn drain_reports_unavailable_when_transport_is_not_configured() {
        let dir = tempfile::tempdir().unwrap();
        let mut client = offline_client(dir.path(), None).await;

        let report = client
            .drain_private_note_backlog()
            .await
            .expect("disabled transport is reported, not thrown");

        assert_eq!(report.status, TransportRecoveryStatus::Unavailable);
        assert_eq!(report.imported, 0);
        assert!(!report.retryable);
        assert!(report.reason.unwrap().contains("not configured"));
    }

    #[tokio::test]
    async fn drain_recovers_the_backlog_for_a_tracked_account_and_is_idempotent() {
        let dir = tempfile::tempdir().unwrap();
        let (node, api) = mock_transport();
        let mut client = offline_client(dir.path(), Some(api)).await;

        // The recovered account is in the store (as after pull_account), so
        // its standard note tag is tracked.
        let account = test_wallet(1);
        client.add_or_update_account(&account, false).await.unwrap();
        add_to_transport(&node, private_note_for(&account, 1));

        let report = client.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Completed);
        assert_eq!(report.imported, 1);
        assert!(!report.retryable);
        assert_eq!(report.reason, None);

        // Idempotence: draining again re-fetches the same backlog but imports
        // nothing new.
        let report = client.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Completed);
        assert_eq!(report.imported, 0);
    }

    #[tokio::test]
    async fn drain_is_tag_scoped_a_store_with_no_tracked_tags_imports_nothing() {
        let dir = tempfile::tempdir().unwrap();
        let (node, api) = mock_transport();
        let mut client = offline_client(dir.path(), Some(api)).await;

        // A note is waiting on the transport, but no account is in the store,
        // so no tag is tracked.
        add_to_transport(&node, private_note_for(&test_wallet(2), 1));

        let report = client.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Completed);
        assert_eq!(report.imported, 0);
    }

    /// Clearing the covered set makes every tracked tag a backfill candidate
    /// at once, and upstream backfills at most
    /// `MAX_BACKFILL_TAGS_PER_SYNC` tags per transport sync — the drain must
    /// keep syncing until every tag's backlog is recovered instead of
    /// reporting `Completed` after the first 64.
    #[tokio::test]
    async fn drain_covers_more_tags_than_the_per_sync_backfill_cap() {
        let dir = tempfile::tempdir().unwrap();
        let (node, api) = mock_transport();
        let mut client = offline_client(dir.path(), Some(api)).await;

        let count = MidenSdkClient::MAX_BACKFILL_TAGS_PER_SYNC + 1;
        for seed in 0..count {
            let account = test_wallet(seed as u8);
            client.add_or_update_account(&account, false).await.unwrap();
            add_to_transport(
                &node,
                p2id_note_for(&account, seed as u32 + 1, NoteType::Private),
            );
        }

        let report = client.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Completed);
        assert_eq!(
            report.imported, count,
            "every tag's backlog must be recovered, not just the first backfill batch"
        );
    }

    /// A failed drain must restore the pre-drain covered-tags set: leaving
    /// it cleared would make every subsequent normal sync re-attempt (and
    /// fail) the same per-tag backfill, breaking a client that synced fine
    /// before the drain.
    #[tokio::test]
    async fn failed_drain_restores_the_covered_tags_snapshot() {
        let dir = tempfile::tempdir().unwrap();
        let node = Arc::new(RwLock::new(MockNoteTransportNode::new()));
        let api: Arc<dyn NoteTransportClient> = Arc::new(InterruptibleTransport {
            inner: MockNoteTransportApi::new(node.clone()),
            fetches_before_failure: std::sync::atomic::AtomicUsize::new(0),
        });
        let mut client = offline_client(dir.path(), Some(api)).await;

        let account = test_wallet(7);
        client.add_or_update_account(&account, false).await.unwrap();
        add_to_transport(&node, private_note_for(&account, 1));

        // Pre-drain state: the account's tag is covered, so normal sync
        // skips its backfill and never touches the failing transport.
        let covered: BTreeSet<NoteTag> =
            std::iter::once(NoteTag::with_account_target(account.id())).collect();
        client
            .miden_client
            .set_setting(NOTE_TRANSPORT_COVERED_TAGS_KEY.to_string(), covered.clone())
            .await
            .unwrap();

        let report = client.drain_private_note_backlog().await.unwrap();
        assert_ne!(report.status, TransportRecoveryStatus::Completed);
        assert_eq!(report.imported, 0);

        let restored: BTreeSet<NoteTag> = client
            .miden_client
            .get_setting(NOTE_TRANSPORT_COVERED_TAGS_KEY.to_string())
            .await
            .unwrap()
            .expect("covered set must be restored after a failed drain");
        assert!(
            restored.is_superset(&covered),
            "pre-drain covered tags must survive a failed drain: {restored:?}"
        );
    }

    #[tokio::test]
    async fn drain_never_regresses_an_advanced_transport_cursor() {
        use miden_client::note_transport::NOTE_TRANSPORT_CURSOR_STORE_SETTING;

        let dir = tempfile::tempdir().unwrap();
        let (node, api) = mock_transport();
        let mut client = offline_client(dir.path(), Some(api)).await;

        let account = test_wallet(3);
        client.add_or_update_account(&account, false).await.unwrap();
        add_to_transport(&node, private_note_for(&account, 1));

        // Simulate a store whose cursor another account's sync has already
        // advanced far past this backlog. The store encodes the cursor as raw
        // big-endian u64 bytes.
        let advanced: u64 = u64::MAX;
        client
            .miden_client
            .set_setting(
                NOTE_TRANSPORT_CURSOR_STORE_SETTING.to_string(),
                advanced.to_be_bytes(),
            )
            .await
            .unwrap();

        // The drain ignores the stored cursor for scanning (the backlogged
        // note is still recovered)...
        let report = client.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Completed);
        assert_eq!(report.imported, 1);

        // ...but persists max(drain_cursor, stored_cursor): the advanced
        // cursor survives.
        let stored: [u8; 8] = client
            .miden_client
            .get_setting(NOTE_TRANSPORT_CURSOR_STORE_SETTING.to_string())
            .await
            .unwrap()
            .expect("cursor setting exists");
        assert_eq!(u64::from_be_bytes(stored), advanced);
    }

    /// Transport that serves a limited number of successful fetches, then
    /// fails with a connection-shaped error — a connection dropped mid-drain.
    struct InterruptibleTransport {
        inner: MockNoteTransportApi,
        fetches_before_failure: std::sync::atomic::AtomicUsize,
    }

    #[async_trait::async_trait]
    impl NoteTransportClient for InterruptibleTransport {
        async fn send_note(
            &self,
            header: miden_protocol::note::NoteHeader,
            details: Vec<u8>,
        ) -> std::result::Result<(), NoteTransportError> {
            self.inner.send_note(header, details);
            Ok(())
        }

        async fn fetch_notes(
            &self,
            tags: &[NoteTag],
            cursor: miden_client::note_transport::NoteTransportCursor,
        ) -> std::result::Result<
            (
                Vec<miden_client::note_transport::NoteInfo>,
                miden_client::note_transport::NoteTransportCursor,
            ),
            NoteTransportError,
        > {
            let allowed = self
                .fetches_before_failure
                .fetch_update(
                    std::sync::atomic::Ordering::SeqCst,
                    std::sync::atomic::Ordering::SeqCst,
                    |n| n.checked_sub(1),
                )
                .is_ok();
            if !allowed {
                return Err(NoteTransportError::Network(
                    "connection dropped mid-drain".to_string(),
                ));
            }
            Ok(self.inner.fetch_notes(tags, cursor))
        }

        async fn stream_notes(
            &self,
            _tag: NoteTag,
            _cursor: miden_client::note_transport::NoteTransportCursor,
        ) -> std::result::Result<
            Box<dyn miden_client::note_transport::NoteStream>,
            NoteTransportError,
        > {
            Ok(Box::new(
                miden_client::testing::note_transport::DummyNoteStream {},
            ))
        }
    }

    /// A connection lost mid-drain after partial progress must not report
    /// `Unavailable` ("nothing was imported"): it is an interrupted,
    /// retryable drain and the partial count is kept.
    #[tokio::test]
    async fn interrupted_drain_with_partial_progress_reports_a_retryable_failure() {
        let dir = tempfile::tempdir().unwrap();
        // Cap each response at one note so the two-note backlog needs two
        // fetches; the transport dies after the first.
        let node = Arc::new(RwLock::new(MockNoteTransportNode::with_max_batch(1)));
        let api: Arc<dyn NoteTransportClient> = Arc::new(InterruptibleTransport {
            inner: MockNoteTransportApi::new(node.clone()),
            fetches_before_failure: std::sync::atomic::AtomicUsize::new(1),
        });
        let mut client = offline_client(dir.path(), Some(api)).await;

        let account = test_wallet(6);
        client.add_or_update_account(&account, false).await.unwrap();
        add_to_transport(&node, private_note_for(&account, 1));
        // Distinct cursor values require distinct insertion timestamps.
        std::thread::sleep(std::time::Duration::from_millis(2));
        add_to_transport(&node, private_note_for(&account, 2));

        let report = client.drain_private_note_backlog().await.unwrap();

        assert_eq!(report.status, TransportRecoveryStatus::Failed);
        assert!(report.retryable);
        assert_eq!(report.imported, 1);
        assert!(report.reason.unwrap().contains("connection dropped"));
    }

    // ---------------------------------------------------------------------
    // load/tag regression suite: the store-side behavior pull_account relies
    // on, which gates whether a drain sees the recovered account at all.
    // ---------------------------------------------------------------------

    #[tokio::test]
    async fn adding_an_account_tracks_its_standard_note_tag() {
        let dir = tempfile::tempdir().unwrap();
        let mut client = offline_client(dir.path(), None).await;

        let account = test_wallet(4);
        client.add_or_update_account(&account, false).await.unwrap();

        let expected = NoteTag::with_account_target(account.id());
        let tags = client.miden_client.get_note_tags().await.unwrap();
        assert!(
            tags.iter().any(|record| record.tag == expected),
            "the recovered account's standard note tag must be tracked after insert"
        );
    }

    #[tokio::test]
    async fn re_adding_an_existing_account_is_idempotent() {
        let dir = tempfile::tempdir().unwrap();
        let mut client = offline_client(dir.path(), None).await;

        let account = test_wallet(5);
        client.add_or_update_account(&account, false).await.unwrap();
        // Reload path: pull_account calls add_or_update_account(_, true) on
        // an account that is already present.
        client.add_or_update_account(&account, true).await.unwrap();

        let expected = NoteTag::with_account_target(account.id());
        let tags = client.miden_client.get_note_tags().await.unwrap();
        assert_eq!(
            tags.iter().filter(|record| record.tag == expected).count(),
            1,
            "reload must not duplicate or drop the account's note tag"
        );
        assert!(
            client
                .miden_client
                .get_account(account.id())
                .await
                .unwrap()
                .is_some()
        );
    }

    // ---------------------------------------------------------------------
    // live smoke (real testnet transport) — run explicitly:
    //   cargo test -p miden-multisig-client --lib live_testnet -- --ignored
    // ---------------------------------------------------------------------

    async fn live_client(dir: &Path, with_transport: bool) -> MultisigClient {
        use miden_client::note_transport::NOTE_TRANSPORT_TESTNET_ENDPOINT;

        let mut builder = MultisigClient::builder()
            .miden_endpoint(Endpoint::testnet())
            .guardian_endpoint("http://localhost:1")
            .account_dir(dir)
            .generate_key();
        if with_transport {
            builder = builder.note_transport_endpoint(NOTE_TRANSPORT_TESTNET_ENDPOINT);
        }
        builder.build().await.expect("live client builds")
    }

    /// The full device-loss round trip over the real testnet transport (the
    /// scenario spike #412 validated): relay a private note, recover it into
    /// a fresh store via the drain, and check idempotence plus the
    /// disabled-transport report. Network-dependent, hence ignored in CI.
    #[tokio::test]
    #[ignore = "requires network access to the Miden testnet"]
    async fn live_testnet_transport_drain_round_trip() {
        use miden_protocol::address::Address;

        let dir = tempfile::tempdir().unwrap();
        let account = test_wallet(9);

        // "Old device": relay a private note addressed at the account.
        // Transport delivery needs no on-chain transaction.
        let mut sender = live_client(&dir.path().join("sender"), true).await;
        sender
            .miden_client
            .send_private_note_with_block_hint(
                private_note_for(&account, 77),
                &Address::new(account.id()),
                // The note never commits on the mock chain, so any
                // at-or-below-commitment hint is valid.
                miden_protocol::block::BlockNumber::from(0u32),
            )
            .await
            .expect("transport send succeeds");

        // "New device": fresh store; pulling the account tracks its tag.
        let mut recovered = live_client(&dir.path().join("recovered"), true).await;
        recovered
            .add_or_update_account(&account, false)
            .await
            .unwrap();

        let report = recovered.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Completed);
        assert!(
            report.imported >= 1,
            "the relayed note must be recovered, got {report:?}"
        );

        let report = recovered.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Completed);
        assert_eq!(report.imported, 0, "re-drain must be idempotent");

        // A custom node endpoint derives no transport service (the testnet
        // preset keeps the upstream default transport), so the drain reports
        // Unavailable without touching the network.
        let mut no_transport = MultisigClient::builder()
            .miden_endpoint(Endpoint::new(
                "http".to_string(),
                "node".to_string(),
                Some(1),
            ))
            .guardian_endpoint("http://localhost:1")
            .account_dir(dir.path().join("no-transport"))
            .generate_key()
            .build()
            .await
            .expect("custom-endpoint client builds");
        let report = no_transport.drain_private_note_backlog().await.unwrap();
        assert_eq!(report.status, TransportRecoveryStatus::Unavailable);
        assert!(!report.retryable);
    }
}