openrtc/

native_node.rs

#![cfg(not(target_arch = "wasm32"))]

use anyhow::Result;
use futures::StreamExt;
use iroh::{
    endpoint::{Connection, RecvStream, SendStream},
    protocol::{AcceptError, ProtocolHandler, Router},
    Endpoint, EndpointAddr, EndpointId, Watcher,
};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex as StdMutex, Weak};
use std::time::Instant;
use tokio::sync::{broadcast, mpsc, oneshot, RwLock};
use tokio_stream::wrappers::BroadcastStream;

use crate::heartbeat::{
    classify_incoming_uni, parse_incoming_pong, respond_to_ping, HealthTransition, HeartbeatConfig,
    IncomingUniClassification, IrohConnectionProbe, IrohHeartbeatManager, IrohProbeRegistry,
    PrefixedRecvStream,
};
use crate::iroh_connection_policy::{
    decide_inbound_install, decide_outbound_install, should_start_outbound_dial,
    should_wait_for_canonical_inbound, ExistingConnectionState, IrohConnectionDirection,
    IrohConnectionInstallDecision, NONCANONICAL_OUTBOUND_DIAL_GRACE_MS,
};

/// Native Iroh endpoint policy shared by OpenRTC's primary endpoint and any
/// separately-built custom-transport endpoint. Keeping this configuration in
/// the lifecycle owner prevents a BLE/LAN adapter from inheriting Iroh's short
/// default idle timeout and silently dropping an otherwise healthy route.
pub fn native_iroh_transport_config() -> iroh::endpoint::QuicTransportConfig {
    use iroh::endpoint::VarInt;

    let idle_timeout = std::time::Duration::from_secs(120)
        .try_into()
        .expect("120 seconds is a valid Iroh idle timeout");
    let keep_alive_secs = if cfg!(any(target_os = "ios", target_os = "android")) {
        90
    } else {
        25
    };

    iroh::endpoint::QuicTransportConfig::builder()
        .max_idle_timeout(Some(idle_timeout))
        .keep_alive_interval(std::time::Duration::from_secs(keep_alive_secs))
        .stream_receive_window(VarInt::from_u32(8 * 1024 * 1024))
        .receive_window(VarInt::from_u32(16 * 1024 * 1024))
        .send_window(16 * 1024 * 1024)
        .build()
}

const INSTALL_COLLISION_PROBE_TIMEOUT: std::time::Duration =
    std::time::Duration::from_millis(1_000);

pub(crate) fn transport_generation_for_connection(connection: &Connection) -> u64 {
    crate::transport_generation::for_connection(connection)
}

fn forget_transport_generation(connection: &Connection) {
    crate::transport_generation::forget_connection(connection);
}

async fn existing_transport_is_live_for_install(
    endpoint_id: EndpointId,
    previous: &Connection,
    fresh_stable_id: u64,
    fresh_direction: IrohConnectionDirection,
    probe_registry: &IrohProbeRegistry,
) -> bool {
    let previous_stable_id = transport_generation_for_connection(previous);
    if previous_stable_id == fresh_stable_id {
        return previous.close_reason().is_none();
    }
    if previous.close_reason().is_some() {
        return false;
    }

    let responsive = probe_registry
        .probe(
            &endpoint_id.to_string(),
            previous_stable_id,
            previous,
            INSTALL_COLLISION_PROBE_TIMEOUT,
        )
        .await;
    if !responsive {
        let reason = match fresh_direction {
            IrohConnectionDirection::Inbound => {
                crate::lifecycle_reason::REASON_ZOMBIE_REPLACED_BY_FRESH_ACCEPT
            }
            IrohConnectionDirection::Outbound => {
                crate::lifecycle_reason::REASON_ZOMBIE_REPLACED_BY_FRESH_DIAL
            }
        };
        println!(
            "[OpenRTC][iroh-install] replacing unresponsive transport endpoint_id={} previous_stable_id={} fresh_stable_id={} fresh_direction={:?}",
            endpoint_id,
            previous_stable_id,
            fresh_stable_id,
            fresh_direction,
        );
        previous.close(0u8.into(), reason.as_bytes());
    }
    responsive
}

async fn should_break_accept_loop(connection: &Connection) -> bool {
    if connection.close_reason().is_some() {
        return true;
    }
    tokio::time::sleep(std::time::Duration::from_millis(25)).await;
    connection.close_reason().is_some()
}

fn local_prefers_outbound(local_endpoint_id: EndpointId, remote_endpoint_id: EndpointId) -> bool {
    local_endpoint_id.to_string() > remote_endpoint_id.to_string()
}

async fn outbound_dial_precheck(
    endpoint: &Endpoint,
    endpoint_id: EndpointId,
    connections: &Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: &Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
) -> bool {
    let prefers_outbound = local_prefers_outbound(endpoint.id(), endpoint_id);
    let existing = {
        let conns = connections.read().await;
        if let Some(existing) = conns.get(&endpoint_id) {
            let direction = connection_inserted_at
                .read()
                .await
                .get(&endpoint_id)
                .map(|metadata| metadata.direction)
                .unwrap_or(IrohConnectionDirection::Outbound);
            Some((
                existing.close_reason().is_none(),
                direction,
                selected_custom_transport_id(existing),
            ))
        } else {
            None
        }
    };
    let existing_state = existing.map(|(alive, direction, _)| (alive, direction));
    let existing_custom_transport_id =
        existing.and_then(|(_, _, custom_transport_id)| custom_transport_id);

    if !should_start_outbound_dial_with_transport_preference(
        existing_state,
        existing_custom_transport_id,
        prefers_outbound,
    ) {
        return true;
    }
    if !should_wait_for_canonical_inbound(existing_state, prefers_outbound) {
        return false;
    }

    tokio::time::sleep(std::time::Duration::from_millis(
        NONCANONICAL_OUTBOUND_DIAL_GRACE_MS,
    ))
    .await;
    connections
        .read()
        .await
        .get(&endpoint_id)
        .is_some_and(|connection| connection.close_reason().is_none())
}

fn manual_disconnect_notice_key(endpoint_id: EndpointId, transport_stable_id: u64) -> String {
    format!("{endpoint_id}:{transport_stable_id}")
}

fn heartbeat_connection_key(endpoint_id: EndpointId, transport_stable_id: u64) -> String {
    format!("{endpoint_id}:{transport_stable_id}")
}

#[derive(Debug, Clone, Copy)]
struct IrohConnectionMetadata {
    inserted_at: Instant,
    direction: IrohConnectionDirection,
}

#[derive(Debug)]
pub enum IncomingStreamType {
    Bi(SendStream, RecvStream),
    Uni(PrefixedRecvStream),
}

