openrtc 1.0.1

#![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;
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,
};

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))
        } else {
            None
        }
    };

    if !should_start_outbound_dial(existing, prefers_outbound) {
        return true;
    }
    if !should_wait_for_canonical_inbound(existing, 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,
    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 AcceptEvent {
    Accepted {
        endpoint_id: EndpointId,
        transport_stable_id: u64,
    },
    Closed {
        endpoint_id: EndpointId,
        transport_stable_id: u64,
        error: Option<String>,
        /// 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,
    },
    KeptExisting {
        fresh_transport_stable_id: u64,
        kept_transport_stable_id: u64,
    },
}

#[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,
}

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 stable_id = connection.stable_id();
    let endpoint_key = endpoint_id.to_string();
    let heartbeat_key = heartbeat_connection_key(endpoint_id, stable_id as u64);

    let mut kept_existing = 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 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() == stable_id,
                alive: previous.close_reason().is_none(),
                direction: previous_direction,
                age_ms: previous_age_ms,
            });
            let prefers_outbound = local_prefers_outbound(local_endpoint_id, endpoint_id);
            let selected_custom_transport_id = connection.paths().iter().find_map(|path| {
                if !path.is_selected() {
                    return None;
                }
                match path.remote_addr() {
                    iroh::TransportAddr::Custom(addr) => Some(addr.id()),
                    _ => None,
                }
            });
            let forced_replacement = replacement_transport_id
                .is_some_and(|expected| selected_custom_transport_id == Some(expected));
            let decision = if forced_replacement && previous.stable_id() == stable_id {
                IrohConnectionInstallDecision::Install
            } else if forced_replacement {
                IrohConnectionInstallDecision::ReplaceExisting {
                    close_existing_reason: "native-custom-transport-upgrade",
                }
            } else {
                match connection_direction {
                    IrohConnectionDirection::Inbound => {
                        decide_inbound_install(existing, prefers_outbound)
                    }
                    IrohConnectionDirection::Outbound => {
                        decide_outbound_install(existing, prefers_outbound)
                    }
                }
            };
            match decision {
                IrohConnectionInstallDecision::Install => {}
                IrohConnectionInstallDecision::ReplaceExisting {
                    close_existing_reason,
                } => {
                    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 as u64,
                        kept_transport_stable_id: previous.stable_id() as u64,
                    });
                    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);
        }
        return Ok(());
    }

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

    // 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 as u64,
        })
        .ok();

    loop {
        tokio::select! {
            biased;
            res = connection.accept_bi() => {
                match res {
                    Ok((send, recv)) => {
                        let _ = stream_sender.send(IncomingStream {
                            endpoint_id,
                            transport_stable_id: stable_id as u64,
                            stream: IncomingStreamType::Bi(send, recv),
                        }).await;
                    }
                    Err(_) => {
                        if should_break_accept_loop(&connection).await {
                            break;
                        }
                    },
                }
            }
            res = connection.accept_uni() => {
                match res {
                    Ok(recv) => {
                        // 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 as u64;
                        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,
                                            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;
    }

    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)
    );
    event_sender
        .send(AcceptEvent::Closed {
            endpoint_id,
            transport_stable_id: stable_id as u64,
            error: close_reason_debug,
            was_locally_closed,
        })
        .ok();

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

    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,
    ) -> 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(),
        }
    }

    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> {
        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(),
            None,
            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 connect(
    endpoint: &Endpoint,
    endpoint_id: EndpointId,
    event_sender: async_channel::Sender<ConnectEvent>,
    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<()> {
    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 stable_id = connection.stable_id();

    {
        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_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);
            match decide_outbound_install(
                Some(ExistingConnectionState {
                    same_stable_id: previous.stable_id() == stable_id,
                    alive: previous.close_reason().is_none(),
                    direction: previous_direction,
                    age_ms: previous_age_ms,
                }),
                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());
                    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,
            },
        );
    }

    event_sender.send(ConnectEvent::Connected).await?;

    // 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>,
    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<()> {
    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 stable_id = connection.stable_id();

    {
        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_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);
            match decide_outbound_install(
                Some(ExistingConnectionState {
                    same_stable_id: previous.stable_id() == stable_id,
                    alive: previous.close_reason().is_none(),
                    direction: previous_direction,
                    age_ms: previous_age_ms,
                }),
                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());
                    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,
            },
        );
    }

    event_sender.send(ConnectEvent::Connected).await?;

    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<ConnectEvent>,
    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)
    );

    event_sender.send(ConnectEvent::Connected).await?;
    let (accept_tx, _) = broadcast::channel(1);
    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),
        None,
    )
    .await
    .ok();

    let _ = event_sender
        .send(ConnectEvent::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<ConnectEvent>,
    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>>>,
    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,
}

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 manual_disconnect_notices = Arc::new(RwLock::new(HashSet::new()));
        let probe_registry = IrohProbeRegistry::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(),
            );
            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,
            incoming_streams: stream_sender,
            incoming_streams_receiver: stream_receiver,
            manual_disconnect_notices,
            probe_registry,
            heartbeat_manager,
            heartbeat_health_tx,
            heartbeat_config,
        })
    }

    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 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(),
                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 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(),
                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, 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 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(ConnectEvent::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| {
                connection.stable_id() as u64 == 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 = conn.stable_id() as u64;
            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 = connection.stable_id() as u64;
        let responsive = self
            .probe_registry
            .probe(
                &endpoint_id.to_string(),
                probed_stable_id,
                &connection,
                timeout,
            )
            .await;
        let current_stable_id = self
            .connections
            .read()
            .await
            .get(&endpoint_id)
            .map(|current| current.stable_id() as u64)?;
        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,
            });
        }
        Some(IrohConnectionProbe {
            transport_stable_id: probed_stable_id,
            responsive,
        })
    }

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

    /// 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<()> {
        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(),
            None,
            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<()> {
        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(),
            None,
            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,
            ))
    }
}

#[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(),
        );

        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));
    }

    #[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);
    }

    /// 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);
    }

    #[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);
    }
}