openrtc 0.2.0

#![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, RwLock};
use tokio_stream::wrappers::BroadcastStream;

use crate::heartbeat::{
    handle_incoming_ping, parse_incoming_pong, try_read_heartbeat_frame, HealthTransition,
    HeartbeatConfig, IrohHeartbeatManager,
};
use crate::iroh_connection_policy::{
    decide_inbound_install, decide_outbound_install, should_redial_without_precheck,
    ExistingConnectionState, IrohConnectionInstallDecision,
};

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 elapsed_ms_since(inserted_at: Option<Instant>) -> u64 {
    inserted_at
        .map(|instant| instant.elapsed().as_millis().min(u128::from(u64::MAX)) as u64)
        .unwrap_or(0)
}

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

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

#[derive(Debug)]
pub struct IncomingStream {
    pub endpoint_id: EndpointId,
    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,
    },
    Closed {
        endpoint_id: EndpointId,
        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,
    },
}

#[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, Instant>>>,
    local_endpoint_id: EndpointId,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    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, Instant>>>,
    local_endpoint_id: EndpointId,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
) -> std::result::Result<(), AcceptError> {
    let endpoint_id = connection.remote_id();
    let stable_id = connection.stable_id();
    let endpoint_key = endpoint_id.to_string();

    {
        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_age_ms = elapsed_ms_since(insert_times.get(&endpoint_id).copied());
            match decide_inbound_install(Some(ExistingConnectionState {
                same_stable_id: previous.stable_id() == stable_id,
                alive: previous.close_reason().is_none(),
                age_ms: previous_age_ms,
                prefer_fresh_duplicate: !local_prefers_outbound(local_endpoint_id, endpoint_id),
            })) {
                IrohConnectionInstallDecision::Install => {}
                IrohConnectionInstallDecision::KeepExisting { close_fresh_reason } => {
                    connection.close(0u8.into(), close_fresh_reason.as_bytes());
                    return Ok(());
                }
                IrohConnectionInstallDecision::ReplaceExisting {
                    close_previous_reason,
                } => {
                    previous.close(0u8.into(), close_previous_reason.as_bytes());
                }
            }
        }
        conns.insert(endpoint_id, connection.clone());
        insert_times.insert(endpoint_id, Instant::now());
    }

    // 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(
            endpoint_key.clone(),
            connection.clone(),
            heartbeat_config.clone(),
            tx.clone(),
        )
        .await;
    }

    event_sender
        .send(AcceptEvent::Accepted { endpoint_id })
        .ok();

    loop {
        tokio::select! {
            biased;
            res = connection.accept_bi() => {
                match res {
                    Ok((send, recv)) => {
                        let _ = stream_sender.send(IncomingStream {
                            endpoint_id,
                            stream: IncomingStreamType::Bi(send, recv),
                        }).await;
                    }
                    Err(_) => {
                        if should_break_accept_loop(&connection).await {
                            break;
                        }
                    },
                }
            }
            res = connection.accept_uni() => {
                match res {
                    Ok(mut recv) => {
                        // Peek the stream: if it's a heartbeat frame, handle it inline
                        // without forwarding to the app-layer stream consumer.
                        if let Some(ref mgr) = heartbeat_manager {
                            if let Some((type_id, payload)) = try_read_heartbeat_frame(&mut recv).await {
                                let conn_clone = connection.clone();
                                let mgr_clone = mgr.clone();
                                let key = endpoint_key.clone();
                                let notices = manual_disconnect_notices.clone();
                                tokio::spawn(async move {
                                    use crate::heartbeat::codec;
                                    if type_id == codec::TYPE_PING {
                                        handle_incoming_ping(&conn_clone, &payload).await;
                                    } else if type_id == codec::TYPE_PONG {
                                        if let Some(pong) = parse_incoming_pong(&payload) {
                                            mgr_clone.deliver_pong(&key, pong).await;
                                        }
                                    } else if type_id == codec::TYPE_MANUAL_DISCONNECT {
                                        notices.write().await.insert(key.clone());
                                        conn_clone.close(
                                            0u8.into(),
                                            crate::lifecycle_reason::REASON_MANUAL_DISCONNECT
                                                .as_bytes(),
                                        );
                                    }
                                });
                                continue;
                            }
                            // Non-heartbeat Uni stream: recv has been partially consumed —
                            // we cannot rewind it. Since heartbeat frames are the only Uni
                            // content in the protocol today, this branch should not be hit.
                            // If it is, forward the (now-consumed) stream handle anyway so
                            // the host can at least observe the stream event.
                        }
                        let _ = stream_sender.send(IncomingStream {
                            endpoint_id,
                            stream: IncomingStreamType::Uni(recv),
                        }).await;
                    }
                    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(&endpoint_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,
            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";

    pub 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, Instant>>>,
        local_endpoint_id: EndpointId,
        manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    ) -> Self {
        Self {
            event_sender,
            stream_sender,
            connections,
            connection_inserted_at,
            local_endpoint_id,
            manual_disconnect_notices,
            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,
            self.manual_disconnect_notices.clone(),
            self.heartbeat_manager.clone(),
            self.heartbeat_health_tx.clone(),
            self.heartbeat_config.clone(),
        )
        .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, Instant>>>,
    stream_sender: async_channel::Sender<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
) -> Result<()> {
    {
        let conns = connections.read().await;
        if let Some(existing) = conns.get(&endpoint_id) {
            if !should_redial_without_precheck(existing.close_reason().is_none()) {
                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_age_ms = elapsed_ms_since(insert_times.get(&endpoint_id).copied());
            match decide_outbound_install(Some(ExistingConnectionState {
                same_stable_id: previous.stable_id() == stable_id,
                alive: previous.close_reason().is_none(),
                age_ms: previous_age_ms,
                prefer_fresh_duplicate: local_prefers_outbound(endpoint.id(), endpoint_id),
            })) {
                IrohConnectionInstallDecision::Install => {}
                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(());
                }
                IrohConnectionInstallDecision::ReplaceExisting {
                    close_previous_reason,
                } => {
                    previous.close(0u8.into(), close_previous_reason.as_bytes());
                }
            }
        }
        conns.insert(endpoint_id, connection.clone());
        insert_times.insert(endpoint_id, Instant::now());
    }

    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(),
        manual_disconnect_notices,
        heartbeat_manager,
        heartbeat_health_tx,
        heartbeat_config,
    )
    .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, Instant>>>,
    stream_sender: async_channel::Sender<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    heartbeat_manager: Option<IrohHeartbeatManager>,
    heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
    heartbeat_config: HeartbeatConfig,
) -> Result<()> {
    {
        let conns = connections.read().await;
        if let Some(existing) = conns.get(&endpoint_id) {
            if !should_redial_without_precheck(existing.close_reason().is_none()) {
                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_age_ms = elapsed_ms_since(insert_times.get(&endpoint_id).copied());
            match decide_outbound_install(Some(ExistingConnectionState {
                same_stable_id: previous.stable_id() == stable_id,
                alive: previous.close_reason().is_none(),
                age_ms: previous_age_ms,
                prefer_fresh_duplicate: local_prefers_outbound(endpoint.id(), endpoint_id),
            })) {
                IrohConnectionInstallDecision::Install => {}
                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(());
                }
                IrohConnectionInstallDecision::ReplaceExisting {
                    close_previous_reason,
                } => {
                    previous.close(0u8.into(), close_previous_reason.as_bytes());
                }
            }
        }
        conns.insert(endpoint_id, connection.clone());
        insert_times.insert(endpoint_id, Instant::now());
    }

    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(),
        manual_disconnect_notices,
        heartbeat_manager,
        heartbeat_health_tx,
        heartbeat_config,
    )
    .await
    .ok();

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

    Ok(())
}

#[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, Instant>>>,
    incoming_streams: async_channel::Sender<IncomingStream>,
    incoming_streams_receiver: async_channel::Receiver<IncomingStream>,
    manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
    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 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(),
            );
            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,
            heartbeat_manager,
            heartbeat_health_tx,
            heartbeat_config,
        })
    }

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

    pub async fn node_addr(&self) -> Result<EndpointAddr> {
        // Browsers (WASM) require relay URLs in tickets. Relay connectivity
        // completes asynchronously after bind(); wait briefly so tickets minted
        // right after startup include relay addrs when possible.
        //
        // If the relay doesn't come online within the timeout, we still return
        // the watched addr (which will be updated later, and presence refresh
        // will republish).
        let _ =
            tokio::time::timeout(std::time::Duration::from_secs(30), self.endpoint.online()).await;
        // Use the watched address so relay URLs appear as soon as the endpoint
        // connects (rather than capturing only the initial local socket addrs).
        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();

        tokio::spawn(async move {
            let result = connect(
                &endpoint,
                endpoint_id,
                event_sender.clone(),
                connections,
                connection_inserted_at,
                stream_sender,
                manual_notices,
                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();

        tokio::spawn(async move {
            let result = connect_addr(
                &endpoint,
                endpoint_id,
                endpoint_addr,
                event_sender.clone(),
                connections,
                connection_inserted_at,
                stream_sender,
                manual_notices,
                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 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 open_bi(&self, endpoint_id: EndpointId) -> Result<(SendStream, RecvStream)> {
        let connection = {
            let conns = self.connections.read().await;
            conns.get(&endpoint_id).cloned()
        };

        if let Some(conn) = connection {
            let (send, recv) = conn.open_bi().await?;
            Ok((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))
        }
    }

    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(),
            self.manual_disconnect_notices.clone(),
            self.heartbeat_manager.clone(),
            self.heartbeat_health_tx.clone(),
            self.heartbeat_config.clone(),
        )
        .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) -> bool {
        self.manual_disconnect_notices
            .write()
            .await
            .remove(&endpoint_id.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use iroh::endpoint::Endpoint;
    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())),
        );

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

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

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