#[derive(Debug)]
pub struct IncomingStream {
    pub endpoint_id: EndpointId,
    /// Physical Iroh connection generation that accepted this stream. This is
    /// distinct from the logical peer/session id and must follow the stream
    /// through admission so stale accept loops cannot replace current owners.
    pub transport_stable_id: u64,
    /// Wire bytes inspected by the Rust admission router and preserved for the
    /// eventual application-crypto decoder. Empty for ordinary streams.
    pub recv_prefix: Vec<u8>,
    pub stream: IncomingStreamType,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ConnectEvent {
    Connected,
    Closed { error: Option<String> },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ReplacementConnectEvent {
    Connected { transport_stable_id: u64 },
    Closed { error: Option<String> },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum AcceptEvent {
    Accepted {
        endpoint_id: EndpointId,
        transport_stable_id: u64,
        replaced_transport_stable_id: Option<u64>,
    },
    Closed {
        endpoint_id: EndpointId,
        transport_stable_id: u64,
        error: Option<String>,
        /// True when this runtime dialed the physical leg. External-router
        /// integrations still need OpenRTC to reconcile closes for outbound
        /// legs, while their host-owned inbound handler performs its own
        /// exact-generation reconciliation.
        was_outbound: bool,
        /// True when the local side initiated the close (e.g. `disconnect()`
        /// after auth rejection). Mirrors the wasm32 variant so consumers can
        /// short-circuit replacement-wait logic that only makes sense for
        /// peer-initiated or transport-failure closes.
        was_locally_closed: bool,
    },
}

/// Immediate result of the native node's atomic physical-connection arbitration.
/// Higher lifecycle layers must only promote `Installed`; `KeptExisting` means
/// the fresh connection lost and the existing stable ID remains authoritative.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ExternalConnectionInstallOutcome {
    Installed {
        transport_stable_id: u64,
        replaced_transport_stable_id: Option<u64>,
    },
    KeptExisting {
        fresh_transport_stable_id: u64,
        kept_transport_stable_id: u64,
    },
}

#[derive(Debug, Clone, Copy)]
struct InboundReplacementAuthorization {
    transport_id: u64,
    expires_at: Instant,
}

type InboundReplacementAuthorizations =
    Arc<RwLock<HashMap<EndpointId, InboundReplacementAuthorization>>>;

type OutboundDialGate = tokio::sync::Mutex<()>;
type OutboundDialGates = Arc<StdMutex<HashMap<EndpointId, Weak<OutboundDialGate>>>>;

struct OutboundDialGateLease {
    endpoint_id: EndpointId,
    gate: Arc<OutboundDialGate>,
    gates: OutboundDialGates,
}

impl Drop for OutboundDialGateLease {
    fn drop(&mut self) {
        let mut gates = self
            .gates
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        let is_current = gates
            .get(&self.endpoint_id)
            .is_some_and(|current| current.ptr_eq(&Arc::downgrade(&self.gate)));
        // The registry holds a Weak reference, so one strong reference means
        // this is the last active/waiting lease. Remove its key immediately;
        // cancelled dials follow the same Drop path.
        if is_current && Arc::strong_count(&self.gate) == 1 {
            gates.remove(&self.endpoint_id);
        }
    }
}

fn outbound_dial_gate(
    endpoint_id: EndpointId,
    registry: &OutboundDialGates,
) -> OutboundDialGateLease {
    let mut gates = registry
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    // The registry owns only weak references. Pruning on every acquisition
    // means completed or cancelled peer dials cannot leave an ever-growing
    // per-endpoint map behind.
    gates.retain(|_, gate| gate.strong_count() > 0);
    let gate = gates
        .get(&endpoint_id)
        .and_then(Weak::upgrade)
        .unwrap_or_else(|| {
            let gate = Arc::new(OutboundDialGate::new(()));
            gates.insert(endpoint_id, Arc::downgrade(&gate));
            gate
        });
    drop(gates);
    OutboundDialGateLease {
        endpoint_id,
        gate,
        gates: Arc::clone(registry),
    }
}

fn selected_custom_transport_id(connection: &Connection) -> Option<u64> {
    connection.paths().iter().find_map(|path| {
        if !path.is_selected() {
            return None;
        }
        match path.remote_addr() {
            iroh::TransportAddr::Custom(addr) => Some(addr.id()),
            _ => None,
        }
    })
}

fn custom_transport_arbitration_override(
    existing_alive: bool,
    existing_custom_transport_id: Option<u64>,
    fresh_custom_transport_id: Option<u64>,
    authorized_replacement_transport_id: Option<u64>,
) -> Option<IrohConnectionInstallDecision> {
    if authorized_replacement_transport_id
        .is_some_and(|expected| fresh_custom_transport_id == Some(expected))
    {
        return Some(IrohConnectionInstallDecision::ReplaceExisting {
            close_existing_reason: crate::lifecycle_reason::REASON_NATIVE_CUSTOM_TRANSPORT_UPGRADE,
        });
    }
    if existing_alive
        && existing_custom_transport_id.is_some()
        && fresh_custom_transport_id != existing_custom_transport_id
    {
        return Some(IrohConnectionInstallDecision::KeepExisting {
            close_fresh_reason: "existing-custom-transport-preferred",
        });
    }
    None
}

fn should_start_outbound_dial_with_transport_preference(
    existing: Option<(bool, IrohConnectionDirection)>,
    existing_custom_transport_id: Option<u64>,
    local_prefers_outbound: bool,
) -> bool {
    if existing.is_some_and(|(alive, _)| alive) && existing_custom_transport_id.is_some() {
        return false;
    }
    should_start_outbound_dial(existing, local_prefers_outbound)
}

fn decide_outbound_install_with_transport_preference(
    existing: Option<ExistingConnectionState>,
    existing_custom_transport_id: Option<u64>,
    fresh_custom_transport_id: Option<u64>,
    local_prefers_outbound: bool,
) -> IrohConnectionInstallDecision {
    custom_transport_arbitration_override(
        existing.as_ref().is_some_and(|state| state.alive),
        existing_custom_transport_id,
        fresh_custom_transport_id,
        None,
    )
    .unwrap_or_else(|| decide_outbound_install(existing, local_prefers_outbound))
}

async fn consume_authorized_inbound_replacement(
    connection: &Connection,
    authorizations: &InboundReplacementAuthorizations,
) -> Option<u64> {
    let endpoint_id = connection.remote_id();
    let deadline = Instant::now() + std::time::Duration::from_secs(1);
    loop {
        let authorization = authorizations.read().await.get(&endpoint_id).copied();
        let Some(authorization) = authorization else {
            return None;
        };
        if Instant::now() >= authorization.expires_at {
            let mut authorizations = authorizations.write().await;
            if authorizations.get(&endpoint_id).is_some_and(|current| {
                current.transport_id == authorization.transport_id
                    && current.expires_at == authorization.expires_at
            }) {
                authorizations.remove(&endpoint_id);
            }
            return None;
        }
        if selected_custom_transport_id(connection) == Some(authorization.transport_id) {
            let mut authorizations = authorizations.write().await;
            if authorizations.get(&endpoint_id).is_some_and(|current| {
                current.transport_id == authorization.transport_id
                    && current.expires_at == authorization.expires_at
            }) {
                authorizations.remove(&endpoint_id);
                return Some(authorization.transport_id);
            }
            return None;
        }
        if Instant::now() >= deadline || connection.close_reason().is_some() {
            return None;
        }
        tokio::time::sleep(std::time::Duration::from_millis(25)).await;
    }
}

#[derive(Debug, Clone)]
pub struct PlutoniumProtocol {
    event_sender: broadcast::Sender<AcceptEvent>,
    stream_sender: async_channel::Sender<IncomingStream>,
    connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
    local_endpoint_id: EndpointId,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    probe_registry: IrohProbeRegistry,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
    inbound_replacement_authorizations: InboundReplacementAuthorizations,
}

async fn run_connection_loop(
    connection: Connection,
    event_sender: broadcast::Sender<AcceptEvent>,
    stream_sender: async_channel::Sender<IncomingStream>,
    connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
    local_endpoint_id: EndpointId,
    connection_direction: IrohConnectionDirection,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    probe_registry: IrohProbeRegistry,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
    replacement_transport_id: Option<u64>,
    mut install_outcome_sender: Option<oneshot::Sender<ExternalConnectionInstallOutcome>>,
) -> std::result::Result<(), AcceptError> {
    let endpoint_id = connection.remote_id();
    let raw_stable_id = connection.stable_id();
    let stable_id = transport_generation_for_connection(&connection);
    let endpoint_key = endpoint_id.to_string();
    let heartbeat_key = heartbeat_connection_key(endpoint_id, stable_id);

    let mut kept_existing = None;
    let mut replaced_transport_stable_id = None;
    {
        let mut conns = connections.write().await;
        let mut insert_times = connection_inserted_at.write().await;
        if let Some(previous) = conns.get(&endpoint_id).cloned() {
            let fresh_custom_transport_id = selected_custom_transport_id(&connection);
            let existing_custom_transport_id = selected_custom_transport_id(&previous);
            let forced_replacement = replacement_transport_id
                .is_some_and(|expected| fresh_custom_transport_id == Some(expected));
            let previous_alive = if forced_replacement {
                previous.close_reason().is_none()
            } else {
                existing_transport_is_live_for_install(
                    endpoint_id,
                    &previous,
                    stable_id,
                    connection_direction,
                    &probe_registry,
                )
                .await
            };
            let previous_direction = insert_times
                .get(&endpoint_id)
                .map(|metadata| metadata.direction)
                .unwrap_or(connection_direction);
            let previous_age_ms = insert_times
                .get(&endpoint_id)
                .map(|metadata| metadata.inserted_at.elapsed().as_millis() as u64)
                .unwrap_or(0);
            let existing = Some(ExistingConnectionState {
                same_stable_id: previous.stable_id() == raw_stable_id,
                alive: previous_alive,
                direction: previous_direction,
                age_ms: previous_age_ms,
            });
            let prefers_outbound = local_prefers_outbound(local_endpoint_id, endpoint_id);
            let decision = if forced_replacement && previous.stable_id() == raw_stable_id {
                IrohConnectionInstallDecision::Install
            } else if let Some(decision) = custom_transport_arbitration_override(
                previous_alive,
                existing_custom_transport_id,
                fresh_custom_transport_id,
                replacement_transport_id,
            ) {
                decision
            } else {
                match connection_direction {
                    IrohConnectionDirection::Inbound => {
                        decide_inbound_install(existing, prefers_outbound)
                    }
                    IrohConnectionDirection::Outbound => {
                        decide_outbound_install(existing, prefers_outbound)
                    }
                }
            };
            match decision {
                IrohConnectionInstallDecision::Install => {
                    let previous_stable_id = transport_generation_for_connection(&previous);
                    if previous_stable_id != stable_id {
                        replaced_transport_stable_id = Some(previous_stable_id);
                    }
                }
                IrohConnectionInstallDecision::ReplaceExisting {
                    close_existing_reason,
                } => {
                    replaced_transport_stable_id =
                        Some(transport_generation_for_connection(&previous));
                    previous.close(0u8.into(), close_existing_reason.as_bytes());
                }
                IrohConnectionInstallDecision::KeepExisting { close_fresh_reason } => {
                    kept_existing = Some(ExternalConnectionInstallOutcome::KeptExisting {
                        fresh_transport_stable_id: stable_id,
                        kept_transport_stable_id: transport_generation_for_connection(&previous),
                    });
                    connection.close(0u8.into(), close_fresh_reason.as_bytes());
                }
            }
        }
        if kept_existing.is_none() {
            conns.insert(endpoint_id, connection.clone());
            insert_times.insert(
                endpoint_id,
                IrohConnectionMetadata {
                    inserted_at: Instant::now(),
                    direction: connection_direction,
                },
            );
        }
    }

    if let Some(outcome) = kept_existing {
        if let Some(sender) = install_outcome_sender.take() {
            let _ = sender.send(outcome);
        }
        forget_transport_generation(&connection);
        return Ok(());
    }

    if let Some(sender) = install_outcome_sender.take() {
        let _ = sender.send(ExternalConnectionInstallOutcome::Installed {
            transport_stable_id: stable_id,
            replaced_transport_stable_id,
        });
    }

    // Start iroh heartbeat for this connection if the manager is configured.
    if let (Some(mgr), Some(tx)) = (&heartbeat_manager, &heartbeat_health_tx) {
        mgr.start_connection(
            heartbeat_key.clone(),
            connection.clone(),
            heartbeat_config.clone(),
            tx.clone(),
        )
        .await;
    }

    event_sender
        .send(AcceptEvent::Accepted {
            endpoint_id,
            transport_stable_id: stable_id,
            replaced_transport_stable_id,
        })
        .ok();

    loop {
        tokio::select! {
            // Application bi streams and control/application uni streams share
            // this physical connection. A biased bi-first dispatcher can
            // indefinitely starve heartbeat pings under sustained application
            // traffic, causing the health owner to tear down a live transport.
            // Tokio's default fair selection keeps both protocol classes
            // progressing without introducing another liveness owner.
            res = connection.accept_bi() => {
                match res {
                    Ok((send, recv)) => {
                        probe_registry
                            .record_inbound_activity(&endpoint_key, stable_id)
                            .await;
                        let _ = stream_sender.send(IncomingStream {
                            endpoint_id,
                            transport_stable_id: stable_id,
                            recv_prefix: Vec::new(),
                            stream: IncomingStreamType::Bi(send, recv),
                        }).await;
                    }
                    Err(_) => {
                        if should_break_accept_loop(&connection).await {
                            break;
                        }
                    },
                }
            }
            res = connection.accept_uni() => {
                match res {
                    Ok(recv) => {
                        probe_registry
                            .record_inbound_activity(&endpoint_key, stable_id)
                            .await;
                        // Classification runs independently so a quiet application
                        // uni stream cannot stall acceptance of later streams.
                        let conn_clone = connection.clone();
                        let mgr_clone = heartbeat_manager.clone();
                        let probe_registry = probe_registry.clone();
                        let endpoint_key = endpoint_key.clone();
                        let heartbeat_key = heartbeat_key.clone();
                        let transport_stable_id = stable_id;
                        let notice_key = manual_disconnect_notice_key(
                            endpoint_id,
                            transport_stable_id,
                        );
                        let notices = manual_disconnect_notices.clone();
                        let application_streams = stream_sender.clone();
                        tokio::spawn(async move {
                            match classify_incoming_uni(recv).await {
                                IncomingUniClassification::Control { type_id, payload } => {
                                use crate::heartbeat::codec;
                                if type_id == codec::TYPE_PING {
                                    if !respond_to_ping(&conn_clone, &payload).await {
                                        eprintln!(
                                            "[OpenRTC][iroh-probe] failed to send pong endpoint_id={} transport_stable_id={}",
                                            endpoint_key,
                                            transport_stable_id,
                                        );
                                    }
                                } else if type_id == codec::TYPE_PONG {
                                    if let Some(pong) = parse_incoming_pong(&payload) {
                                        probe_registry
                                            .deliver_pong(
                                                &endpoint_key,
                                                transport_stable_id,
                                                &pong,
                                            )
                                            .await;
                                        if let Some(manager) = mgr_clone {
                                            manager.deliver_pong(&heartbeat_key, pong).await;
                                        }
                                    }
                                } else if type_id == codec::TYPE_MANUAL_DISCONNECT {
                                    notices.write().await.insert(notice_key);
                                    conn_clone.close(
                                        0u8.into(),
                                        crate::lifecycle_reason::REASON_MANUAL_DISCONNECT
                                            .as_bytes(),
                                    );
                                }
                                }
                                IncomingUniClassification::Application(recv) => {
                                    let _ = application_streams
                                        .send(IncomingStream {
                                            endpoint_id,
                                            transport_stable_id,
                                            recv_prefix: Vec::new(),
                                            stream: IncomingStreamType::Uni(recv),
                                        })
                                        .await;
                                }
                                IncomingUniClassification::MalformedControl => {}
                            }
                        });
                    }
                    Err(_) => {
                        if should_break_accept_loop(&connection).await {
                            break;
                        }
                    },
                }
            }
            _ = connection.closed() => break,
        }
    }

    // Stop heartbeat for this connection.
    if let Some(ref mgr) = heartbeat_manager {
        mgr.stop_connection(&heartbeat_key).await;
    }
    probe_registry
        .forget_transport(&endpoint_key, stable_id)
        .await;

    let close_reason = connection.close_reason();
    let close_reason_debug = close_reason.as_ref().map(|reason| format!("{:?}", reason));
    let was_locally_closed = matches!(
        close_reason,
        Some(iroh::endpoint::ConnectionError::LocallyClosed)
    );
    let current_transport_stable_id = connections
        .read()
        .await
        .get(&endpoint_id)
        .map(transport_generation_for_connection);
    eprintln!(
        "[OpenRTC][iroh-connection][closed] endpoint_id={} transport_stable_id={} direction={:?} close_reason={:?} current_transport_stable_id={:?}",
        endpoint_id,
        stable_id,
        connection_direction,
        close_reason_debug,
        current_transport_stable_id,
    );
    event_sender
        .send(AcceptEvent::Closed {
            endpoint_id,
            transport_stable_id: stable_id,
            error: close_reason_debug,
            was_outbound: matches!(connection_direction, IrohConnectionDirection::Outbound),
            was_locally_closed,
        })
        .ok();

    {
        let mut conns = connections.write().await;
        let should_remove = conns
            .get(&endpoint_id)
            .map(|current| current.stable_id() == raw_stable_id)
            .unwrap_or(false);
        if should_remove {
            conns.remove(&endpoint_id);
            connection_inserted_at.write().await.remove(&endpoint_id);
        }
    }

    forget_transport_generation(&connection);
    Ok(())
}

impl PlutoniumProtocol {
    pub const ALPN: &[u8] = b"plutonium/p2p/1";

    fn new(
        event_sender: broadcast::Sender<AcceptEvent>,
        stream_sender: async_channel::Sender<IncomingStream>,
        connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
        connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
        local_endpoint_id: EndpointId,
        manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
        probe_registry: IrohProbeRegistry,
        inbound_replacement_authorizations: InboundReplacementAuthorizations,
    ) -> Self {
        Self {
            event_sender,
            stream_sender,
            connections,
            connection_inserted_at,
            local_endpoint_id,
            manual_disconnect_notices,
            probe_registry,
            heartbeat_manager: None,
            heartbeat_health_tx: None,
            heartbeat_config: HeartbeatConfig::default(),
            inbound_replacement_authorizations,
        }
    }

    pub fn with_heartbeat(
        mut self,
        manager: IrohHeartbeatManager,
        health_tx: mpsc::Sender<HealthTransition>,
        config: HeartbeatConfig,
    ) -> Self {
        self.heartbeat_manager = Some(manager);
        self.heartbeat_health_tx = Some(health_tx);
        self.heartbeat_config = config;
        self
    }

    async fn handle_connection(
        self,
        connection: Connection,
    ) -> std::result::Result<(), AcceptError> {
        let replacement_transport_id = consume_authorized_inbound_replacement(
            &connection,
            &self.inbound_replacement_authorizations,
        )
        .await;
        run_connection_loop(
            connection,
            self.event_sender.clone(),
            self.stream_sender.clone(),
            self.connections.clone(),
            self.connection_inserted_at.clone(),
            self.local_endpoint_id,
            IrohConnectionDirection::Inbound,
            self.manual_disconnect_notices.clone(),
            self.probe_registry.clone(),
            self.heartbeat_manager.clone(),
            self.heartbeat_health_tx.clone(),
            self.heartbeat_config.clone(),
            replacement_transport_id,
            None,
        )
        .await
    }
}

impl ProtocolHandler for PlutoniumProtocol {
    #[allow(refining_impl_trait)]
    fn accept(
        &self,
        connection: Connection,
    ) -> impl n0_future::Future<Output = std::result::Result<(), AcceptError>> + std::marker::Send
    {
        let proto = self.clone();
        async move { proto.handle_connection(connection).await }
    }
}

async fn retire_unobserved_outbound_connection(
    endpoint_id: EndpointId,
    connection: &Connection,
    connections: &Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: &Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
) {
    connection.close(0u8.into(), b"connect-event-receiver-closed");
    let raw_stable_id = connection.stable_id();
    let mut conns = connections.write().await;
    if conns
        .get(&endpoint_id)
        .is_some_and(|current| current.stable_id() == raw_stable_id)
    {
        conns.remove(&endpoint_id);
        connection_inserted_at.write().await.remove(&endpoint_id);
    }
    forget_transport_generation(connection);
}

async fn connect(
    endpoint: &Endpoint,
    endpoint_id: EndpointId,
    event_sender: async_channel::Sender<ConnectEvent>,
    outbound_dial_gates: OutboundDialGates,
    connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
    stream_sender: async_channel::Sender<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    probe_registry: IrohProbeRegistry,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
) -> Result<()> {
    let connection = {
        // Only the physical Iroh owner can make the precheck + dial + install
        // decision atomic. The gate is per endpoint, so unrelated peer dials
        // remain fully parallel.
        let dial_gate = outbound_dial_gate(endpoint_id, &outbound_dial_gates);
        let _dial_guard = dial_gate.gate.lock().await;
        if outbound_dial_precheck(endpoint, endpoint_id, &connections, &connection_inserted_at)
            .await
        {
            event_sender.send(ConnectEvent::Connected).await?;
            return Ok(());
        }

        let connection = endpoint
            .connect(endpoint_id, PlutoniumProtocol::ALPN)
            .await?;
        let raw_stable_id = connection.stable_id();
        let stable_id = transport_generation_for_connection(&connection);

        {
            let mut conns = connections.write().await;
            let mut insert_times = connection_inserted_at.write().await;
            if let Some(previous) = conns.get(&endpoint_id).cloned() {
                let previous_alive = existing_transport_is_live_for_install(
                    endpoint_id,
                    &previous,
                    stable_id,
                    IrohConnectionDirection::Outbound,
                    &probe_registry,
                )
                .await;
                let previous_direction = insert_times
                    .get(&endpoint_id)
                    .map(|metadata| metadata.direction)
                    .unwrap_or(IrohConnectionDirection::Outbound);
                let previous_age_ms = insert_times
                    .get(&endpoint_id)
                    .map(|metadata| metadata.inserted_at.elapsed().as_millis() as u64)
                    .unwrap_or(0);
                let existing = Some(ExistingConnectionState {
                    same_stable_id: previous.stable_id() == raw_stable_id,
                    alive: previous_alive,
                    direction: previous_direction,
                    age_ms: previous_age_ms,
                });
                match decide_outbound_install_with_transport_preference(
                    existing,
                    selected_custom_transport_id(&previous),
                    selected_custom_transport_id(&connection),
                    local_prefers_outbound(endpoint.id(), endpoint_id),
                ) {
                    IrohConnectionInstallDecision::Install => {}
                    IrohConnectionInstallDecision::ReplaceExisting {
                        close_existing_reason,
                    } => {
                        previous.close(0u8.into(), close_existing_reason.as_bytes());
                    }
                    IrohConnectionInstallDecision::KeepExisting { close_fresh_reason } => {
                        connection.close(0u8.into(), close_fresh_reason.as_bytes());
                        forget_transport_generation(&connection);
                        event_sender.send(ConnectEvent::Connected).await?;
                        let _ = event_sender
                            .send(ConnectEvent::Closed { error: None })
                            .await;
                        return Ok(());
                    }
                }
            }
            conns.insert(endpoint_id, connection.clone());
            insert_times.insert(
                endpoint_id,
                IrohConnectionMetadata {
                    inserted_at: Instant::now(),
                    direction: IrohConnectionDirection::Outbound,
                },
            );
        }
        connection
    };

    if let Err(error) = event_sender.send(ConnectEvent::Connected).await {
        retire_unobserved_outbound_connection(
            endpoint_id,
            &connection,
            &connections,
            &connection_inserted_at,
        )
        .await;
        return Err(error.into());
    }

    // Use the shared loop (handles heartbeat interception for Uni streams).
    let (accept_tx, _) = broadcast::channel(1);
    run_connection_loop(
        connection,
        accept_tx,
        stream_sender,
        connections.clone(),
        connection_inserted_at,
        endpoint.id(),
        IrohConnectionDirection::Outbound,
        manual_disconnect_notices,
        probe_registry,
        heartbeat_manager,
        heartbeat_health_tx,
        heartbeat_config,
        None,
        None,
    )
    .await
    .ok();

    event_sender
        .send(ConnectEvent::Closed { error: None })
        .await?;

    Ok(())
}

async fn connect_addr(
    endpoint: &Endpoint,
    endpoint_id: EndpointId,
    endpoint_addr: EndpointAddr,
    event_sender: async_channel::Sender<ConnectEvent>,
    outbound_dial_gates: OutboundDialGates,
    connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
    stream_sender: async_channel::Sender<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    probe_registry: IrohProbeRegistry,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
) -> Result<()> {
    let connection = {
        let dial_gate = outbound_dial_gate(endpoint_id, &outbound_dial_gates);
        let _dial_guard = dial_gate.gate.lock().await;
        if outbound_dial_precheck(endpoint, endpoint_id, &connections, &connection_inserted_at)
            .await
        {
            event_sender.send(ConnectEvent::Connected).await?;
            return Ok(());
        }

        let connection = endpoint
            .connect(endpoint_addr, PlutoniumProtocol::ALPN)
            .await?;
        let raw_stable_id = connection.stable_id();
        let stable_id = transport_generation_for_connection(&connection);

        {
            let mut conns = connections.write().await;
            let mut insert_times = connection_inserted_at.write().await;
            if let Some(previous) = conns.get(&endpoint_id).cloned() {
                let previous_alive = existing_transport_is_live_for_install(
                    endpoint_id,
                    &previous,
                    stable_id,
                    IrohConnectionDirection::Outbound,
                    &probe_registry,
                )
                .await;
                let previous_direction = insert_times
                    .get(&endpoint_id)
                    .map(|metadata| metadata.direction)
                    .unwrap_or(IrohConnectionDirection::Outbound);
                let previous_age_ms = insert_times
                    .get(&endpoint_id)
                    .map(|metadata| metadata.inserted_at.elapsed().as_millis() as u64)
                    .unwrap_or(0);
                let existing = Some(ExistingConnectionState {
                    same_stable_id: previous.stable_id() == raw_stable_id,
                    alive: previous_alive,
                    direction: previous_direction,
                    age_ms: previous_age_ms,
                });
                match decide_outbound_install_with_transport_preference(
                    existing,
                    selected_custom_transport_id(&previous),
                    selected_custom_transport_id(&connection),
                    local_prefers_outbound(endpoint.id(), endpoint_id),
                ) {
                    IrohConnectionInstallDecision::Install => {}
                    IrohConnectionInstallDecision::ReplaceExisting {
                        close_existing_reason,
                    } => {
                        previous.close(0u8.into(), close_existing_reason.as_bytes());
                    }
                    IrohConnectionInstallDecision::KeepExisting { close_fresh_reason } => {
                        connection.close(0u8.into(), close_fresh_reason.as_bytes());
                        forget_transport_generation(&connection);
                        event_sender.send(ConnectEvent::Connected).await?;
                        let _ = event_sender
                            .send(ConnectEvent::Closed { error: None })
                            .await;
                        return Ok(());
                    }
                }
            }
            conns.insert(endpoint_id, connection.clone());
            insert_times.insert(
                endpoint_id,
                IrohConnectionMetadata {
                    inserted_at: Instant::now(),
                    direction: IrohConnectionDirection::Outbound,
                },
            );
        }
        connection
    };

    if let Err(error) = event_sender.send(ConnectEvent::Connected).await {
        retire_unobserved_outbound_connection(
            endpoint_id,
            &connection,
            &connections,
            &connection_inserted_at,
        )
        .await;
        return Err(error.into());
    }

    let (accept_tx, _) = broadcast::channel(1);
    run_connection_loop(
        connection,
        accept_tx,
        stream_sender,
        connections.clone(),
        connection_inserted_at,
        endpoint.id(),
        IrohConnectionDirection::Outbound,
        manual_disconnect_notices,
        probe_registry,
        heartbeat_manager,
        heartbeat_health_tx,
        heartbeat_config,
        None,
        None,
    )
    .await
    .ok();

    event_sender
        .send(ConnectEvent::Closed { error: None })
        .await?;

    Ok(())
}

async fn install_replacement_connection(
    endpoint: &Endpoint,
    connection: Connection,
    replacement_transport_id: u64,
    event_sender: async_channel::Sender<ReplacementConnectEvent>,
    connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
    stream_sender: async_channel::Sender<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    probe_registry: IrohProbeRegistry,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
) -> Result<()> {
    fn replacement_path_summary(connection: &Connection) -> Vec<String> {
        connection
            .paths()
            .iter()
            .map(|path| {
                format!(
                    "selected={} remote_addr={:?}",
                    path.is_selected(),
                    path.remote_addr()
                )
            })
            .collect()
    }

    eprintln!(
        "[OpenRTC][custom-transport] replacement admission waiting stable_id={} remote_endpoint_id={} expected_transport_id={} paths={:?}",
        connection.stable_id(),
        connection.remote_id(),
        replacement_transport_id,
        replacement_path_summary(&connection)
    );
    let selected_expected_transport = || {
        connection.paths().iter().any(|path| {
            path.is_selected()
                && matches!(
                    path.remote_addr(),
                    iroh::TransportAddr::Custom(addr) if addr.id() == replacement_transport_id
                )
        })
    };
    let selected = tokio::time::timeout(std::time::Duration::from_secs(5), async {
        loop {
            if selected_expected_transport() {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        }
    })
    .await
    .is_ok();
    if !selected {
        eprintln!(
            "[OpenRTC][custom-transport] replacement admission rejected stable_id={} remote_endpoint_id={} expected_transport_id={} paths={:?}",
            connection.stable_id(),
            connection.remote_id(),
            replacement_transport_id,
            replacement_path_summary(&connection)
        );
        connection.close(0u8.into(), b"custom-transport-not-selected");
        anyhow::bail!(
            "replacement connection did not select custom transport id {replacement_transport_id}"
        );
    }

    eprintln!(
        "[OpenRTC][custom-transport] replacement admission accepted stable_id={} remote_endpoint_id={} transport_id={} paths={:?}",
        connection.stable_id(),
        connection.remote_id(),
        replacement_transport_id,
        replacement_path_summary(&connection)
    );

    let (accept_tx, _) = broadcast::channel(1);
    let (install_outcome_sender, install_outcome_receiver) = oneshot::channel();
    let connection_loop = tokio::spawn(run_connection_loop(
        connection,
        accept_tx,
        stream_sender,
        connections,
        connection_inserted_at,
        endpoint.id(),
        IrohConnectionDirection::Outbound,
        manual_disconnect_notices,
        probe_registry,
        heartbeat_manager,
        heartbeat_health_tx,
        heartbeat_config,
        Some(replacement_transport_id),
        Some(install_outcome_sender),
    ));
    let install_outcome = install_outcome_receiver
        .await
        .map_err(|_| anyhow::anyhow!("replacement connection install task ended"))?;
    let transport_stable_id = match install_outcome {
        ExternalConnectionInstallOutcome::Installed {
            transport_stable_id,
            ..
        } => transport_stable_id,
        ExternalConnectionInstallOutcome::KeptExisting {
            fresh_transport_stable_id,
            kept_transport_stable_id,
        } => {
            anyhow::bail!(
                "replacement connection lost arbitration fresh_stable_id={} kept_stable_id={}",
                fresh_transport_stable_id,
                kept_transport_stable_id,
            );
        }
    };
    event_sender
        .send(ReplacementConnectEvent::Connected {
            transport_stable_id,
        })
        .await?;
    let _ = connection_loop.await;

    let _ = event_sender
        .send(ReplacementConnectEvent::Closed { error: None })
        .await;
    Ok(())
}

#[allow(unused_variables)]
async fn connect_replacement_addr(
    endpoint: &Endpoint,
    _endpoint_id: EndpointId,
    endpoint_addr: EndpointAddr,
    replacement_transport_id: u64,
    event_sender: async_channel::Sender<ReplacementConnectEvent>,
    connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
    stream_sender: async_channel::Sender<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    probe_registry: IrohProbeRegistry,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
) -> Result<()> {
    #[cfg(not(openrtc_iroh_preferred_transport_api))]
    anyhow::bail!(
        "custom transport replacement requires a host Iroh build with the preferred-transport API"
    );

    #[cfg(openrtc_iroh_preferred_transport_api)]
    {
        let preferred_transport_addr = endpoint_addr
        .addrs
        .iter()
        .find(|addr| {
            matches!(
                addr,
                iroh::TransportAddr::Custom(custom)
                    if custom.id() == replacement_transport_id
            )
        })
        .cloned()
        .ok_or_else(|| {
            anyhow::anyhow!(
                "replacement address does not contain custom transport id {replacement_transport_id}"
            )
        })?;
        eprintln!(
        "[OpenRTC][custom-transport] replacement dial pinned remote_endpoint_id={} transport_id={} addr={:?}",
        endpoint_addr.id, replacement_transport_id, preferred_transport_addr
    );
        let connection = endpoint
            .connect_with_opts(
                endpoint_addr,
                PlutoniumProtocol::ALPN,
                iroh::endpoint::ConnectOptions::new()
                    .with_preferred_transport_addr(preferred_transport_addr),
            )
            .await?
            .await?;
        install_replacement_connection(
            endpoint,
            connection,
            replacement_transport_id,
            event_sender,
            connections,
            connection_inserted_at,
            stream_sender,
            manual_disconnect_notices,
            probe_registry,
            heartbeat_manager,
            heartbeat_health_tx,
            heartbeat_config,
        )
        .await
    }
}

#[derive(Debug, Clone)]
pub struct IrohNativeNode {
    endpoint: Endpoint,
    // Keep the router alive for endpoints that run pluto-rtc's internal accept loop.
    #[allow(dead_code)]
    router: Option<Router>,
    accept_events: broadcast::Sender<AcceptEvent>,
    connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
    connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
    outbound_dial_gates: OutboundDialGates,
    incoming_streams: async_channel::Sender<IncomingStream>,
    incoming_streams_receiver: async_channel::Receiver<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    probe_registry: IrohProbeRegistry,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
    inbound_replacement_authorizations: InboundReplacementAuthorizations,
}

impl IrohNativeNode {
    pub async fn spawn_with_endpoint(endpoint: Endpoint) -> Result<Self> {
        Self::spawn_with_endpoint_config(endpoint, true).await
    }

    /// Spawn a native node wrapper without creating an internal Router accept loop.
    ///
    /// This mode is used when the embedding application owns the single iroh
    /// Router and forwards accepted plutonium connections into pluto-rtc via
    /// `Client::handle_incoming_connection`.
    pub async fn spawn_with_endpoint_no_router(endpoint: Endpoint) -> Result<Self> {
        Self::spawn_with_endpoint_config(endpoint, false).await
    }

    async fn spawn_with_endpoint_config(endpoint: Endpoint, spawn_router: bool) -> Result<Self> {
        Self::spawn_with_endpoint_config_and_heartbeat(
            endpoint,
            spawn_router,
            None,
            None,
            HeartbeatConfig::default(),
        )
        .await
    }

    /// Spawn with an explicit heartbeat manager for iroh-level liveness monitoring.
    pub async fn spawn_with_heartbeat(
        endpoint: Endpoint,
        heartbeat_manager: IrohHeartbeatManager,
        heartbeat_health_tx: mpsc::Sender<HealthTransition>,
        heartbeat_config: HeartbeatConfig,
    ) -> Result<Self> {
        Self::spawn_with_endpoint_config_and_heartbeat(
            endpoint,
            true,
            Some(heartbeat_manager),
            Some(heartbeat_health_tx),
            heartbeat_config,
        )
        .await
    }

    async fn spawn_with_endpoint_config_and_heartbeat(
        endpoint: Endpoint,
        spawn_router: bool,
        heartbeat_manager: Option<IrohHeartbeatManager>,
        heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
        heartbeat_config: HeartbeatConfig,
    ) -> Result<Self> {
        let (event_sender, _) = broadcast::channel(128);
        let (stream_sender, stream_receiver) = async_channel::bounded(64);
        let connections = Arc::new(RwLock::new(HashMap::new()));
        let connection_inserted_at = Arc::new(RwLock::new(HashMap::new()));
        let outbound_dial_gates = Arc::new(StdMutex::new(HashMap::new()));
        let manual_disconnect_notices = Arc::new(RwLock::new(HashSet::new()));
        let probe_registry = IrohProbeRegistry::new();
        let inbound_replacement_authorizations = Arc::new(RwLock::new(HashMap::new()));

        let router = if spawn_router {
            let proto = PlutoniumProtocol::new(
                event_sender.clone(),
                stream_sender.clone(),
                connections.clone(),
                connection_inserted_at.clone(),
                endpoint.id(),
                manual_disconnect_notices.clone(),
                probe_registry.clone(),
                inbound_replacement_authorizations.clone(),
            );
            let proto = if let (Some(mgr), Some(tx)) =
                (heartbeat_manager.clone(), heartbeat_health_tx.clone())
            {
                proto.with_heartbeat(mgr, tx, heartbeat_config.clone())
            } else {
                proto
            };
            Some(
                Router::builder(endpoint.clone())
                    .accept(PlutoniumProtocol::ALPN, proto)
                    .spawn(),
            )
        } else {
            None
        };

        Ok(Self {
            endpoint,
            router,
            accept_events: event_sender,
            connections,
            connection_inserted_at,
            outbound_dial_gates,
            incoming_streams: stream_sender,
            incoming_streams_receiver: stream_receiver,
            manual_disconnect_notices,
            probe_registry,
            heartbeat_manager,
            heartbeat_health_tx,
            heartbeat_config,
            inbound_replacement_authorizations,
        })
    }

    pub fn endpoint(&self) -> &Endpoint {
        &self.endpoint
    }

    pub async fn node_addr(&self) -> Result<EndpointAddr> {
        // Client initialization owns the bounded relay-online wait and arranges
        // a presence republish when a relay arrives later. Ticket reads must use
        // the latest address snapshot immediately: waiting here again stalls
        // presence and every ticket refresh during relay or DNS outages.
        Ok(self.endpoint.watch_addr().get())
    }

    pub async fn is_connected(&self, endpoint_id: EndpointId) -> bool {
        let conns = self.connections.read().await;
        if let Some(conn) = conns.get(&endpoint_id) {
            conn.close_reason().is_none()
        } else {
            false
        }
    }

    pub fn accept_events(&self) -> futures::stream::BoxStream<'static, AcceptEvent> {
        let receiver = self.accept_events.subscribe();
        Box::pin(
            BroadcastStream::new(receiver).filter_map(|event| futures::future::ready(event.ok())),
        )
    }

    pub fn connect(
        &self,
        endpoint_id: EndpointId,
    ) -> futures::stream::BoxStream<'static, ConnectEvent> {
        let (event_sender, event_receiver) = async_channel::bounded(16);
        let endpoint = self.endpoint.clone();
        let connections = self.connections.clone();
        let connection_inserted_at = self.connection_inserted_at.clone();
        let outbound_dial_gates = self.outbound_dial_gates.clone();
        let stream_sender = self.incoming_streams.clone();
        let hb_mgr = self.heartbeat_manager.clone();
        let hb_tx = self.heartbeat_health_tx.clone();
        let hb_cfg = self.heartbeat_config.clone();
        let manual_notices = self.manual_disconnect_notices.clone();
        let probe_registry = self.probe_registry.clone();

        tokio::spawn(async move {
            let result = connect(
                &endpoint,
                endpoint_id,
                event_sender.clone(),
                outbound_dial_gates,
                connections,
                connection_inserted_at,
                stream_sender,
                manual_notices,
                probe_registry,
                hb_mgr,
                hb_tx,
                hb_cfg,
            )
            .await;

            if let Err(error) = result {
                let _ = event_sender
                    .send(ConnectEvent::Closed {
                        error: Some(error.to_string()),
                    })
                    .await;
            }
        });

        Box::pin(event_receiver)
    }

    pub fn connect_addr(
        &self,
        endpoint_id: EndpointId,
        endpoint_addr: EndpointAddr,
    ) -> futures::stream::BoxStream<'static, ConnectEvent> {
        let (event_sender, event_receiver) = async_channel::bounded(16);
        let endpoint = self.endpoint.clone();
        let connections = self.connections.clone();
        let connection_inserted_at = self.connection_inserted_at.clone();
        let outbound_dial_gates = self.outbound_dial_gates.clone();
        let stream_sender = self.incoming_streams.clone();
        let hb_mgr = self.heartbeat_manager.clone();
        let hb_tx = self.heartbeat_health_tx.clone();
        let hb_cfg = self.heartbeat_config.clone();
        let manual_notices = self.manual_disconnect_notices.clone();
        let probe_registry = self.probe_registry.clone();

        tokio::spawn(async move {
            let result = connect_addr(
                &endpoint,
                endpoint_id,
                endpoint_addr,
                event_sender.clone(),
                outbound_dial_gates,
                connections,
                connection_inserted_at,
                stream_sender,
                manual_notices,
                probe_registry,
                hb_mgr,
                hb_tx,
                hb_cfg,
            )
            .await;

            if let Err(error) = result {
                let _ = event_sender
                    .send(ConnectEvent::Closed {
                        error: Some(error.to_string()),
                    })
                    .await;
            }
        });

        Box::pin(event_receiver)
    }

    /// Establish a fresh Iroh generation through one prepared custom transport.
    /// The existing base generation remains available until the replacement is
    /// authenticated and the requested custom path is selected.
    pub fn connect_replacement_addr(
        &self,
        endpoint_id: EndpointId,
        endpoint_addr: EndpointAddr,
        replacement_transport_id: u64,
    ) -> futures::stream::BoxStream<'static, ReplacementConnectEvent> {
        let (event_sender, event_receiver) = async_channel::bounded(16);
        let endpoint = self.endpoint.clone();
        let connections = self.connections.clone();
        let connection_inserted_at = self.connection_inserted_at.clone();
        let stream_sender = self.incoming_streams.clone();
        let hb_mgr = self.heartbeat_manager.clone();
        let hb_tx = self.heartbeat_health_tx.clone();
        let hb_cfg = self.heartbeat_config.clone();
        let manual_notices = self.manual_disconnect_notices.clone();
        let probe_registry = self.probe_registry.clone();

        tokio::spawn(async move {
            if let Err(error) = connect_replacement_addr(
                &endpoint,
                endpoint_id,
                endpoint_addr,
                replacement_transport_id,
                event_sender.clone(),
                connections,
                connection_inserted_at,
                stream_sender,
                manual_notices,
                probe_registry,
                hb_mgr,
                hb_tx,
                hb_cfg,
            )
            .await
            {
                let _ = event_sender
                    .send(ReplacementConnectEvent::Closed {
                        error: Some(error.to_string()),
                    })
                    .await;
            }
        });

        Box::pin(event_receiver)
    }

    pub async fn disconnect(&self, endpoint_id: EndpointId) -> Result<()> {
        self.disconnect_with_reason(
            endpoint_id,
            crate::lifecycle_reason::REASON_DISCONNECTED_BY_USER,
        )
        .await
    }

    pub async fn disconnect_with_reason(
        &self,
        endpoint_id: EndpointId,
        reason: &str,
    ) -> Result<()> {
        let connection = {
            let mut conns = self.connections.write().await;
            conns.remove(&endpoint_id)
        };

        if let Some(conn) = connection {
            if std::env::var("PLUTO_RTC_TEARDOWN_TRACE").is_ok() {
                eprintln!(
                    "[PlutoRTC][teardown-trace] NativeNode::disconnect endpoint_id={}",
                    endpoint_id
                );
            }
            conn.close(1u8.into(), reason.as_bytes());
        }

        Ok(())
    }

    pub async fn disconnect_with_reason_if_current(
        &self,
        endpoint_id: EndpointId,
        expected_transport_stable_id: u64,
        reason: &str,
    ) -> Result<bool> {
        let connection = {
            let mut connections = self.connections.write().await;
            let is_current = connections.get(&endpoint_id).is_some_and(|connection| {
                transport_generation_for_connection(connection) == expected_transport_stable_id
            });
            if !is_current {
                return Ok(false);
            }
            connections.remove(&endpoint_id)
        };

        if let Some(connection) = connection {
            connection.close(1u8.into(), reason.as_bytes());
            return Ok(true);
        }
        Ok(false)
    }

    pub async fn open_bi(&self, endpoint_id: EndpointId) -> Result<(SendStream, RecvStream)> {
        let (_, send, recv) = self.open_bi_with_transport_stable_id(endpoint_id).await?;
        Ok((send, recv))
    }

    /// Open a stream and return the physical connection generation that owns
    /// it. Reading the id from the same cloned `Connection` avoids labeling a
    /// stream with a replacement that raced into the endpoint map.
    pub async fn open_bi_with_transport_stable_id(
        &self,
        endpoint_id: EndpointId,
    ) -> Result<(u64, SendStream, RecvStream)> {
        let connection = {
            let conns = self.connections.read().await;
            conns.get(&endpoint_id).cloned()
        };

        if let Some(conn) = connection {
            let transport_stable_id = transport_generation_for_connection(&conn);
            let (send, recv) = conn.open_bi().await?;
            Ok((transport_stable_id, send, recv))
        } else {
            Err(anyhow::anyhow!("No active connection to {}", endpoint_id))
        }
    }

    pub async fn open_uni(&self, endpoint_id: EndpointId) -> Result<SendStream> {
        let connection = {
            let conns = self.connections.read().await;
            conns.get(&endpoint_id).cloned()
        };

        if let Some(conn) = connection {
            let send = conn.open_uni().await?;
            Ok(send)
        } else {
            Err(anyhow::anyhow!("No active connection to {}", endpoint_id))
        }
    }

    /// Perform a real remote round-trip on the current physical generation.
    /// A replacement that wins during the probe is preserved and checked on
    /// the next health pass instead of inheriting the retired leg's failure.
    pub async fn probe_connection(
        &self,
        endpoint_id: EndpointId,
        timeout: std::time::Duration,
    ) -> Option<IrohConnectionProbe> {
        let connection = self.connections.read().await.get(&endpoint_id).cloned()?;
        let probed_stable_id = transport_generation_for_connection(&connection);
        let responsive = self
            .probe_registry
            .probe(
                &endpoint_id.to_string(),
                probed_stable_id,
                &connection,
                timeout,
            )
            .await;
        let last_inbound_activity_age = self
            .probe_registry
            .last_inbound_activity_age(&endpoint_id.to_string(), probed_stable_id)
            .await;
        let current_stable_id = self
            .connections
            .read()
            .await
            .get(&endpoint_id)
            .map(transport_generation_for_connection)?;
        if current_stable_id != probed_stable_id {
            return Some(IrohConnectionProbe {
                // No pong was observed on the replacement generation. Return
                // the generation that was actually probed so callers can treat
                // the race as indeterminate and fence any delayed failure.
                transport_stable_id: probed_stable_id,
                responsive: false,
                last_inbound_activity_age,
            });
        }
        Some(IrohConnectionProbe {
            transport_stable_id: probed_stable_id,
            responsive,
            last_inbound_activity_age,
        })
    }

    pub fn incoming_streams_stream(&self) -> async_channel::Receiver<IncomingStream> {
        self.incoming_streams_receiver.clone()
    }

    /// Authorize one protected, peer-scoped inbound custom-transport replacement.
    /// The authorization is consumed only by the exact peer and transport id.
    pub(crate) async fn authorize_inbound_replacement(
        &self,
        endpoint_id: EndpointId,
        transport_id: u64,
        ttl: std::time::Duration,
    ) -> Instant {
        let expires_at = Instant::now() + ttl;
        self.inbound_replacement_authorizations
            .write()
            .await
            .insert(
                endpoint_id,
                InboundReplacementAuthorization {
                    transport_id,
                    expires_at,
                },
            );
        expires_at
    }

    /// Revoke an authorization only when it still names this exact transport.
    /// A delayed cleanup task therefore cannot remove a newer replacement grant.
    pub(crate) async fn revoke_inbound_replacement_if_current(
        &self,
        endpoint_id: EndpointId,
        transport_id: u64,
        expires_at: Instant,
    ) -> bool {
        let mut authorizations = self.inbound_replacement_authorizations.write().await;
        if authorizations.get(&endpoint_id).is_some_and(|current| {
            current.transport_id == transport_id && current.expires_at == expires_at
        }) {
            authorizations.remove(&endpoint_id);
            true
        } else {
            false
        }
    }

    /// Ingest an already-accepted connection into the pluto-rtc native node's
    /// internal connection/event/stream pipeline.
    pub async fn accept_external_connection(&self, connection: Connection) -> Result<()> {
        let replacement_transport_id = consume_authorized_inbound_replacement(
            &connection,
            &self.inbound_replacement_authorizations,
        )
        .await;
        run_connection_loop(
            connection,
            self.accept_events.clone(),
            self.incoming_streams.clone(),
            self.connections.clone(),
            self.connection_inserted_at.clone(),
            self.endpoint.id(),
            IrohConnectionDirection::Inbound,
            self.manual_disconnect_notices.clone(),
            self.probe_registry.clone(),
            self.heartbeat_manager.clone(),
            self.heartbeat_health_tx.clone(),
            self.heartbeat_config.clone(),
            replacement_transport_id,
            None,
        )
        .await
        .map_err(|e| anyhow::anyhow!(e.to_string()))
    }

    /// Ingest an externally accepted connection while reporting the atomic
    /// install decision before the stream loop starts. The caller must keep
    /// polling this future while awaiting `install_outcome_sender`.
    pub(crate) async fn accept_external_connection_with_install_notifier(
        &self,
        connection: Connection,
        install_outcome_sender: oneshot::Sender<ExternalConnectionInstallOutcome>,
    ) -> Result<()> {
        let replacement_transport_id = consume_authorized_inbound_replacement(
            &connection,
            &self.inbound_replacement_authorizations,
        )
        .await;
        run_connection_loop(
            connection,
            self.accept_events.clone(),
            self.incoming_streams.clone(),
            self.connections.clone(),
            self.connection_inserted_at.clone(),
            self.endpoint.id(),
            IrohConnectionDirection::Inbound,
            self.manual_disconnect_notices.clone(),
            self.probe_registry.clone(),
            self.heartbeat_manager.clone(),
            self.heartbeat_health_tx.clone(),
            self.heartbeat_config.clone(),
            replacement_transport_id,
            Some(install_outcome_sender),
        )
        .await
        .map_err(|e| anyhow::anyhow!(e.to_string()))
    }

    /// Get a raw iroh::Connection for a given EndpointId, if one exists.
    /// Used by external protocol handlers (handshake, bucket_sync) to attach
    /// application-level logic on top of pluto-rtc-managed connections.
    pub async fn get_connection(&self, endpoint_id: EndpointId) -> Option<Connection> {
        let conns = self.connections.read().await;
        conns.get(&endpoint_id).cloned()
    }

    /// List all currently active EndpointIds with connections.
    pub async fn active_endpoint_ids(&self) -> Vec<EndpointId> {
        let conns = self.connections.read().await;
        conns.keys().cloned().collect()
    }

    pub async fn take_manual_disconnect_notice(
        &self,
        endpoint_id: EndpointId,
        transport_stable_id: u64,
    ) -> bool {
        self.manual_disconnect_notices
            .write()
            .await
            .remove(&manual_disconnect_notice_key(
                endpoint_id,
                transport_stable_id,
            ))
    }

    /// Observe an explicit peer disconnect without consuming the generation-
    /// scoped notice. Both the generic physical-close watcher and the
    /// incoming-handler lifecycle may observe the same close; only the latter
    /// removes the notice after its terminal transition is fenced.
    pub async fn has_manual_disconnect_notice(
        &self,
        endpoint_id: EndpointId,
        transport_stable_id: u64,
    ) -> bool {
        self.manual_disconnect_notices
            .read()
            .await
            .contains(&manual_disconnect_notice_key(
                endpoint_id,
                transport_stable_id,
            ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use iroh::endpoint::Endpoint;
    use tokio::io::AsyncWriteExt;
    use tokio::time::{sleep, timeout, Duration};

    async fn setup_protocol_endpoint() -> (
        Router,
        PlutoniumProtocol,
        broadcast::Receiver<AcceptEvent>,
        async_channel::Receiver<IncomingStream>,
        Arc<RwLock<HashMap<EndpointId, Connection>>>,
    ) {
        let (event_tx, event_rx) = broadcast::channel(16);
        let (stream_tx, stream_rx) = async_channel::unbounded();
        let connections = Arc::new(RwLock::new(HashMap::new()));
        let connection_inserted_at = Arc::new(RwLock::new(HashMap::new()));

        let endpoint = Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .unwrap();

        let protocol = PlutoniumProtocol::new(
            event_tx.clone(),
            stream_tx.clone(),
            connections.clone(),
            connection_inserted_at,
            endpoint.id(),
            Arc::new(RwLock::new(HashSet::new())),
            IrohProbeRegistry::new(),
            Arc::new(RwLock::new(HashMap::new())),
        );

        let router = Router::builder(endpoint)
            .accept(PlutoniumProtocol::ALPN, Arc::new(protocol.clone()))
            .spawn();

        (router, protocol, event_rx, stream_rx, connections)
    }

    #[test]
    fn manual_disconnect_notice_is_scoped_to_transport_generation() {
        let endpoint_id = iroh::SecretKey::generate().public();
        let old_transport = manual_disconnect_notice_key(endpoint_id, 41);
        let replacement_transport = manual_disconnect_notice_key(endpoint_id, 42);
        let mut notices = HashSet::from([old_transport.clone()]);

        assert_ne!(old_transport, replacement_transport);
        assert!(!notices.remove(&replacement_transport));
        assert!(notices.remove(&old_transport));
    }

    #[test]
    fn authorized_custom_transport_replaces_live_base_connection() {
        assert_eq!(
            custom_transport_arbitration_override(true, None, Some(4344901), Some(4344901)),
            Some(IrohConnectionInstallDecision::ReplaceExisting {
                close_existing_reason:
                    crate::lifecycle_reason::REASON_NATIVE_CUSTOM_TRANSPORT_UPGRADE,
            })
        );
    }

    #[test]
    fn live_custom_transport_wins_against_late_non_custom_duplicate() {
        assert_eq!(
            custom_transport_arbitration_override(true, Some(4344901), None, None),
            Some(IrohConnectionInstallDecision::KeepExisting {
                close_fresh_reason: "existing-custom-transport-preferred",
            })
        );
    }

    #[test]
    fn canonical_outbound_dial_cannot_preempt_live_custom_transport() {
        let existing = Some(ExistingConnectionState {
            same_stable_id: false,
            alive: true,
            direction: IrohConnectionDirection::Inbound,
            age_ms: 0,
        });

        assert!(!should_start_outbound_dial_with_transport_preference(
            Some((true, IrohConnectionDirection::Inbound)),
            Some(4344901),
            true,
        ));
        assert_eq!(
            decide_outbound_install_with_transport_preference(existing, Some(4344901), None, true,),
            IrohConnectionInstallDecision::KeepExisting {
                close_fresh_reason: "existing-custom-transport-preferred",
            },
        );
    }

    #[test]
    fn unauthorized_custom_transport_uses_normal_direction_arbitration() {
        assert_eq!(
            custom_transport_arbitration_override(true, None, Some(4344901), None),
            None
        );
    }

    #[tokio::test]
    async fn delayed_replacement_cleanup_cannot_revoke_a_newer_grant() {
        let endpoint = Endpoint::builder(iroh::endpoint::presets::N0)
            .relay_mode(iroh::RelayMode::Disabled)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("endpoint");
        let node = IrohNativeNode::spawn_with_endpoint_no_router(endpoint)
            .await
            .expect("native node");
        let remote = iroh::SecretKey::generate().public();
        let old_expiry = node
            .authorize_inbound_replacement(remote, 4344901, Duration::from_secs(1))
            .await;
        let new_expiry = node
            .authorize_inbound_replacement(remote, 4344901, Duration::from_secs(2))
            .await;

        assert!(
            !node
                .revoke_inbound_replacement_if_current(remote, 4344901, old_expiry)
                .await
        );
        assert!(
            node.revoke_inbound_replacement_if_current(remote, 4344901, new_expiry)
                .await
        );
    }

    #[tokio::test]
    async fn node_addr_does_not_wait_for_an_unavailable_relay() {
        let endpoint = Endpoint::builder(iroh::endpoint::presets::N0)
            .relay_mode(iroh::RelayMode::Disabled)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind relay-disabled endpoint");
        let expected_id = endpoint.id();
        let node = IrohNativeNode::spawn_with_endpoint(endpoint)
            .await
            .expect("spawn relay-disabled node");

        let address = timeout(Duration::from_secs(1), node.node_addr())
            .await
            .expect("node_addr must not wait for relay readiness")
            .expect("read current endpoint address");

        assert_eq!(address.id, expected_id);
    }

    #[tokio::test]
    async fn test_two_nodes_connect_and_exchange_streams() {
        let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
        let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;

        let ep1 = r1.endpoint();
        let ep2 = r2.endpoint();

        // Exchange addr
        let addr2 = ep2.addr();

        // Node 1 connects to Node 2
        let conn_res = ep1.connect(addr2, PlutoniumProtocol::ALPN).await.unwrap();

        // Wait for connect event on Node 2
        let event = timeout(Duration::from_secs(5), events2.recv())
            .await
            .unwrap()
            .unwrap();
        match event {
            AcceptEvent::Accepted { endpoint_id, .. } => assert_eq!(endpoint_id, ep1.id()),
            _ => panic!("Expected AcceptEvent::Accepted"),
        }

        // Node 1 opens stream
        let (mut send1, _recv1) = conn_res.open_bi().await.unwrap();
        send1.write_all(b"hello node2").await.unwrap();

        // Node 2 receives stream
        let incoming = timeout(Duration::from_secs(5), streams2.recv())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(incoming.endpoint_id, ep1.id());

        let mut recv2 = match incoming.stream {
            IncomingStreamType::Bi(_, r) => r,
            _ => panic!("Expected Bi stream"),
        };

        let mut buf = [0u8; 11];
        recv2.read_exact(&mut buf).await.unwrap();
        assert_eq!(&buf, b"hello node2");
    }

    #[tokio::test]
    async fn retiring_losing_control_candidate_preserves_buffered_verdict() {
        let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
        let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;
        let ep1 = r1.endpoint();
        let ep2 = r2.endpoint();
        let connection = ep1
            .connect(ep2.addr(), PlutoniumProtocol::ALPN)
            .await
            .expect("connect protocol endpoints");
        timeout(Duration::from_secs(5), events2.recv())
            .await
            .expect("host accept timeout")
            .expect("host accept event");

        let (mut dialer_send, mut dialer_recv) =
            connection.open_bi().await.expect("open candidate stream");
        dialer_send
            .write_all(b"candidate")
            .await
            .expect("activate candidate stream");
        dialer_send.flush().await.expect("flush candidate stream");
        let incoming = timeout(Duration::from_secs(5), streams2.recv())
            .await
            .expect("host stream timeout")
            .expect("host stream");
        let IncomingStreamType::Bi(mut host_send, host_recv) = incoming.stream else {
            panic!("expected bidirectional candidate stream");
        };
        let verdict = b"session-token-approved";
        host_send
            .write_all(verdict)
            .await
            .expect("write approval verdict");
        host_send.flush().await.expect("flush approval verdict");

        let (retired, received) = tokio::join!(
            crate::client::retire_losing_native_control_candidate(host_send, host_recv),
            dialer_recv.read_to_end(256),
        );
        retired.expect("retire candidate after peer acknowledges verdict");
        assert_eq!(
            received.expect("dialer reads losing-stream verdict"),
            verdict,
        );
    }

    #[tokio::test]
    async fn retiring_replaced_control_send_preserves_buffered_verdict() {
        let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
        let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;
        let connection = r1
            .endpoint()
            .connect(r2.endpoint().addr(), PlutoniumProtocol::ALPN)
            .await
            .expect("connect protocol endpoints");
        timeout(Duration::from_secs(5), events2.recv())
            .await
            .expect("host accept timeout")
            .expect("host accept event");

        let (mut dialer_send, mut dialer_recv) =
            connection.open_bi().await.expect("open displaced stream");
        dialer_send
            .write_all(b"candidate")
            .await
            .expect("activate displaced stream");
        dialer_send.flush().await.expect("flush displaced stream");
        let incoming = timeout(Duration::from_secs(5), streams2.recv())
            .await
            .expect("host stream timeout")
            .expect("host stream");
        let IncomingStreamType::Bi(mut host_send, _host_recv) = incoming.stream else {
            panic!("expected bidirectional displaced stream");
        };
        let verdict = b"session-token-approved";
        host_send
            .write_all(verdict)
            .await
            .expect("write approval verdict");
        host_send.flush().await.expect("host should flush approval");
        let displaced_send = Arc::new(tokio::sync::Mutex::new(host_send));

        let (retired, received) = tokio::join!(
            crate::client::retire_replaced_native_control_send(displaced_send),
            dialer_recv.read_to_end(256),
        );
        retired.expect("retire replaced stream after peer acknowledges verdict");
        assert_eq!(
            received.expect("dialer reads replaced-stream verdict"),
            verdict,
        );
    }

    #[tokio::test]
    async fn test_healthy_connection_not_replaced() {
        let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
        let (r2, _proto2, mut events2, _streams2, conns2) = setup_protocol_endpoint().await;

        let ep1 = r1.endpoint();
        let ep2 = r2.endpoint();

        let addr2 = ep2.addr();
        let addr1 = ep1.addr();

        // 1. First connection
        let _conn1 = ep1.connect(addr2, PlutoniumProtocol::ALPN).await.unwrap();

        // Wait for connection to be registered in Protocol 2
        let _ = timeout(Duration::from_secs(5), events2.recv())
            .await
            .unwrap()
            .unwrap();

        let active_count = conns2.read().await.len();
        assert_eq!(active_count, 1);

        let original_stable_id = conns2.read().await.get(&ep1.id()).unwrap().stable_id();

        // 2. Dual-dial: Node 2 connects to Node 1 while connection is still healthy
        let _conn2 = ep2.connect(addr1, PlutoniumProtocol::ALPN).await.unwrap();

        // Allow time for the second connection to process
        sleep(Duration::from_millis(100)).await;

        // The original connection should still be intact because it wasn't closed
        let current_conn = conns2.read().await.get(&ep1.id()).unwrap().clone();
        assert_eq!(current_conn.stable_id(), original_stable_id);
    }

    #[tokio::test]
    async fn fresh_connection_replaces_an_unresponsive_same_identity_transport() {
        let (host, _protocol, mut host_events, _streams, host_connections) =
            setup_protocol_endpoint().await;
        let host_addr = host.endpoint().addr();
        let remote_secret = iroh::SecretKey::generate();

        let first_remote = Endpoint::builder(iroh::endpoint::presets::N0)
            .secret_key(remote_secret.clone())
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind first remote runtime");
        let _first_connection = first_remote
            .connect(host_addr.clone(), PlutoniumProtocol::ALPN)
            .await
            .expect("connect first remote runtime");
        timeout(Duration::from_secs(5), host_events.recv())
            .await
            .expect("first host accept timeout")
            .expect("first host accept");
        let first_stable_id = host_connections
            .read()
            .await
            .get(&remote_secret.public())
            .expect("first transport installed")
            .stable_id();

        // Model a process crash: the old QUIC handle still looks locally open,
        // but the old runtime is no longer consuming OpenRTC control streams.
        // A replacement process starts with the same durable endpoint identity.
        let replacement_remote = Endpoint::builder(iroh::endpoint::presets::N0)
            .secret_key(remote_secret.clone())
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind replacement remote runtime");
        let _replacement_connection = replacement_remote
            .connect(host_addr, PlutoniumProtocol::ALPN)
            .await
            .expect("connect replacement remote runtime");
        timeout(Duration::from_secs(5), host_events.recv())
            .await
            .expect("replacement host accept timeout")
            .expect("replacement host accept");

        let replacement_stable_id = host_connections
            .read()
            .await
            .get(&remote_secret.public())
            .expect("replacement transport installed")
            .stable_id();
        assert_ne!(
            replacement_stable_id, first_stable_id,
            "an actively unresponsive old leg must not block a same-identity process restart",
        );
    }

    /// Symmetric idle iroh heartbeat should open a **bounded** number of uni streams
    /// (ping + pong on each side). Bursts far above this model usually mean non-heartbeat
    /// traffic or regressions.
    #[tokio::test]
    async fn idle_symmetric_iroh_heartbeat_send_uni_open_rate_bounded() {
        use crate::heartbeat::idle_symmetric_heartbeat_max_send_uni_opens_upper_bound;
        use crate::heartbeat::iroh_heartbeat::test_counters;
        use std::sync::atomic::Ordering;

        test_counters::reset_heartbeat_send_uni_count();

        let tick = Duration::from_millis(200);
        let heartbeat_config = HeartbeatConfig {
            tick_interval: tick,
            suspect_after: Duration::from_secs(10),
            stale_after: Duration::from_secs(30),
            send_timeout: Duration::from_secs(2),
        };

        let (tx1, _rx1) = mpsc::channel::<HealthTransition>(32);
        let (tx2, _rx2) = mpsc::channel::<HealthTransition>(32);
        let mgr1 = IrohHeartbeatManager::new();
        let mgr2 = IrohHeartbeatManager::new();

        let ep1 = Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .unwrap();
        let ep2 = Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .unwrap();

        let node1 = IrohNativeNode::spawn_with_heartbeat(ep1, mgr1, tx1, heartbeat_config.clone())
            .await
            .unwrap();
        let node2 = IrohNativeNode::spawn_with_heartbeat(ep2, mgr2, tx2, heartbeat_config)
            .await
            .unwrap();

        let remote_id = node2.endpoint().id();
        let remote_addr = node2.node_addr().await.unwrap();

        let mut conn_stream = node1.connect_addr(remote_id, remote_addr);
        let connected = timeout(Duration::from_secs(5), async {
            while let Some(ev) = conn_stream.next().await {
                match ev {
                    ConnectEvent::Connected => return true,
                    ConnectEvent::Closed { .. } => return false,
                }
            }
            false
        })
        .await
        .unwrap();
        assert!(connected, "expected outbound connect to reach Connected");

        let observe = Duration::from_millis(900);
        sleep(observe).await;

        let observed = test_counters::HEARTBEAT_SEND_UNI_COUNT.load(Ordering::SeqCst);
        let bound = idle_symmetric_heartbeat_max_send_uni_opens_upper_bound(observe, tick);
        assert!(
            observed <= bound,
            "heartbeat send_uni opens should stay within idle symmetric model (observed={} bound={})",
            observed,
            bound
        );
        assert!(
            observed >= 4,
            "expected some heartbeat uni traffic after idle window (observed={})",
            observed
        );
    }

    #[tokio::test]
    async fn active_probe_requires_remote_control_loop_round_trip() {
        let ep1 = Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind first endpoint");
        let ep2 = Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind second endpoint");

        let node1 = IrohNativeNode::spawn_with_endpoint(ep1)
            .await
            .expect("spawn first node");
        let node2 = IrohNativeNode::spawn_with_endpoint(ep2)
            .await
            .expect("spawn second node");
        let remote_id = node2.endpoint().id();
        let mut events = node1.connect_addr(
            remote_id,
            node2.node_addr().await.expect("second node address"),
        );

        let connected = timeout(Duration::from_secs(5), async {
            while let Some(event) = events.next().await {
                if matches!(event, ConnectEvent::Connected) {
                    return true;
                }
            }
            false
        })
        .await
        .expect("connect timeout");
        assert!(connected, "expected physical connection");

        let probe = node1
            .probe_connection(remote_id, Duration::from_secs(2))
            .await
            .expect("current physical generation");
        assert!(
            probe.responsive,
            "remote runtime must return the typed pong"
        );
        assert_ne!(probe.transport_stable_id, 0);
        assert!(
            probe.last_inbound_activity_age.is_some(),
            "the accepted pong stream must also be recorded as generation-bound activity",
        );
    }

    #[tokio::test]
    async fn active_probe_is_not_starved_by_bidirectional_stream_load() {
        let ep1 = Endpoint::builder(iroh::endpoint::presets::N0)
            .relay_mode(iroh::RelayMode::Disabled)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind first endpoint");
        let ep2 = Endpoint::builder(iroh::endpoint::presets::N0)
            .relay_mode(iroh::RelayMode::Disabled)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind second endpoint");

        let node1 = IrohNativeNode::spawn_with_endpoint(ep1)
            .await
            .expect("spawn first node");
        let node2 = IrohNativeNode::spawn_with_endpoint(ep2)
            .await
            .expect("spawn second node");
        let remote_id = node2.endpoint().id();
        let mut events = node1.connect_addr(
            remote_id,
            node2.node_addr().await.expect("second node address"),
        );
        timeout(Duration::from_secs(5), async {
            while let Some(event) = events.next().await {
                if matches!(event, ConnectEvent::Connected) {
                    return;
                }
            }
            panic!("connection event stream ended");
        })
        .await
        .expect("connect timeout");

        let incoming = node2.incoming_streams_stream();
        let drain = tokio::spawn(async move { while incoming.recv().await.is_ok() {} });
        let flood_node = node1.clone();
        let flood = tokio::spawn(async move {
            loop {
                let Ok((mut send, _recv)) = flood_node.open_bi(remote_id).await else {
                    break;
                };
                if send.write_all(b"x").await.is_err() {
                    break;
                }
                let _ = send.finish();
                tokio::task::yield_now().await;
            }
        });

        sleep(Duration::from_millis(25)).await;
        let probe = node1
            .probe_connection(remote_id, Duration::from_secs(2))
            .await
            .expect("current physical generation");
        flood.abort();
        drain.abort();

        assert!(
            probe.responsive,
            "bidirectional application/control load must not starve the uni-stream liveness control plane",
        );
    }

    #[tokio::test]
    async fn concurrent_native_address_dials_install_one_physical_generation() {
        let ep1 = Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind first endpoint");
        let ep2 = Endpoint::builder(iroh::endpoint::presets::N0)
            .alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
            .bind()
            .await
            .expect("bind second endpoint");

        let node1 = IrohNativeNode::spawn_with_endpoint(ep1)
            .await
            .expect("spawn first node");
        let node2 = IrohNativeNode::spawn_with_endpoint(ep2)
            .await
            .expect("spawn second node");
        let remote_id = node2.endpoint().id();
        let remote_addr = node2.node_addr().await.expect("second node address");
        let mut remote_events = node2.accept_events();

        // Exercise the physical owner directly. All callers begin before any
        // one of them has observed Connected, matching concurrent recovery
        // wakes after a network change.
        let mut dial_streams = (0..16)
            .map(|_| node1.connect_addr(remote_id, remote_addr.clone()))
            .collect::<Vec<_>>();
        let connected = futures::future::join_all(dial_streams.iter_mut().map(|events| async {
            timeout(Duration::from_secs(5), async {
                while let Some(event) = events.next().await {
                    match event {
                        ConnectEvent::Connected => return true,
                        ConnectEvent::Closed { .. } => return false,
                    }
                }
                false
            })
            .await
            .expect("concurrent dial timeout")
        }))
        .await;
        assert!(
            connected.iter().all(|connected| *connected),
            "every caller should reuse the one installed physical connection"
        );

        let accepted = timeout(Duration::from_secs(5), remote_events.next())
            .await
            .expect("remote accept timeout")
            .expect("remote accept event");
        assert!(matches!(accepted, AcceptEvent::Accepted { .. }));
        let unexpected_second_accept = timeout(Duration::from_millis(300), async {
            while let Some(event) = remote_events.next().await {
                if matches!(event, AcceptEvent::Accepted { .. }) {
                    return Some(event);
                }
            }
            None
        })
        .await;
        assert!(
            unexpected_second_accept.is_err(),
            "concurrent callers must not create a second remote generation"
        );
        assert_eq!(
            node1.active_endpoint_ids().await,
            vec![remote_id],
            "exactly one physical endpoint mapping should remain"
        );
        assert!(
            node1
                .outbound_dial_gates
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
                .is_empty(),
            "completed single-flight entries must be retired"
        );
    }

    #[tokio::test]
    async fn application_uni_stream_prefix_is_replayed_after_control_classification() {
        let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
        let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;
        let connection = r1
            .endpoint()
            .connect(r2.endpoint().addr(), PlutoniumProtocol::ALPN)
            .await
            .expect("connect endpoints");
        timeout(Duration::from_secs(5), events2.recv())
            .await
            .expect("accept event timeout")
            .expect("accept event");

        let payload = b"application-uni-payload-that-is-not-control";
        let mut send = connection.open_uni().await.expect("open app uni stream");
        send.write_all(payload)
            .await
            .expect("write app uni payload");
        send.finish().expect("finish app uni payload");

        let incoming = timeout(Duration::from_secs(5), streams2.recv())
            .await
            .expect("application stream timeout")
            .expect("application stream");
        let IncomingStreamType::Uni(mut recv) = incoming.stream else {
            panic!("expected application uni stream");
        };
        let mut received = Vec::new();
        tokio::io::AsyncReadExt::read_to_end(&mut recv, &mut received)
            .await
            .expect("read replayed application payload");
        assert_eq!(received, payload);
    }
}