ghostlink_core/

runtime.rs

//! Distributed inference runtime scaffolding.
//!
//! This module provides lightweight execution planning primitives that bridge
//! placement plans to token-step pipeline schedules across heterogeneous devices.

use crate::circuit_breaker::CircuitBreaker;
use crate::planning::LayerAssignment;
use hmac::{Hmac, Mac};
use sha2::Sha256;
use std::io::{self, Read, Write};
use std::io::{BufReader, BufWriter};
use std::net::{Shutdown as TcpShutdown, SocketAddr, TcpListener, TcpStream};
use std::path::PathBuf;
use std::slice;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use std::time::Instant;

#[cfg(unix)]
use std::net::Shutdown as UnixSocketShutdown;
#[cfg(unix)]
use std::os::unix::net::{UnixListener, UnixStream};

/// Execution target selected for a pipeline stage.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeviceKind {
    Npu,
    Gpu,
    Cpu,
}

impl DeviceKind {
    /// Human-friendly label.
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Npu => "NPU",
            Self::Gpu => "GPU",
            Self::Cpu => "CPU",
        }
    }
}

/// Placement of one contiguous layer range onto a node and device.
#[derive(Clone, Debug, PartialEq)]
pub struct StagePlacement {
    pub node_id: String,
    pub start_layer: usize,
    pub end_layer: usize,
    pub device: DeviceKind,
    /// Estimated compute latency for one micro-batch token step.
    pub est_latency_ms: f32,
}

impl StagePlacement {
    pub fn num_layers(&self) -> usize {
        self.end_layer.saturating_sub(self.start_layer)
    }
}

/// Pipeline execution plan for one model.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct PipelinePlan {
    pub stages: Vec<StagePlacement>,
}

impl PipelinePlan {
    /// Build a pipeline plan from layer assignments and a node/device map.
    pub fn from_assignments(
        assignments: &[LayerAssignment],
        device_by_node: &std::collections::HashMap<String, DeviceKind>,
    ) -> Self {
        Self::from_assignments_with_measured(assignments, device_by_node, &Default::default())
    }

    /// Build a pipeline plan using measured latencies from cluster metrics if available.
    pub fn from_assignments_with_measured(
        assignments: &[LayerAssignment],
        device_by_node: &std::collections::HashMap<String, DeviceKind>,
        cluster: &crate::cluster::ClusterState,
    ) -> Self {
        let metrics_guard = cluster
            .metrics
            .lock()
            .unwrap_or_else(|poison| poison.into_inner());

        let stages = assignments
            .iter()
            .map(|assignment| {
                let device = device_by_node
                    .get(&assignment.node_id)
                    .copied()
                    .unwrap_or(DeviceKind::Cpu);

                // Use measured latency if available (converted from us to ms per layer)
                let per_layer_cost_ms =
                    if let Some(metrics) = metrics_guard.get(&assignment.node_id) {
                        if metrics.latency_samples > 0 && assignment.num_layers > 0 {
                            (metrics.avg_latency_us / 1000.0) / assignment.num_layers as f32
                        } else {
                            Self::default_cost_for_device(device)
                        }
                    } else {
                        Self::default_cost_for_device(device)
                    };

                StagePlacement {
                    node_id: assignment.node_id.clone(),
                    start_layer: assignment.start_layer,
                    end_layer: assignment.end_layer,
                    device,
                    est_latency_ms: per_layer_cost_ms * assignment.num_layers as f32,
                }
            })
            .collect();

        Self { stages }
    }

    fn default_cost_for_device(device: DeviceKind) -> f32 {
        match device {
            DeviceKind::Npu => 0.42,
            DeviceKind::Gpu => 0.55,
            DeviceKind::Cpu => 1.25,
        }
    }

    /// Return aggregate token-step latency estimate for one micro-batch.
    pub fn est_step_latency_ms(&self) -> f32 {
        self.stages.iter().map(|s| s.est_latency_ms).sum()
    }

    pub fn summary(&self) -> String {
        let mut out = String::from("Pipeline Plan\n");
        out.push_str("============\n");

        for (idx, stage) in self.stages.iter().enumerate() {
            out.push_str(&format!(
                "Stage {}: {} layers {}-{} on {} ({:.2} ms)\n",
                idx,
                stage.node_id,
                stage.start_layer,
                stage.end_layer,
                stage.device.as_str(),
                stage.est_latency_ms
            ));
        }

        out.push_str(&format!(
            "Estimated per-token step latency: {:.2} ms\n",
            self.est_step_latency_ms()
        ));
        out
    }
}

fn run_stage_compute(payload: &mut [f32], stage: &StagePlacement) {
    let base_rounds = stage.num_layers().max(1) / 4 + 1;
    let rounds = match stage.device {
        DeviceKind::Npu => base_rounds,
        DeviceKind::Gpu => base_rounds * 2,
        DeviceKind::Cpu => base_rounds * 3,
    };

    let alpha = match stage.device {
        DeviceKind::Npu => 1.001_f32,
        DeviceKind::Gpu => 1.003_f32,
        DeviceKind::Cpu => 1.005_f32,
    };

    for _ in 0..rounds {
        for value in payload.iter_mut() {
            *value = ((*value * alpha) + 0.125).sin();
        }
    }
}

/// Per-stage runtime telemetry captured from real in-process execution.
#[derive(Clone, Debug, PartialEq)]
pub struct StageExecutionStats {
    pub stage_idx: usize,
    pub processed_batches: usize,
    pub avg_compute_ms: f32,
    pub avg_recv_wait_ms: f32,
    pub avg_send_wait_ms: f32,
    pub avg_bridge_write_ms: f32,
    pub avg_bridge_read_ms: f32,
}

/// Aggregate runtime execution output.
#[derive(Clone, Debug, PartialEq)]
pub struct ExecutionResult {
    pub token_count: usize,
    pub micro_batch: usize,
    pub batch_count: usize,
    pub stage_count: usize,
    pub total_time_ms: f32,
    pub throughput_tokens_per_sec: f32,
    pub avg_token_latency_ms: f32,
    pub p95_token_latency_ms: f32,
    pub p99_token_latency_ms: f32,
    /// One entry per token, in completion order — the raw samples P95/P99
    /// are computed from. Exposed (rather than kept purely internal) so
    /// callers such as `ghost-link flow` can dump the real per-token
    /// latency distribution for a histogram, not just precomputed
    /// percentiles.
    pub token_latencies_ms: Vec<f32>,
    pub stage_stats: Vec<StageExecutionStats>,
}

/// Transport protocol selection for inter-stage bridges.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum TransportKind {
    /// Standard TCP transport (works across machines).
    #[default]
    Tcp,
    /// Unix domain socket transport (loopback/local only, lower overhead).
    /// Supported on Linux, macOS, and Windows 10 build 1803+.
    Unix,
}

/// Unified transport listener wrapping platform-specific socket types.
#[derive(Debug)]
pub enum BridgeListener {
    Tcp(TcpListener),
    #[cfg(unix)]
    Unix(UnixListener),
}

impl BridgeListener {
    /// Accept an incoming connection. Returns the connected stream (peer addr discarded).
    pub fn accept(&self) -> io::Result<BridgeStream> {
        match self {
            BridgeListener::Tcp(l) => l.accept().map(|(s, _)| BridgeStream::Tcp(s)),
            #[cfg(unix)]
            BridgeListener::Unix(l) => l.accept().map(|(s, _)| BridgeStream::Unix(s)),
        }
    }
}

/// Unified transport stream wrapping platform-specific socket types.
pub enum BridgeStream {
    Tcp(TcpStream),
    #[cfg(unix)]
    Unix(UnixStream),
}

impl Read for BridgeStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self {
            BridgeStream::Tcp(s) => s.read(buf),
            #[cfg(unix)]
            BridgeStream::Unix(s) => s.read(buf),
        }
    }
}

impl Write for BridgeStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            BridgeStream::Tcp(s) => s.write(buf),
            #[cfg(unix)]
            BridgeStream::Unix(s) => s.write(buf),
        }
    }
    fn flush(&mut self) -> io::Result<()> {
        match self {
            BridgeStream::Tcp(s) => s.flush(),
            #[cfg(unix)]
            BridgeStream::Unix(s) => s.flush(),
        }
    }
}

impl BridgeStream {
    /// Set TCP_NODELAY (no-op for Unix domain sockets).
    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
        match self {
            BridgeStream::Tcp(s) => s.set_nodelay(nodelay),
            #[cfg(unix)]
            BridgeStream::Unix(_) => Ok(()),
        }
    }

    /// Set socket send buffer size (SO_SNDBUF). No-op on platforms without support.
    pub fn set_send_buffer_size(&self, _size: usize) -> io::Result<()> {
        match self {
            BridgeStream::Tcp(_s) => {
                // Platform-specific socket buffer tuning would go here
                // Omitted for portability - most OSes auto-tune well
                Ok(())
            }
            #[cfg(unix)]
            BridgeStream::Unix(_) => Ok(()),
        }
    }

    /// Set socket receive buffer size (SO_RCVBUF). No-op on platforms without support.
    pub fn set_recv_buffer_size(&self, _size: usize) -> io::Result<()> {
        match self {
            BridgeStream::Tcp(_s) => {
                // Platform-specific socket buffer tuning would go here
                // Omitted for portability - most OSes auto-tune well
                Ok(())
            }
            #[cfg(unix)]
            BridgeStream::Unix(_) => Ok(()),
        }
    }

    /// Shut down the write half of the connection.
    pub fn shutdown_write(&self) -> io::Result<()> {
        match self {
            BridgeStream::Tcp(s) => s.shutdown(TcpShutdown::Write),
            #[cfg(unix)]
            BridgeStream::Unix(s) => s.shutdown(UnixSocketShutdown::Write),
        }
    }
}

/// Unified transport address for TCP or Unix domain sockets.
#[derive(Clone, Debug)]
pub enum BridgeAddr {
    Tcp(SocketAddr),
    Unix(PathBuf),
}

/// Build a collision-free Unix socket path for an inter-stage bridge.
///
/// Paths must be unique per process AND per pipeline: a fixed name like
/// `ghostlink-bridge-0.sock` lets two concurrent pipelines delete each
/// other's live sockets and cross-connect their streams, which surfaces
/// as garbled frames (auth failures / bogus payload lengths) far from
/// the cause. PID plus a process-wide counter guarantees uniqueness;
/// callers still remove stale files defensively before binding.
#[cfg(unix)]
fn unique_bridge_socket_path(source_stage: usize) -> std::path::PathBuf {
    use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
    static BRIDGE_SOCKET_SEQ: AtomicU64 = AtomicU64::new(0);
    let seq = BRIDGE_SOCKET_SEQ.fetch_add(1, AtomicOrdering::Relaxed);
    let mut p = std::env::temp_dir();
    p.push(format!(
        "ghostlink-bridge-{}-{}-{}.sock",
        std::process::id(),
        seq,
        source_stage
    ));
    p
}

impl BridgeAddr {
    /// Connect to the address (with a timeout for TCP, instant for Unix).
    fn connect(&self, timeout: Duration) -> io::Result<BridgeStream> {
        match self {
            BridgeAddr::Tcp(addr) => {
                TcpStream::connect_timeout(addr, timeout).map(BridgeStream::Tcp)
            }
            #[cfg(unix)]
            BridgeAddr::Unix(path) => UnixStream::connect(path).map(BridgeStream::Unix),
            #[cfg(not(unix))]
            BridgeAddr::Unix(_) => Err(io::Error::new(
                io::ErrorKind::Unsupported,
                "Unix domain sockets require a Unix platform",
            )),
        }
    }
}

/// Runtime controls for TCP/Unix transport execution.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TcpTransportConfig {
    pub max_inflight_batches: usize,
    pub reconnect_attempts: usize,
    pub reconnect_backoff_ms: u64,
    /// Per-batch HMAC-SHA256 auth token (legacy, use `auth_session_token` for performance).
    pub auth_token: Option<String>,
    /// Session-level auth token. Does a one-time HMAC handshake at connection setup,
    /// then verifies a lightweight monotonic sequence counter per batch. Much faster
    /// than per-batch HMAC while maintaining connection-level security.
    pub auth_session_token: Option<String>,
    pub use_mtls: bool,
    pub cert_chain_path: Option<String>,
    /// Transport protocol: `Tcp` (LAN) or `Unix` (loopback/local, lower overhead).
    pub transport: TransportKind,
    /// Socket send buffer size (SO_SNDBUF). 0 = system default.
    pub send_buffer_size: usize,
    /// Socket receive buffer size (SO_RCVBUF). 0 = system default.
    pub recv_buffer_size: usize,
    /// Use vectored I/O (writev/readv) for zero-copy header+payload writes.
    pub use_vectored_io: bool,
    /// Synthetic `f32` elements per token in the benchmark/demo payload these
    /// pipeline-execution paths generate — they move a stand-in payload
    /// through the real transport rather than real model activations (see
    /// `ghost-link::main`'s `flow` command). Defaults to 16 elements (64
    /// bytes/token), matching the historical hardcoded value so every
    /// existing caller keeps identical behavior; the `flow` CLI raises this
    /// to reflect a real FP16/BF16 hidden-state size (e.g. 4096 × 2 bytes ÷
    /// 4 bytes/f32 = 2048 elements/token for a 7B-class model).
    pub elems_per_token: usize,
}

impl Default for TcpTransportConfig {
    fn default() -> Self {
        Self {
            max_inflight_batches: 512,
            reconnect_attempts: 3,
            reconnect_backoff_ms: 25,
            auth_token: None,
            auth_session_token: None,
            use_mtls: false,
            cert_chain_path: None,
            transport: TransportKind::Tcp,
            send_buffer_size: 256 * 1024, // 256 KB
            recv_buffer_size: 256 * 1024, // 256 KB
            use_vectored_io: true,
            elems_per_token: 16,
        }
    }
}

/// Result of a one-time session auth handshake. Shared between reader and writer threads.
#[derive(Clone, Debug)]
struct SessionAuth {
    session_id: u64,
    seq: Arc<AtomicU64>,
    peer_seq: Arc<AtomicU64>,
}

impl ExecutionResult {
    pub fn summary(&self) -> String {
        let mut out = String::from("Execution Runtime\n");
        out.push_str("=================\n");
        out.push_str(&format!(
            "Tokens: {} | Micro-batch: {} | Batches: {} | Stages: {}\n",
            self.token_count, self.micro_batch, self.batch_count, self.stage_count
        ));
        out.push_str(&format!(
            "Measured wall-clock time: {:.2} ms\n",
            self.total_time_ms
        ));
        out.push_str(&format!(
            "Throughput: {:.2} tokens/sec\n",
            self.throughput_tokens_per_sec
        ));
        out.push_str(&format!(
            "Avg token latency: {:.2} ms | P95: {:.2} ms | P99: {:.2} ms\n",
            self.avg_token_latency_ms, self.p95_token_latency_ms, self.p99_token_latency_ms
        ));

        for stage in &self.stage_stats {
            out.push_str(&format!(
                "Stage {} batches={} compute={:.2} ms recv-wait={:.2} ms send-wait={:.2} ms bridge-write={:.2} ms bridge-read={:.2} ms\n",
                stage.stage_idx,
                stage.processed_batches,
                stage.avg_compute_ms,
                stage.avg_recv_wait_ms,
                stage.avg_send_wait_ms,
                stage.avg_bridge_write_ms,
                stage.avg_bridge_read_ms
            ));
        }

        out
    }
}

/// One scheduled operation in steady-state token generation.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TokenStep {
    pub token_idx: usize,
    pub stage_idx: usize,
}

/// Generate a simple steady-state pipeline schedule preview.
///
/// For each token, all stages execute in order; this preview is useful for
/// verifying stage count and queue depth wiring before integrating transport.
pub fn build_token_schedule(stage_count: usize, token_count: usize) -> Vec<TokenStep> {
    let mut schedule = Vec::with_capacity(stage_count.saturating_mul(token_count));
    for token_idx in 0..token_count {
        for stage_idx in 0..stage_count {
            schedule.push(TokenStep {
                token_idx,
                stage_idx,
            });
        }
    }
    schedule
}

/// Execute pipeline stages with real channel/thread wiring for `token_count` tokens.
///
/// `micro_batch` controls how many tokens are packaged per batch. Values below
/// 1 are clamped to 1.
pub fn execute_pipeline(
    plan: &PipelinePlan,
    token_count: usize,
    micro_batch: usize,
) -> Result<ExecutionResult, String> {
    execute_pipeline_with_rebalance(plan, token_count, micro_batch, None)
}

pub fn execute_pipeline_with_rebalance(
    plan: &PipelinePlan,
    token_count: usize,
    micro_batch: usize,
    rebalance: Option<&crate::planning::RebalanceTrigger>,
) -> Result<ExecutionResult, String> {
    execute_pipeline_with_rebalance_and_measured(
        plan,
        token_count,
        micro_batch,
        rebalance,
        None,
        None,
    )
}

pub fn execute_pipeline_with_rebalance_and_measured(
    plan: &PipelinePlan,
    token_count: usize,
    micro_batch: usize,
    rebalance: Option<&crate::planning::RebalanceTrigger>,
    cluster: Option<&crate::cluster::ClusterState>,
    placement_context: Option<&crate::planning::PlacementPlan>,
) -> Result<ExecutionResult, String> {
    let stage_count = plan.stages.len();
    let micro_batch = micro_batch.max(1);
    let batch_count = token_count.div_ceil(micro_batch);
    if stage_count == 0 || token_count == 0 {
        return Ok(ExecutionResult {
            token_count,
            micro_batch,
            batch_count,
            stage_count,
            total_time_ms: 0.0,
            throughput_tokens_per_sec: 0.0,
            avg_token_latency_ms: 0.0,
            p95_token_latency_ms: 0.0,
            p99_token_latency_ms: 0.0,
            token_latencies_ms: Vec::new(),
            stage_stats: Vec::new(),
        });
    }

    #[derive(Debug)]
    struct BatchWork {
        tokens_in_batch: usize,
        started_at: Instant,
        payload: Vec<f32>,
    }

    #[derive(Debug)]
    struct StageAccumulator {
        stage_idx: usize,
        processed_batches: usize,
        total_compute_ms: f32,
        total_recv_wait_ms: f32,
        total_send_wait_ms: f32,
    }

    use crate::ring::{RingConfig, SpscRingBuffer};
    use std::sync::Arc;

    // Use zero-copy SPSC ring buffers for high-throughput in-memory execution.
    // Use default RingConfig (4095 capacity, 2867 backpressure threshold) for maximum batching.
    let ring_cfg = RingConfig::default();

    let mut rings = Vec::with_capacity(stage_count + 1);
    for _ in 0..=stage_count {
        rings.push(Arc::new(SpscRingBuffer::<BatchWork>::new(ring_cfg)));
    }

    // Return ring for recycling payload vectors to reduce allocation churn.
    let recycler_ring = Arc::new(SpscRingBuffer::<Vec<f32>>::new(ring_cfg));

    let mut stage_handles = Vec::with_capacity(stage_count);
    for stage_idx in 0..stage_count {
        let stage = plan.stages[stage_idx].clone();
        let rx_ring = Arc::clone(&rings[stage_idx]);
        let tx_ring = Arc::clone(&rings[stage_idx + 1]);

        let handle = thread::spawn(move || {
            let mut total_compute_ms = 0.0_f32;
            let mut total_recv_wait_ms = 0.0_f32;
            let mut total_send_wait_ms = 0.0_f32;

            // Fixed batch count avoids OS scheduler rendezvous — each stage spins
            // on wait_for_data() / wait_for_space() with exponential backoff, staying
            // hot on its core.
            for _ in 0..batch_count {
                let recv_start = Instant::now();
                rx_ring.wait_for_data();
                let mut batch = rx_ring.pop().unwrap();
                total_recv_wait_ms += recv_start.elapsed().as_secs_f32() * 1000.0;

                let compute_start = Instant::now();
                run_stage_compute(&mut batch.payload, &stage);
                total_compute_ms += compute_start.elapsed().as_secs_f32() * 1000.0;

                let send_start = Instant::now();
                tx_ring.wait_for_space();
                let _ = tx_ring.push(batch);
                total_send_wait_ms += send_start.elapsed().as_secs_f32() * 1000.0;
            }

            StageAccumulator {
                stage_idx,
                processed_batches: batch_count,
                total_compute_ms,
                total_recv_wait_ms,
                total_send_wait_ms,
            }
        });

        stage_handles.push(handle);
    }

    let completion_ring = Arc::clone(&rings[stage_count]);
    let entry_ring = Arc::clone(&rings[0]);
    // Clear rings vector so strong_count can drop once threads finish
    drop(rings);

    let recycler_ring_c = Arc::clone(&recycler_ring);
    let latencies_handle = thread::spawn(move || {
        let mut latencies = Vec::with_capacity(token_count);
        for _ in 0..batch_count {
            completion_ring.wait_for_data();
            if let Some(done_batch) = completion_ring.pop() {
                let batch_latency_ms = done_batch.started_at.elapsed().as_secs_f32() * 1000.0;
                for _ in 0..done_batch.tokens_in_batch {
                    latencies.push(batch_latency_ms);
                }
                recycler_ring_c.wait_for_space();
                let _ = recycler_ring_c.push(done_batch.payload);
            }
        }
        latencies
    });

    let exec_start = Instant::now();
    for batch_idx in 0..batch_count {
        let batch_start_token = batch_idx * micro_batch;
        let tokens_in_batch = (token_count - batch_start_token).min(micro_batch);
        let payload_len = (tokens_in_batch.max(1) * 16).max(32);

        let mut payload = recycler_ring
            .pop()
            .unwrap_or_else(|| Vec::with_capacity(payload_len));
        payload.resize(payload_len, 0.0);
        for (idx, val) in payload.iter_mut().enumerate() {
            *val = (batch_idx as f32 * 0.01) + (idx as f32 * 0.0001);
        }

        // Evaluate dynamic rebalance trigger during execution
        if let (Some(trigger), Some(cluster_state), Some(placement)) =
            (rebalance, cluster, placement_context)
        {
            // Check for rebalance every 25th batch to allow stability between checks
            if batch_idx % 25 == 0 && batch_idx > 0 {
                if let Some(migration) = trigger.evaluate(cluster_state, placement) {
                    tracing::info!(
                        "Dynamic Runtime: Rebalance Triggered! Moving layers {:?} from {} to {}",
                        migration.layers,
                        migration.source_node,
                        migration.target_node
                    );
                    for step in migration.generate_handoff_plan() {
                        tracing::info!("  -> {}", step);
                    }
                }
            }
        }

        entry_ring.wait_for_space();
        let _ = entry_ring.push(BatchWork {
            tokens_in_batch,
            started_at: Instant::now(),
            payload,
        });
    }
    // Drop entry ring to signal completion to stage threads
    drop(entry_ring);

    let token_latencies = latencies_handle.join().unwrap_or_default();
    let total_time_ms = exec_start.elapsed().as_secs_f32() * 1000.0;

    let mut stage_stats = Vec::with_capacity(stage_count);
    for handle in stage_handles {
        if let Ok(stats) = handle.join() {
            let divisor = stats.processed_batches.max(1) as f32;
            let avg_compute_ms = stats.total_compute_ms / divisor;

            // Feedback loop: update cluster state with actual compute latency
            if let Some(cluster_state) = cluster {
                if let Some(stage) = plan.stages.get(stats.stage_idx) {
                    cluster_state.get_metrics_mut(&stage.node_id, |m| {
                        // Convert ms back to us for ClusterState (avg_latency_us)
                        m.record_latency(avg_compute_ms * 1000.0);
                    });
                }
            }

            stage_stats.push(StageExecutionStats {
                stage_idx: stats.stage_idx,
                processed_batches: stats.processed_batches,
                avg_compute_ms,
                avg_recv_wait_ms: stats.total_recv_wait_ms / divisor,
                avg_send_wait_ms: stats.total_send_wait_ms / divisor,
                avg_bridge_write_ms: 0.0,
                avg_bridge_read_ms: 0.0,
            });
        }
    }
    stage_stats.sort_by_key(|s| s.stage_idx);

    let throughput_tokens_per_sec = if total_time_ms > 0.0 {
        token_count as f32 / (total_time_ms / 1000.0)
    } else {
        0.0
    };

    let avg_token_latency_ms = if token_latencies.is_empty() {
        0.0
    } else {
        token_latencies.iter().sum::<f32>() / token_latencies.len() as f32
    };

    let mut sorted = token_latencies.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let p95_idx = (sorted.len().saturating_sub(1) as f32 * 0.95).round() as usize;
    let p95_token_latency_ms = sorted.get(p95_idx).copied().unwrap_or(0.0);
    let p99_idx = (sorted.len().saturating_sub(1) as f32 * 0.99).round() as usize;
    let p99_token_latency_ms = sorted.get(p99_idx).copied().unwrap_or(0.0);

    Ok(ExecutionResult {
        token_count,
        micro_batch,
        batch_count,
        stage_count,
        total_time_ms,
        throughput_tokens_per_sec,
        avg_token_latency_ms,
        p95_token_latency_ms,
        p99_token_latency_ms,
        token_latencies_ms: token_latencies,
        stage_stats,
    })
}

#[derive(Debug, Clone)]
pub struct TransportBatch {
    pub batch_id: usize,
    pub tokens_in_batch: usize,
    pub payload: Vec<f32>,
}

fn auth_tag(
    source_stage: usize,
    batch_id: usize,
    tokens_in_batch: usize,
    payload: &[f32],
    token: &str,
) -> [u8; 32] {
    let mut mac = Hmac::<Sha256>::new_from_slice(token.as_bytes())
        .expect("HMAC key setup for transport auth failed");
    mac.update(&(source_stage as u32).to_le_bytes());
    mac.update(&(batch_id as u64).to_le_bytes());
    mac.update(&(tokens_in_batch as u32).to_le_bytes());
    mac.update(&(payload.len() as u32).to_le_bytes());
    let payload_bytes = payload_as_le_bytes(payload);
    mac.update(payload_bytes.as_ref());
    mac.finalize().into_bytes().into()
}

fn payload_as_le_bytes(payload: &[f32]) -> std::borrow::Cow<'_, [u8]> {
    if cfg!(target_endian = "little") {
        // SAFETY: f32 is POD; casting [f32] to its contiguous byte representation is valid.
        let bytes = unsafe {
            slice::from_raw_parts(
                payload.as_ptr() as *const u8,
                std::mem::size_of_val(payload),
            )
        };
        std::borrow::Cow::Borrowed(bytes)
    } else {
        let mut bytes = Vec::with_capacity(payload.len() * 4);
        for value in payload {
            bytes.extend_from_slice(&value.to_le_bytes());
        }
        std::borrow::Cow::Owned(bytes)
    }
}

fn write_transport_batch(
    writer: &mut impl Write,
    batch: &TransportBatch,
    source_stage: usize,
    token: Option<&str>,
    frame_buf: &mut Vec<u8>,
) -> io::Result<()> {
    _write_transport_batch_inner(writer, batch, source_stage, token, None, frame_buf)
}

fn write_transport_batch_session(
    writer: &mut impl Write,
    batch: &TransportBatch,
    source_stage: usize,
    session: &SessionAuth,
    frame_buf: &mut Vec<u8>,
) -> io::Result<()> {
    _write_transport_batch_inner(writer, batch, source_stage, None, Some(session), frame_buf)
}

fn _write_transport_batch_inner(
    writer: &mut impl Write,
    batch: &TransportBatch,
    source_stage: usize,
    token: Option<&str>,
    session: Option<&SessionAuth>,
    frame_buf: &mut Vec<u8>,
) -> io::Result<()> {
    let batch_id = batch.batch_id as u64;
    let tokens = batch.tokens_in_batch as u32;
    let payload_len = batch.payload.len() as u32;
    let source_stage_u16 = source_stage as u16;

    frame_buf.clear();
    let header_size = if session.is_some() {
        // source_stage(2) + batch_id(8) + tokens(4) + payload_len(4) + tag_present(1) + session_id(8) + generation(8)
        2 + 8 + 4 + 4 + 1 + 8 + 8
    } else if token.is_some() {
        // legacy HMAC: 2 + 8 + 4 + 4 + 1 + 32
        2 + 8 + 4 + 4 + 1 + 32
    } else {
        2 + 8 + 4 + 4 + 1
    };
    // Reserve space ONLY for the header to avoid large intermediate allocation/capacity checks.
    frame_buf.reserve(header_size);
    frame_buf.extend_from_slice(&source_stage_u16.to_le_bytes());
    frame_buf.extend_from_slice(&batch_id.to_le_bytes());
    frame_buf.extend_from_slice(&tokens.to_le_bytes());
    frame_buf.extend_from_slice(&payload_len.to_le_bytes());

    if let Some(s) = session {
        frame_buf.push(2); // session-level auth
        let gen = s.seq.fetch_add(1, Ordering::Relaxed);
        frame_buf.extend_from_slice(&s.session_id.to_le_bytes());
        frame_buf.extend_from_slice(&gen.to_le_bytes());
    } else if let Some(t) = token {
        frame_buf.push(1); // legacy HMAC present
        let tag = auth_tag(
            source_stage,
            batch.batch_id,
            batch.tokens_in_batch,
            &batch.payload,
            t,
        );
        frame_buf.extend_from_slice(&tag);
    } else {
        frame_buf.push(0); // No auth
    }

    // Zero-Copy Transport Optimization: By writing the header and the float payload
    // separately directly to the underlying BufWriter, we completely avoid copying
    // the entire multi-kilobyte/megabyte float slice payload into the intermediate
    // scratch frame_buf vector. This yields a massive reduction in memory bandwidth and CPU cycles.
    let payload_bytes = payload_as_le_bytes(&batch.payload);
    writer.write_all(frame_buf)?;
    writer.write_all(payload_bytes.as_ref())?;

    Ok(())
}

fn read_transport_batch(
    reader: &mut impl Read,
    expected_source_stage: usize,
    token: Option<&str>,
    payload_buf: &mut Vec<f32>,
) -> io::Result<Option<TransportBatch>> {
    _read_transport_batch_inner(reader, expected_source_stage, token, None, payload_buf)
}

fn read_transport_batch_session(
    reader: &mut impl Read,
    expected_source_stage: usize,
    session: &SessionAuth,
    payload_buf: &mut Vec<f32>,
) -> io::Result<Option<TransportBatch>> {
    _read_transport_batch_inner(
        reader,
        expected_source_stage,
        None,
        Some(session),
        payload_buf,
    )
}

fn _read_transport_batch_inner(
    reader: &mut impl Read,
    expected_source_stage: usize,
    token: Option<&str>,
    session: Option<&SessionAuth>,
    payload_buf: &mut Vec<f32>,
) -> io::Result<Option<TransportBatch>> {
    // Coalesce multiple read_exact calls into a single 19-byte read of the static packet header.
    // Length breakdown: source_stage(2) + batch_id(8) + tokens(4) + payload_len(4) + tag_type(1) = 19 bytes.
    let mut header_buf = [0u8; 19];
    match reader.read_exact(&mut header_buf) {
        Ok(()) => {}
        Err(err) if err.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
        Err(err) => return Err(err),
    }

    let source_stage = u16::from_le_bytes(header_buf[0..2].try_into().unwrap()) as usize;
    if source_stage != expected_source_stage {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "unexpected source stage in transport frame",
        ));
    }

    let batch_id = u64::from_le_bytes(header_buf[2..10].try_into().unwrap()) as usize;
    let tokens_in_batch = u32::from_le_bytes(header_buf[10..14].try_into().unwrap()) as usize;
    let payload_len = u32::from_le_bytes(header_buf[14..18].try_into().unwrap()) as usize;

    // Bound the wire-declared payload length BEFORE allocating. A corrupted or
    // cross-wired frame can otherwise declare up to u32::MAX elements (~17 GB
    // of f32), and the resulting failed allocation aborts the process (or gets
    // the backend OOM-killed on Linux) instead of surfacing an error.
    const MAX_WIRE_PAYLOAD_ELEMS: usize = 64 * 1024 * 1024; // 256 MB of f32
    if payload_len > MAX_WIRE_PAYLOAD_ELEMS {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "wire payload length {} exceeds maximum {} (corrupted or cross-wired frame?)",
                payload_len, MAX_WIRE_PAYLOAD_ELEMS
            ),
        ));
    }

    let tag_type = header_buf[18];

    match tag_type {
        2 => {
            // Session-level auth: read session_id + generation in a single 16-byte read instead of two 8-byte reads.
            let mut session_bytes = [0u8; 16];
            reader.read_exact(&mut session_bytes)?;
            let frame_session_id = u64::from_le_bytes(session_bytes[0..8].try_into().unwrap());
            let generation = u64::from_le_bytes(session_bytes[8..16].try_into().unwrap());

            if session.is_none() {
                return Err(io::Error::new(
                    io::ErrorKind::PermissionDenied,
                    "session auth required but no session established",
                ));
            }
            let s = session.unwrap();
            if frame_session_id != s.session_id {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "session auth: session_id mismatch",
                ));
            }
            let prev = s.peer_seq.fetch_max(generation, Ordering::AcqRel);
            if generation < prev {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "session auth: non-monotonic generation",
                ));
            }
        }
        1 => {
            // Legacy HMAC auth
            if token.is_none() {
                return Err(io::Error::new(
                    io::ErrorKind::PermissionDenied,
                    "HMAC auth present but no token configured",
                ));
            }
            let mut received_tag = [0u8; 32];
            reader.read_exact(&mut received_tag)?;
            payload_buf.resize(payload_len, 0.0);
            _read_payload(reader, payload_buf, payload_len)?;

            let expected_tag = auth_tag(
                source_stage,
                batch_id,
                tokens_in_batch,
                payload_buf,
                token.unwrap(),
            );
            if received_tag != expected_tag {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "transport auth tag mismatch",
                ));
            }
            return Ok(Some(TransportBatch {
                batch_id,
                tokens_in_batch,
                payload: payload_buf.clone(),
            }));
        }
        0 => {
            // No auth
            if token.is_some() || session.is_some() {
                return Err(io::Error::new(
                    io::ErrorKind::PermissionDenied,
                    "transport authentication required but no auth present",
                ));
            }
        }
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "unknown transport auth type",
            ));
        }
    }

    // Common path: read payload after auth verification
    payload_buf.resize(payload_len, 0.0);
    _read_payload(reader, payload_buf, payload_len)?;

    Ok(Some(TransportBatch {
        batch_id,
        tokens_in_batch,
        payload: payload_buf.clone(),
    }))
}

fn _read_payload(
    reader: &mut impl Read,
    payload_buf: &mut [f32],
    payload_len: usize,
) -> io::Result<()> {
    if payload_len == 0 {
        return Ok(());
    }
    if cfg!(target_endian = "little") {
        let payload_bytes = unsafe {
            slice::from_raw_parts_mut(
                payload_buf.as_mut_ptr() as *mut u8,
                payload_len * std::mem::size_of::<f32>(),
            )
        };
        reader.read_exact(payload_bytes)?;
    } else {
        let mut payload_bytes = vec![0u8; payload_len * 4];
        reader.read_exact(&mut payload_bytes)?;
        for (i, chunk) in payload_bytes.as_chunks::<4>().0.iter().enumerate() {
            payload_buf[i] = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
        }
    }
    Ok(())
}

/// Spawn a TCP transport bridge between pipeline stages.
///
/// In a multi-node deployment, the `bind_addr` and `connect_addr` can point to
/// different physical interfaces or remote hosts.
///
/// `circuit_breaker` is optional and persists failure history *across* calls
/// (unlike `config.reconnect_attempts`, which only bounds retries within this
/// one call) — pass `ClusterState::circuit_breaker_for(node_id)` for a real
/// remote node so a chronically-unreachable node fails fast on subsequent
/// pipeline runs instead of repeating the full connect/backoff sequence every
/// time. `None` (e.g. loopback benchmarking, where there's no real remote
/// failure mode to protect against) reproduces the prior always-retry behavior
/// exactly.
pub fn spawn_tcp_bridge(
    source_stage: usize,
    input_rx: mpsc::Receiver<TransportBatch>,
    output_tx: mpsc::SyncSender<TransportBatch>,
    config: TcpTransportConfig,
    listener: BridgeListener,
    addr: BridgeAddr,
    circuit_breaker: Option<CircuitBreaker>,
) -> thread::JoinHandle<BridgeAccumulator> {
    thread::spawn(move || {
        let client_stream = {
            let breaker_open = circuit_breaker
                .as_ref()
                .is_some_and(|breaker| !breaker.should_attempt());

            let mut connected = None;
            if breaker_open {
                tracing::warn!(
                    "Bridge: circuit breaker open for {:?}; skipping connection attempts",
                    addr
                );
            } else {
                let attempts = config.reconnect_attempts.max(1);
                for attempt in 0..attempts {
                    match addr.connect(Duration::from_secs(1)) {
                        Ok(stream) => {
                            connected = Some(stream);
                            break;
                        }
                        Err(e) => {
                            tracing::debug!(
                                "Bridge: Connection attempt {}/{} failed for {:?}: {}",
                                attempt + 1,
                                attempts,
                                addr,
                                e
                            );
                            thread::sleep(Duration::from_millis(config.reconnect_backoff_ms));
                        }
                    }
                }
            }

            match connected {
                Some(stream) => {
                    if let Some(breaker) = &circuit_breaker {
                        breaker.record_success();
                    }
                    stream
                }
                None => {
                    if let Some(breaker) = &circuit_breaker {
                        breaker.record_failure();
                    }
                    tracing::error!(
                        "Bridge: Exhausted retries for {:?}. Falling back to passthrough.",
                        addr
                    );
                    for batch in input_rx {
                        if output_tx.send(batch).is_err() {
                            break;
                        }
                    }
                    return BridgeAccumulator::default_with_stage(source_stage);
                }
            }
        };

        // Inter-stage frames are small and latency-sensitive. For TCP, Nagle's
        // algorithm would coalesce these and stall on peer delayed-ACKs (~40ms),
        // which compounds across hops and dominates multi-stage paths.
        if let Err(e) = client_stream.set_nodelay(true) {
            tracing::debug!("Bridge: failed to set TCP_NODELAY on client: {}", e);
        }

        let server_stream = match listener.accept() {
            Ok(stream) => {
                // Once the peer is connected the Unix socket file is no longer
                // needed; unlink immediately so unique per-pipeline socket
                // paths do not accumulate in the temp directory.
                #[cfg(unix)]
                if let BridgeAddr::Unix(path) = &addr {
                    let _ = std::fs::remove_file(path);
                }
                stream
            }
            Err(_) => {
                #[cfg(unix)]
                if let BridgeAddr::Unix(path) = &addr {
                    let _ = std::fs::remove_file(path);
                }
                for batch in input_rx {
                    if output_tx.send(batch).is_err() {
                        break;
                    }
                }
                return BridgeAccumulator::default_with_stage(source_stage);
            }
        };

        if let Err(e) = server_stream.set_nodelay(true) {
            tracing::debug!("Bridge: failed to set TCP_NODELAY on server: {}", e);
        }

        // Set socket buffer sizes for better throughput
        if let Err(e) = client_stream.set_send_buffer_size(config.send_buffer_size) {
            tracing::debug!("Bridge: failed to set SO_SNDBUF on client: {}", e);
        }
        if let Err(e) = client_stream.set_recv_buffer_size(config.recv_buffer_size) {
            tracing::debug!("Bridge: failed to set SO_RCVBUF on client: {}", e);
        }
        if let Err(e) = server_stream.set_send_buffer_size(config.send_buffer_size) {
            tracing::debug!("Bridge: failed to set SO_SNDBUF on server: {}", e);
        }
        if let Err(e) = server_stream.set_recv_buffer_size(config.recv_buffer_size) {
            tracing::debug!("Bridge: failed to set SO_RCVBUF on server: {}", e);
        }

        let session = config.auth_session_token.as_ref().map(|session_token| {
            // Derive a deterministic session_id from the token (no handshake round-trip).
            // The TCP connect/accept already authenticates this connection.
            // Use SHA-256 of the token (first 8 bytes) for cross-platform determinism.
            let token_hash = {
                let mut mac = Hmac::<Sha256>::new_from_slice(session_token.as_bytes())
                    .expect("HMAC key for session_id derivation");
                mac.update(b"session-id");
                mac.finalize().into_bytes()
            };
            let session_id = u64::from_le_bytes([
                token_hash[0],
                token_hash[1],
                token_hash[2],
                token_hash[3],
                token_hash[4],
                token_hash[5],
                token_hash[6],
                token_hash[7],
            ]);
            tracing::debug!(
                "TCP Bridge: session auth established, session_id={}",
                session_id
            );
            SessionAuth {
                session_id,
                seq: Arc::new(AtomicU64::new(0)),
                peer_seq: Arc::new(AtomicU64::new(0)),
            }
        });

        let use_session = session.clone();
        let writer_use_session = use_session.clone();
        let reader_use_session = use_session.clone();
        let writer_auth_token = if writer_use_session.is_some() {
            None
        } else {
            config
                .auth_session_token
                .clone()
                .or_else(|| config.auth_token.clone())
        };
        let reader_auth_token = if reader_use_session.is_some() {
            None
        } else {
            config
                .auth_session_token
                .clone()
                .or_else(|| config.auth_token.clone())
        };

        let writer = thread::spawn(move || {
            let mut writer = BufWriter::with_capacity(64 * 1024, client_stream);
            let mut processed_batches = 0usize;
            let mut total_write_ms = 0.0_f32;
            let mut frame_buf = Vec::with_capacity(64 * 1024);
            let write_result: Result<(), ()> = (|| {
                if let Some(ref s) = writer_use_session {
                    for batch in input_rx {
                        let write_start = Instant::now();
                        write_transport_batch_session(
                            &mut writer,
                            &batch,
                            source_stage,
                            s,
                            &mut frame_buf,
                        )
                        .map_err(|_| ())?;
                        // BufWriter only spills to the socket once its 64KB
                        // buffer fills — for the small, sub-buffer-size frames
                        // this bridge carries, that meant every batch sat in
                        // userspace until the whole input channel drained, so
                        // the receiver saw nothing until one final burst at
                        // stream end instead of the pipelined, per-batch
                        // delivery TCP_NODELAY above was set up to provide.
                        // Flushing per frame is what actually puts NODELAY to
                        // work for this latency-sensitive path.
                        writer.flush().map_err(|_| ())?;
                        total_write_ms += write_start.elapsed().as_secs_f32() * 1000.0;
                        processed_batches += 1;
                    }
                } else {
                    for batch in input_rx {
                        let write_start = Instant::now();
                        write_transport_batch(
                            &mut writer,
                            &batch,
                            source_stage,
                            writer_auth_token.as_deref(),
                            &mut frame_buf,
                        )
                        .map_err(|_| ())?;
                        writer.flush().map_err(|_| ())?;
                        total_write_ms += write_start.elapsed().as_secs_f32() * 1000.0;
                        processed_batches += 1;
                    }
                }
                Ok(())
            })();
            let _ = writer.flush();
            let _ = writer.get_ref().shutdown_write();
            let _ = write_result;
            (processed_batches, total_write_ms)
        });

        let mut reader = BufReader::with_capacity(64 * 1024, server_stream);
        let mut read_batches = 0usize;
        let mut total_read_ms = 0.0_f32;
        let mut payload_buf = Vec::with_capacity(16 * 1024);

        let _: Result<(), ()> = if let Some(ref s) = reader_use_session {
            loop {
                let read_start = Instant::now();
                match read_transport_batch_session(&mut reader, source_stage, s, &mut payload_buf) {
                    Ok(Some(batch)) => {
                        total_read_ms += read_start.elapsed().as_secs_f32() * 1000.0;
                        if output_tx.send(batch).is_err() {
                            break Ok(());
                        }
                        read_batches += 1;
                    }
                    Ok(None) => break Ok(()),
                    Err(_) => break Err(()),
                }
            }
        } else {
            loop {
                let read_start = Instant::now();
                match read_transport_batch(
                    &mut reader,
                    source_stage,
                    reader_auth_token.as_deref(),
                    &mut payload_buf,
                ) {
                    Ok(Some(batch)) => {
                        total_read_ms += read_start.elapsed().as_secs_f32() * 1000.0;
                        if output_tx.send(batch).is_err() {
                            break Ok(());
                        }
                        read_batches += 1;
                    }
                    Ok(None) => break Ok(()),
                    Err(_) => break Err(()),
                }
            }
        };

        let (write_batches, total_write_ms) = writer.join().unwrap_or((0, 0.0));
        BridgeAccumulator {
            source_stage,
            processed_batches: read_batches.max(write_batches),
            total_write_ms,
            total_read_ms,
        }
    })
}

#[derive(Debug)]
pub struct BridgeAccumulator {
    pub source_stage: usize,
    pub processed_batches: usize,
    pub total_write_ms: f32,
    pub total_read_ms: f32,
}

impl BridgeAccumulator {
    fn default_with_stage(source_stage: usize) -> Self {
        Self {
            source_stage,
            processed_batches: 0,
            total_write_ms: 0.0,
            total_read_ms: 0.0,
        }
    }
}

/// Execute pipeline stages distributed across the LAN using TCP transport.
pub fn execute_pipeline_distributed(
    plan: &PipelinePlan,
    token_count: usize,
    micro_batch: usize,
    config: TcpTransportConfig,
    cluster: &crate::cluster::ClusterState,
    placement_context: Option<&crate::planning::PlacementPlan>,
    rebalance: Option<&crate::planning::RebalanceTrigger>,
) -> Result<ExecutionResult, String> {
    let stage_count = plan.stages.len();
    let micro_batch = micro_batch.max(1);
    let batch_count = token_count.div_ceil(micro_batch);
    if stage_count == 0 || token_count == 0 {
        return Ok(ExecutionResult {
            token_count,
            micro_batch,
            batch_count,
            stage_count,
            total_time_ms: 0.0,
            throughput_tokens_per_sec: 0.0,
            avg_token_latency_ms: 0.0,
            p95_token_latency_ms: 0.0,
            p99_token_latency_ms: 0.0,
            token_latencies_ms: Vec::new(),
            stage_stats: Vec::new(),
        });
    }

    #[derive(Debug)]
    struct StageAccumulator {
        stage_idx: usize,
        processed_batches: usize,
        total_compute_ms: f32,
        total_recv_wait_ms: f32,
        total_send_wait_ms: f32,
    }

    let (entry_tx, entry_rx) = mpsc::channel::<TransportBatch>();
    let mut stage_inputs: Vec<Option<mpsc::Receiver<TransportBatch>>> =
        Vec::with_capacity(stage_count);
    stage_inputs.push(Some(entry_rx));
    let mut stage_outputs = Vec::with_capacity(stage_count);
    let mut bridge_handles = Vec::new();
    let bridge_capacity = config.max_inflight_batches.max(1);

    for source_stage_idx in 0..stage_count.saturating_sub(1) {
        let (bridge_in_tx, bridge_in_rx) = mpsc::sync_channel::<TransportBatch>(bridge_capacity);
        let (bridge_out_tx, bridge_out_rx) = mpsc::sync_channel::<TransportBatch>(bridge_capacity);
        stage_outputs.push(bridge_in_tx);
        stage_inputs.push(Some(bridge_out_rx));

        let source_stage_p = &plan.stages[source_stage_idx];
        let target_stage_p = &plan.stages[source_stage_idx + 1];
        let same_node = source_stage_p.node_id == target_stage_p.node_id;

        let (listener, addr) = if same_node && config.transport == TransportKind::Unix {
            #[cfg(unix)]
            {
                // Same-node with Unix transport: use Unix domain socket.
                let sock_path = unique_bridge_socket_path(source_stage_idx);
                let _ = std::fs::remove_file(&sock_path);
                let ulistener = UnixListener::bind(&sock_path).map_err(|e| {
                    format!("Failed to bind Unix socket at {}: {e}", sock_path.display())
                })?;
                let addr = BridgeAddr::Unix(sock_path);
                (BridgeListener::Unix(ulistener), addr)
            }
            #[cfg(not(unix))]
            return Err("Unix domain sockets require a Unix platform (Linux/macOS)".to_string());
        } else {
            // TCP transport for cross-node or when configured explicitly.
            let bind_ip = if let Some(m) = cluster.get_metrics(&target_stage_p.node_id) {
                m.ip_address
                    .map(|sa| sa.ip())
                    .unwrap_or_else(|| [127, 0, 0, 1].into())
            } else {
                [127, 0, 0, 1].into()
            };

            let tcp_listener = TcpListener::bind(SocketAddr::new(bind_ip, 0))
                .map_err(|e| format!("Failed to bind TCP listener on {bind_ip}: {e}"))?;
            let actual_addr = tcp_listener
                .local_addr()
                .map_err(|e| format!("failed to get TCP local addr: {e}"))?;

            // For connection, if it's the same node, use loopback.
            let connect_ip = if same_node {
                [127, 0, 0, 1].into()
            } else {
                actual_addr.ip()
            };
            let connect_addr = SocketAddr::new(connect_ip, actual_addr.port());
            (
                BridgeListener::Tcp(tcp_listener),
                BridgeAddr::Tcp(connect_addr),
            )
        };

        let breaker = Some(cluster.circuit_breaker_for(&target_stage_p.node_id));
        bridge_handles.push(spawn_tcp_bridge(
            source_stage_idx,
            bridge_in_rx,
            bridge_out_tx,
            config.clone(),
            listener,
            addr,
            breaker,
        ));
    }

    let (completion_tx, completion_rx) = mpsc::sync_channel::<TransportBatch>(bridge_capacity);
    stage_outputs.push(completion_tx);

    let mut stage_handles = Vec::with_capacity(stage_count);
    for stage_idx in 0..stage_count {
        let stage = plan.stages[stage_idx].clone();
        let rx = stage_inputs[stage_idx]
            .take()
            .expect("stage receiver should exist");
        let tx_next = stage_outputs[stage_idx].clone();

        let handle = thread::spawn(move || {
            let mut processed_batches = 0usize;
            let mut total_compute_ms = 0.0_f32;
            let mut total_recv_wait_ms = 0.0_f32;
            let mut total_send_wait_ms = 0.0_f32;

            loop {
                let recv_start = Instant::now();
                let Ok(mut batch) = rx.recv() else {
                    break;
                };
                total_recv_wait_ms += recv_start.elapsed().as_secs_f32() * 1000.0;

                let compute_start = Instant::now();
                run_stage_compute(&mut batch.payload, &stage);
                total_compute_ms += compute_start.elapsed().as_secs_f32() * 1000.0;

                let send_start = Instant::now();
                if tx_next.send(batch).is_err() {
                    break;
                }
                total_send_wait_ms += send_start.elapsed().as_secs_f32() * 1000.0;
                processed_batches += 1;
            }

            StageAccumulator {
                stage_idx,
                processed_batches,
                total_compute_ms,
                total_recv_wait_ms,
                total_send_wait_ms,
            }
        });

        stage_handles.push(handle);
    }

    drop(stage_outputs);

    let mut batch_started_at = vec![Instant::now(); batch_count];
    let exec_start = Instant::now();
    for (batch_idx, batch_started_slot) in batch_started_at.iter_mut().enumerate() {
        let batch_start_token = batch_idx * micro_batch;
        let tokens_in_batch = (token_count - batch_start_token).min(micro_batch);
        let payload_len = (tokens_in_batch.max(1) * config.elems_per_token.max(1)).max(32);
        let payload = (0..payload_len)
            .map(|idx| (batch_idx as f32 * 0.01) + (idx as f32 * 0.0001))
            .collect();

        *batch_started_slot = Instant::now();
        let _ = entry_tx.send(TransportBatch {
            batch_id: batch_idx,
            tokens_in_batch,
            payload,
        });
    }
    drop(entry_tx);

    let mut token_latencies = Vec::with_capacity(token_count);
    for _ in 0..batch_count {
        let Ok(done_batch) = completion_rx.recv() else {
            break;
        };
        let started_at = batch_started_at[done_batch.batch_id];
        let batch_latency_ms = started_at.elapsed().as_secs_f32() * 1000.0;
        for _ in 0..done_batch.tokens_in_batch {
            token_latencies.push(batch_latency_ms);
        }
    }

    // Evaluate dynamic rebalance trigger during execution
    if let (Some(trigger), Some(placement)) = (rebalance, placement_context) {
        if let Some(migration) = trigger.evaluate(cluster, placement) {
            tracing::info!(
                "Distributed Runtime: Rebalance Triggered! Moving layers {:?} from {} to {}",
                migration.layers,
                migration.source_node,
                migration.target_node
            );
        }
    }
    let total_time_ms = exec_start.elapsed().as_secs_f32() * 1000.0;

    let mut stage_stats = Vec::with_capacity(stage_count);
    for handle in stage_handles {
        if let Ok(stats) = handle.join() {
            let divisor = stats.processed_batches.max(1) as f32;
            let avg_compute_ms = stats.total_compute_ms / divisor;

            // Feedback loop: update cluster state with actual compute latency.
            if let Some(stage) = plan.stages.get(stats.stage_idx) {
                cluster.get_metrics_mut(&stage.node_id, |m| {
                    m.record_latency(avg_compute_ms * 1000.0);
                });
            }

            stage_stats.push(StageExecutionStats {
                stage_idx: stats.stage_idx,
                processed_batches: stats.processed_batches,
                avg_compute_ms,
                avg_recv_wait_ms: stats.total_recv_wait_ms / divisor,
                avg_send_wait_ms: stats.total_send_wait_ms / divisor,
                avg_bridge_write_ms: 0.0,
                avg_bridge_read_ms: 0.0,
            });
        }
    }

    let mut bridge_stats = Vec::with_capacity(bridge_handles.len());
    for handle in bridge_handles {
        if let Ok(stats) = handle.join() {
            bridge_stats.push(stats);
        }
    }

    for bridge in bridge_stats {
        if let Some(stage) = stage_stats
            .iter_mut()
            .find(|s| s.stage_idx == bridge.source_stage)
        {
            let divisor = bridge.processed_batches.max(1) as f32;
            stage.avg_bridge_write_ms = bridge.total_write_ms / divisor;
            stage.avg_bridge_read_ms = bridge.total_read_ms / divisor;
        }
    }

    stage_stats.sort_by_key(|s| s.stage_idx);

    let throughput_tokens_per_sec = if total_time_ms > 0.0 {
        token_count as f32 / (total_time_ms / 1000.0)
    } else {
        0.0
    };

    let avg_token_latency_ms = if token_latencies.is_empty() {
        0.0
    } else {
        token_latencies.iter().sum::<f32>() / token_latencies.len() as f32
    };

    let mut sorted = token_latencies.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let p95_idx = ((sorted.len().saturating_sub(1)) as f32 * 0.95).round() as usize;
    let p95_token_latency_ms = sorted.get(p95_idx).copied().unwrap_or(0.0);
    let p99_idx = ((sorted.len().saturating_sub(1)) as f32 * 0.99).round() as usize;
    let p99_token_latency_ms = sorted.get(p99_idx).copied().unwrap_or(0.0);

    Ok(ExecutionResult {
        token_count,
        micro_batch,
        batch_count,
        stage_count,
        total_time_ms,
        throughput_tokens_per_sec,
        avg_token_latency_ms,
        p95_token_latency_ms,
        p99_token_latency_ms,
        token_latencies_ms: token_latencies,
        stage_stats,
    })
}

pub fn execute_pipeline_tcp_loopback(
    plan: &PipelinePlan,
    token_count: usize,
    micro_batch: usize,
) -> Result<ExecutionResult, String> {
    execute_pipeline_tcp_loopback_with_config(
        plan,
        token_count,
        micro_batch,
        TcpTransportConfig::default(),
    )
}

/// Execute pipeline stages with real TCP loopback transport bridges and explicit transport config.
pub fn execute_pipeline_tcp_loopback_with_config(
    plan: &PipelinePlan,
    token_count: usize,
    micro_batch: usize,
    config: TcpTransportConfig,
) -> Result<ExecutionResult, String> {
    let stage_count = plan.stages.len();
    let micro_batch = micro_batch.max(1);
    let batch_count = token_count.div_ceil(micro_batch);
    if stage_count == 0 || token_count == 0 {
        return Ok(ExecutionResult {
            token_count,
            micro_batch,
            batch_count,
            stage_count,
            total_time_ms: 0.0,
            throughput_tokens_per_sec: 0.0,
            avg_token_latency_ms: 0.0,
            p95_token_latency_ms: 0.0,
            p99_token_latency_ms: 0.0,
            token_latencies_ms: Vec::new(),
            stage_stats: Vec::new(),
        });
    }

    #[derive(Debug)]
    struct StageAccumulator {
        stage_idx: usize,
        processed_batches: usize,
        total_compute_ms: f32,
        total_recv_wait_ms: f32,
        total_send_wait_ms: f32,
    }

    let (entry_tx, entry_rx) = mpsc::channel::<TransportBatch>();
    let mut stage_inputs: Vec<Option<mpsc::Receiver<TransportBatch>>> =
        Vec::with_capacity(stage_count);
    stage_inputs.push(Some(entry_rx));
    let mut stage_outputs = Vec::with_capacity(stage_count);
    let mut bridge_handles = Vec::new();
    let bridge_capacity = config.max_inflight_batches.max(1);

    for source_stage in 0..stage_count.saturating_sub(1) {
        let (bridge_in_tx, bridge_in_rx) = mpsc::sync_channel::<TransportBatch>(bridge_capacity);
        let (bridge_out_tx, bridge_out_rx) = mpsc::sync_channel::<TransportBatch>(bridge_capacity);
        stage_outputs.push(bridge_in_tx);
        stage_inputs.push(Some(bridge_out_rx));

        let (listener, addr) = match config.transport {
            TransportKind::Tcp => {
                // For loopback execution, use distinct ports for each inter-stage bridge.
                let loopback = [127, 0, 0, 1].into();
                let bind_addr = SocketAddr::new(loopback, 0);
                let listener = TcpListener::bind(bind_addr)
                    .map_err(|e| format!("failed to bind TCP loopback listener: {e}"))?;
                let actual_addr = listener
                    .local_addr()
                    .map_err(|e| format!("failed to get TCP loopback addr: {e}"))?;
                (BridgeListener::Tcp(listener), BridgeAddr::Tcp(actual_addr))
            }
            #[cfg(unix)]
            TransportKind::Unix => {
                let sock_path = unique_bridge_socket_path(source_stage);
                let _ = std::fs::remove_file(&sock_path);
                let listener = UnixListener::bind(&sock_path).map_err(|e| {
                    format!("failed to bind Unix socket at {}: {e}", sock_path.display())
                })?;
                let addr = BridgeAddr::Unix(sock_path);
                (BridgeListener::Unix(listener), addr)
            }
            #[cfg(not(unix))]
            TransportKind::Unix => {
                return Err("Unix domain sockets require a Unix platform (Linux/macOS)".to_string());
            }
        };

        bridge_handles.push(spawn_tcp_bridge(
            source_stage,
            bridge_in_rx,
            bridge_out_tx,
            config.clone(),
            listener,
            addr,
            None, // loopback benchmarking: no real remote-node failure mode to track
        ));
    }

    let (completion_tx, completion_rx) = mpsc::sync_channel::<TransportBatch>(bridge_capacity);
    stage_outputs.push(completion_tx);

    let mut stage_handles = Vec::with_capacity(stage_count);
    for stage_idx in 0..stage_count {
        let stage = plan.stages[stage_idx].clone();
        let rx = stage_inputs[stage_idx]
            .take()
            .expect("stage receiver should exist");
        let tx_next = stage_outputs[stage_idx].clone();

        let handle = thread::spawn(move || {
            let mut processed_batches = 0usize;
            let mut total_compute_ms = 0.0_f32;
            let mut total_recv_wait_ms = 0.0_f32;
            let mut total_send_wait_ms = 0.0_f32;

            loop {
                let recv_start = Instant::now();
                let Ok(mut batch) = rx.recv() else {
                    break;
                };
                total_recv_wait_ms += recv_start.elapsed().as_secs_f32() * 1000.0;

                let compute_start = Instant::now();
                run_stage_compute(&mut batch.payload, &stage);
                total_compute_ms += compute_start.elapsed().as_secs_f32() * 1000.0;

                let send_start = Instant::now();
                if tx_next.send(batch).is_err() {
                    break;
                }
                total_send_wait_ms += send_start.elapsed().as_secs_f32() * 1000.0;
                processed_batches += 1;
            }

            StageAccumulator {
                stage_idx,
                processed_batches,
                total_compute_ms,
                total_recv_wait_ms,
                total_send_wait_ms,
            }
        });

        stage_handles.push(handle);
    }

    // Drop setup-time sender clones so channel closure can propagate and
    // stage/bridge threads can terminate once work is drained.
    drop(stage_outputs);

    let mut batch_started_at = vec![Instant::now(); batch_count];
    let exec_start = Instant::now();
    for (batch_idx, batch_started_slot) in batch_started_at.iter_mut().enumerate() {
        let batch_start_token = batch_idx * micro_batch;
        let tokens_in_batch = (token_count - batch_start_token).min(micro_batch);
        let payload_len = (tokens_in_batch.max(1) * config.elems_per_token.max(1)).max(32);
        let payload = (0..payload_len)
            .map(|idx| (batch_idx as f32 * 0.01) + (idx as f32 * 0.0001))
            .collect();

        *batch_started_slot = Instant::now();
        if entry_tx
            .send(TransportBatch {
                batch_id: batch_idx,
                tokens_in_batch,
                payload,
            })
            .is_err()
        {
            break;
        }
    }
    drop(entry_tx);

    let mut token_latencies = Vec::with_capacity(token_count);
    for _ in 0..batch_count {
        let Ok(done_batch) = completion_rx.recv() else {
            break;
        };

        let started_at = batch_started_at
            .get(done_batch.batch_id)
            .copied()
            .unwrap_or(exec_start);
        let batch_latency_ms = started_at.elapsed().as_secs_f32() * 1000.0;
        for _ in 0..done_batch.tokens_in_batch {
            token_latencies.push(batch_latency_ms);
        }
    }
    let total_time_ms = exec_start.elapsed().as_secs_f32() * 1000.0;

    let mut stage_stats = Vec::with_capacity(stage_count);
    for handle in stage_handles {
        if let Ok(stats) = handle.join() {
            let divisor = stats.processed_batches.max(1) as f32;
            stage_stats.push(StageExecutionStats {
                stage_idx: stats.stage_idx,
                processed_batches: stats.processed_batches,
                avg_compute_ms: stats.total_compute_ms / divisor,
                avg_recv_wait_ms: stats.total_recv_wait_ms / divisor,
                avg_send_wait_ms: stats.total_send_wait_ms / divisor,
                avg_bridge_write_ms: 0.0,
                avg_bridge_read_ms: 0.0,
            });
        }
    }

    let mut bridge_stats = Vec::with_capacity(bridge_handles.len());
    for handle in bridge_handles {
        if let Ok(stats) = handle.join() {
            bridge_stats.push(stats);
        }
    }

    for bridge in bridge_stats {
        if let Some(stage) = stage_stats
            .iter_mut()
            .find(|s| s.stage_idx == bridge.source_stage)
        {
            let divisor = bridge.processed_batches.max(1) as f32;
            stage.avg_bridge_write_ms = bridge.total_write_ms / divisor;
            stage.avg_bridge_read_ms = bridge.total_read_ms / divisor;
        }
    }

    stage_stats.sort_by_key(|s| s.stage_idx);

    let throughput_tokens_per_sec = if total_time_ms > 0.0 {
        token_count as f32 / (total_time_ms / 1000.0)
    } else {
        0.0
    };

    let avg_token_latency_ms = if token_latencies.is_empty() {
        0.0
    } else {
        token_latencies.iter().sum::<f32>() / token_latencies.len() as f32
    };

    let mut sorted = token_latencies.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let p95_idx = ((sorted.len().saturating_sub(1)) as f32 * 0.95).round() as usize;
    let p95_token_latency_ms = sorted.get(p95_idx).copied().unwrap_or(0.0);
    let p99_idx = ((sorted.len().saturating_sub(1)) as f32 * 0.99).round() as usize;
    let p99_token_latency_ms = sorted.get(p99_idx).copied().unwrap_or(0.0);

    Ok(ExecutionResult {
        token_count,
        micro_batch,
        batch_count,
        stage_count,
        total_time_ms,
        throughput_tokens_per_sec,
        avg_token_latency_ms,
        p95_token_latency_ms,
        p99_token_latency_ms,
        token_latencies_ms: token_latencies,
        stage_stats,
    })
}

// --- Real remote-worker execution -----------------------------------------
//
// Everything above this point (`execute_pipeline_distributed`,
// `execute_pipeline_tcp_loopback_with_config`, ...) runs entirely within one
// process: every stage's TCP bridge both binds *and* connects to a local
// interface, so even in "TCP transport" mode there is no way to reach a
// genuinely separate machine — `bind()` only ever succeeds for addresses
// this host actually owns. The two functions below are the real thing: a
// worker process that binds and accepts on a second machine, and a
// coordinator function that *connects out* to it (`TcpStream::connect`, not
// `TcpListener::bind`) instead of simulating both ends locally.
//
// Scope is intentionally narrow (exactly one local stage + one remote
// stage, one connection, one-shot) to match `flow`'s existing local+remote
// 2-node model rather than building a general N-worker mesh. The "compute"
// on both ends is still `run_stage_compute`'s synthetic proxy workload
// (real CPU time, not real LLM inference) — what's genuinely new here is
// the *transport*: a real socket between two real OS processes, optionally
// on two real machines, not the numbers `run_stage_compute` produces.

/// A minimal wire handshake sent once when a coordinator connects to a
/// `run_stage_worker`: just enough for the worker to reconstruct the
/// `StagePlacement` it needs to call the same `run_stage_compute` the local
/// path uses, so timing is produced by identical code on both ends.
fn write_stage_handshake(writer: &mut impl Write, stage: &StagePlacement) -> Result<(), String> {
    let num_layers = stage.num_layers() as u32;
    let device_byte: u8 = match stage.device {
        DeviceKind::Npu => 0,
        DeviceKind::Gpu => 1,
        DeviceKind::Cpu => 2,
    };
    writer
        .write_all(&num_layers.to_le_bytes())
        .and_then(|_| writer.write_all(&[device_byte]))
        .map_err(|e| format!("failed to write stage handshake: {e}"))
}

fn read_stage_handshake(reader: &mut impl Read) -> Result<StagePlacement, String> {
    let mut layers_bytes = [0u8; 4];
    reader
        .read_exact(&mut layers_bytes)
        .map_err(|e| format!("failed to read stage handshake: {e}"))?;
    let num_layers = u32::from_le_bytes(layers_bytes) as usize;

    let mut device_byte = [0u8; 1];
    reader
        .read_exact(&mut device_byte)
        .map_err(|e| format!("failed to read stage handshake: {e}"))?;
    let device = match device_byte[0] {
        0 => DeviceKind::Npu,
        1 => DeviceKind::Gpu,
        _ => DeviceKind::Cpu,
    };

    Ok(StagePlacement {
        node_id: "remote".to_string(),
        start_layer: 0,
        end_layer: num_layers,
        device,
        est_latency_ms: 0.0,
    })
}

/// Summary a `stage-worker` process reports (to its own stdout, and back to
/// any caller driving it directly) once its one coordinator connection
/// closes.
#[derive(Clone, Debug, PartialEq)]
pub struct StageWorkerSummary {
    pub batches_processed: usize,
    pub avg_compute_ms: f32,
}

/// Runs as a real, separate OS process (see `ghost-link stage-worker`):
/// accepts exactly one coordinator connection on `listener`, reads the
/// stage handshake, then loops read-batch → compute → write-result until
/// the coordinator closes the connection. One-shot, matching how
/// `cluster-start`'s child `listen` processes already behave.
///
/// Takes an already-bound `TcpListener` rather than binding internally so
/// the caller controls exactly when/how the port is chosen — the CLI binds
/// the address the user passed on `--bind`, while tests bind `:0` to get a
/// real OS-assigned free port and can read it back via `local_addr()`
/// before this function (which blocks until a peer connects) is called.
pub fn run_stage_worker(
    listener: TcpListener,
    config: &TcpTransportConfig,
) -> Result<StageWorkerSummary, String> {
    let (mut stream, peer) = listener
        .accept()
        .map_err(|e| format!("failed to accept a connection: {e}"))?;
    let _ = stream.set_nodelay(true);

    let stage = read_stage_handshake(&mut stream)?;
    eprintln!(
        "[stage-worker] connection from {peer}; assigned {} layers on {}",
        stage.num_layers(),
        stage.device.as_str()
    );

    let mut frame_buf = Vec::new();
    let mut read_buf = Vec::new();
    let mut batches_processed = 0usize;
    let mut total_compute_ms = 0.0_f32;

    loop {
        let batch =
            match read_transport_batch(&mut stream, 0, config.auth_token.as_deref(), &mut read_buf)
            {
                Ok(Some(batch)) => batch,
                Ok(None) => break, // coordinator closed the connection cleanly
                Err(e) => return Err(format!("failed to read batch from coordinator: {e}")),
            };

        let mut payload = batch.payload;
        let compute_start = Instant::now();
        run_stage_compute(&mut payload, &stage);
        total_compute_ms += compute_start.elapsed().as_secs_f32() * 1000.0;

        let result = TransportBatch {
            batch_id: batch.batch_id,
            tokens_in_batch: batch.tokens_in_batch,
            payload,
        };
        write_transport_batch(
            &mut stream,
            &result,
            0,
            config.auth_token.as_deref(),
            &mut frame_buf,
        )
        .map_err(|e| format!("failed to write result batch to coordinator: {e}"))?;
        batches_processed += 1;
    }

    let avg_compute_ms = if batches_processed > 0 {
        total_compute_ms / batches_processed as f32
    } else {
        0.0
    };
    eprintln!(
        "[stage-worker] coordinator disconnected after {batches_processed} batch(es), avg compute {avg_compute_ms:.2} ms"
    );
    Ok(StageWorkerSummary {
        batches_processed,
        avg_compute_ms,
    })
}

/// Collapses every stage assigned to one node into a single logical
/// `StagePlacement` spanning that node's full layer range. The real layer
/// assignment algorithm doesn't hand each node one contiguous chunk — a
/// 60-layer/2-node `flow` plan produces ~11 raw stages (verified: 6 small
/// chunks for the local node, 5 for the remote one, each capped at roughly
/// 3GB), not the 2 this function originally assumed. The chunks for a given
/// node ARE always contiguous relative to each other, so merging them into
/// one placement per node is behaviorally sound: `run_stage_compute`'s
/// synthetic cost model only depends on total layer count and device kind,
/// not on how finely a node's range happens to be subdivided.
fn merge_stages_for_node(
    node_id: &str,
    stages: &[StagePlacement],
) -> Result<StagePlacement, String> {
    let matching: Vec<&StagePlacement> = stages.iter().filter(|s| s.node_id == node_id).collect();
    let Some(first) = matching.first() else {
        return Err(format!("no pipeline stages assigned to node '{node_id}'"));
    };
    let start_layer = matching.iter().map(|s| s.start_layer).min().unwrap_or(0);
    let end_layer = matching.iter().map(|s| s.end_layer).max().unwrap_or(0);
    Ok(StagePlacement {
        node_id: node_id.to_string(),
        start_layer,
        end_layer,
        device: first.device,
        est_latency_ms: matching.iter().map(|s| s.est_latency_ms).sum(),
    })
}

/// Coordinator side of real remote execution. `plan` must have at least one
/// stage assigned to `local_node_id` and at least one assigned to
/// `remote_node_id` — matching `flow`'s existing local+remote 2-*node*
/// model, though each node's stages get merged (see `merge_stages_for_node`)
/// since the real assignment algorithm splits one node's layers into
/// several raw stages, not necessarily one. The merged local stage runs
/// in-process as usual; the merged remote stage's batches are sent over a
/// real outbound connection (`TcpStream::connect(remote_addr)`) to a
/// separate `run_stage_worker` process, so the round-trip time recorded
/// here is genuine network + remote-compute time, not a local simulation.
pub fn execute_pipeline_with_remote_stage(
    plan: &PipelinePlan,
    local_node_id: &str,
    remote_node_id: &str,
    token_count: usize,
    micro_batch: usize,
    config: TcpTransportConfig,
    remote_addr: SocketAddr,
) -> Result<ExecutionResult, String> {
    let local_stage = merge_stages_for_node(local_node_id, &plan.stages)?;
    let remote_stage = merge_stages_for_node(remote_node_id, &plan.stages)?;
    let local_stage = &local_stage;
    let remote_stage = &remote_stage;

    let micro_batch = micro_batch.max(1);
    let batch_count = token_count.div_ceil(micro_batch);
    if token_count == 0 {
        return Ok(ExecutionResult {
            token_count,
            micro_batch,
            batch_count,
            stage_count: 2,
            total_time_ms: 0.0,
            throughput_tokens_per_sec: 0.0,
            avg_token_latency_ms: 0.0,
            p95_token_latency_ms: 0.0,
            p99_token_latency_ms: 0.0,
            token_latencies_ms: Vec::new(),
            stage_stats: Vec::new(),
        });
    }

    let mut stream = TcpStream::connect(remote_addr).map_err(|e| {
        format!(
            "failed to connect to remote stage worker at {remote_addr}: {e} \
             (is `ghost-link stage-worker --bind {remote_addr}` running there?)"
        )
    })?;
    let _ = stream.set_nodelay(true);
    write_stage_handshake(&mut stream, remote_stage)?;

    let mut frame_buf = Vec::new();
    let mut read_buf = Vec::new();
    let mut token_latencies = Vec::with_capacity(token_count);
    let mut local_compute_ms = 0.0_f32;
    let mut remote_round_trip_ms = 0.0_f32;

    let exec_start = Instant::now();
    for batch_idx in 0..batch_count {
        let batch_start_token = batch_idx * micro_batch;
        let tokens_in_batch = (token_count - batch_start_token).min(micro_batch);
        let payload_len = (tokens_in_batch.max(1) * config.elems_per_token.max(1)).max(32);
        let mut payload: Vec<f32> = (0..payload_len)
            .map(|idx| (batch_idx as f32 * 0.01) + (idx as f32 * 0.0001))
            .collect();

        let batch_started = Instant::now();

        let local_start = Instant::now();
        run_stage_compute(&mut payload, local_stage);
        local_compute_ms += local_start.elapsed().as_secs_f32() * 1000.0;

        let outgoing = TransportBatch {
            batch_id: batch_idx,
            tokens_in_batch,
            payload,
        };
        let rt_start = Instant::now();
        write_transport_batch(
            &mut stream,
            &outgoing,
            0,
            config.auth_token.as_deref(),
            &mut frame_buf,
        )
        .map_err(|e| format!("failed to send batch {batch_idx} to remote worker: {e}"))?;
        let result = read_transport_batch(&mut stream, 0, config.auth_token.as_deref(), &mut read_buf)
            .map_err(|e| format!("failed to read batch {batch_idx} result from remote worker: {e}"))?
            .ok_or_else(|| {
                format!("remote worker at {remote_addr} closed the connection before batch {batch_idx} completed")
            })?;
        remote_round_trip_ms += rt_start.elapsed().as_secs_f32() * 1000.0;

        let batch_latency_ms = batch_started.elapsed().as_secs_f32() * 1000.0;
        for _ in 0..result.tokens_in_batch {
            token_latencies.push(batch_latency_ms);
        }
    }
    let total_time_ms = exec_start.elapsed().as_secs_f32() * 1000.0;
    drop(stream); // signals the worker's read loop to exit cleanly

    let divisor = batch_count.max(1) as f32;
    let throughput_tokens_per_sec = if total_time_ms > 0.0 {
        token_count as f32 / (total_time_ms / 1000.0)
    } else {
        0.0
    };
    let avg_token_latency_ms = if token_latencies.is_empty() {
        0.0
    } else {
        token_latencies.iter().sum::<f32>() / token_latencies.len() as f32
    };
    let mut sorted = token_latencies.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let p95_idx = (sorted.len().saturating_sub(1) as f32 * 0.95).round() as usize;
    let p95_token_latency_ms = sorted.get(p95_idx).copied().unwrap_or(0.0);
    let p99_idx = (sorted.len().saturating_sub(1) as f32 * 0.99).round() as usize;
    let p99_token_latency_ms = sorted.get(p99_idx).copied().unwrap_or(0.0);

    Ok(ExecutionResult {
        token_count,
        micro_batch,
        batch_count,
        stage_count: 2,
        total_time_ms,
        throughput_tokens_per_sec,
        avg_token_latency_ms,
        p95_token_latency_ms,
        p99_token_latency_ms,
        token_latencies_ms: token_latencies,
        stage_stats: vec![
            StageExecutionStats {
                stage_idx: 0,
                processed_batches: batch_count,
                avg_compute_ms: local_compute_ms / divisor,
                avg_recv_wait_ms: 0.0,
                avg_send_wait_ms: 0.0,
                avg_bridge_write_ms: 0.0,
                avg_bridge_read_ms: 0.0,
            },
            StageExecutionStats {
                stage_idx: 1,
                processed_batches: batch_count,
                // The remote side's own compute time isn't observable from
                // here without a second telemetry channel — what IS real
                // and observable is the round trip, reported as bridge
                // write time so it's visible in `ExecutionResult::summary()`
                // without inventing a number for compute we didn't measure.
                avg_compute_ms: 0.0,
                avg_recv_wait_ms: 0.0,
                avg_send_wait_ms: 0.0,
                avg_bridge_write_ms: remote_round_trip_ms / divisor,
                avg_bridge_read_ms: 0.0,
            },
        ],
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::planning::LayerAssignment;
    use std::collections::HashMap;
    use std::io::Cursor;

    #[test]
    fn pipeline_plan_uses_node_device_mapping() {
        let assignments = vec![
            LayerAssignment::new("node-a".to_string(), 0, 10, 5.0),
            LayerAssignment::new("node-b".to_string(), 10, 20, 5.0),
        ];

        let mut devices = HashMap::new();
        devices.insert("node-a".to_string(), DeviceKind::Gpu);
        devices.insert("node-b".to_string(), DeviceKind::Cpu);

        let plan = PipelinePlan::from_assignments(&assignments, &devices);

        assert_eq!(plan.stages.len(), 2);
        assert_eq!(plan.stages[0].device, DeviceKind::Gpu);
        assert_eq!(plan.stages[1].device, DeviceKind::Cpu);
        assert!(plan.est_step_latency_ms() > 0.0);
    }

    #[test]
    fn schedule_size_matches_tokens_times_stages() {
        let schedule = build_token_schedule(3, 4);
        assert_eq!(schedule.len(), 12);
        assert_eq!(
            schedule[0],
            TokenStep {
                token_idx: 0,
                stage_idx: 0
            }
        );
        assert_eq!(
            schedule[11],
            TokenStep {
                token_idx: 3,
                stage_idx: 2
            }
        );
    }

    #[test]
    fn execution_reports_basic_stage_metrics() {
        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "node-a".to_string(),
                    start_layer: 0,
                    end_layer: 10,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.0,
                },
                StagePlacement {
                    node_id: "node-b".to_string(),
                    start_layer: 10,
                    end_layer: 20,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.0,
                },
            ],
        };

        let result = execute_pipeline(&plan, 8, 1).expect("failed");
        assert_eq!(result.stage_count, 2);
        assert_eq!(result.token_count, 8);
        assert_eq!(result.micro_batch, 1);
        assert_eq!(result.batch_count, 8);
        assert_eq!(result.stage_stats.len(), 2);
        assert!(result.total_time_ms > 0.0);
        assert!(result.avg_token_latency_ms > 0.0);
        assert!(result.throughput_tokens_per_sec > 0.0);
    }

    #[test]
    fn execution_respects_micro_batch_count() {
        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "node-a".to_string(),
                    start_layer: 0,
                    end_layer: 10,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.0,
                },
                StagePlacement {
                    node_id: "node-b".to_string(),
                    start_layer: 10,
                    end_layer: 20,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.0,
                },
            ],
        };

        let mb1 = execute_pipeline(&plan, 32, 1).expect("failed");
        let mb4 = execute_pipeline(&plan, 32, 4).expect("failed");

        assert_eq!(mb1.batch_count, 32);
        assert_eq!(mb4.batch_count, 8);
        assert!(mb4.total_time_ms > 0.0);
    }

    #[test]
    fn execution_supports_npu_stage_metrics() {
        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "node-npu".to_string(),
                    start_layer: 0,
                    end_layer: 8,
                    device: DeviceKind::Npu,
                    est_latency_ms: 0.9,
                },
                StagePlacement {
                    node_id: "node-gpu".to_string(),
                    start_layer: 8,
                    end_layer: 16,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.1,
                },
            ],
        };

        let result = execute_pipeline(&plan, 12, 3).expect("failed");
        assert_eq!(result.stage_count, 2);
        assert_eq!(result.batch_count, 4);
        assert_eq!(result.stage_stats.len(), 2);
        assert!(result.stage_stats.iter().any(|s| s.stage_idx == 0));
        assert!(result.total_time_ms > 0.0);
        assert!(result.throughput_tokens_per_sec > 0.0);
    }

    #[test]
    fn tcp_loopback_execution_reports_metrics() {
        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "node-a".to_string(),
                    start_layer: 0,
                    end_layer: 10,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.0,
                },
                StagePlacement {
                    node_id: "node-b".to_string(),
                    start_layer: 10,
                    end_layer: 20,
                    device: DeviceKind::Cpu,
                    est_latency_ms: 2.0,
                },
            ],
        };

        let result = execute_pipeline_tcp_loopback(&plan, 16, 2).expect("failed");
        assert_eq!(result.token_count, 16);
        assert_eq!(result.batch_count, 8);
        assert_eq!(result.stage_stats.len(), 2);
        assert!(result.total_time_ms > 0.0);
        assert!(result.throughput_tokens_per_sec > 0.0);
        assert!(
            result.p99_token_latency_ms >= result.p95_token_latency_ms,
            "p99 ({}) must be >= p95 ({})",
            result.p99_token_latency_ms,
            result.p95_token_latency_ms
        );
    }

    #[test]
    fn elems_per_token_config_scales_payload_size_and_still_completes() {
        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "node-a".to_string(),
                    start_layer: 0,
                    end_layer: 10,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.0,
                },
                StagePlacement {
                    node_id: "node-b".to_string(),
                    start_layer: 10,
                    end_layer: 20,
                    device: DeviceKind::Cpu,
                    est_latency_ms: 2.0,
                },
            ],
        };

        // 4096 hidden_dim * 2 bytes (fp16/bf16) / 4 bytes per wire f32 = 2048
        // elements/token — a real 7B-class model's per-token activation size,
        // instead of the default 16-element synthetic stand-in.
        let config = TcpTransportConfig {
            elems_per_token: 2048,
            ..TcpTransportConfig::default()
        };

        let result = execute_pipeline_tcp_loopback_with_config(&plan, 8, 2, config)
            .expect("large-payload loopback execution failed");
        assert_eq!(result.token_count, 8);
        assert!(result.total_time_ms > 0.0);
        assert!(
            result.p99_token_latency_ms >= result.p95_token_latency_ms,
            "p99 ({}) must be >= p95 ({})",
            result.p99_token_latency_ms,
            result.p95_token_latency_ms
        );
    }

    #[test]
    fn tcp_loopback_execution_with_hardening_config_reports_metrics() {
        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "node-a".to_string(),
                    start_layer: 0,
                    end_layer: 12,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.5,
                },
                StagePlacement {
                    node_id: "node-b".to_string(),
                    start_layer: 12,
                    end_layer: 24,
                    device: DeviceKind::Cpu,
                    est_latency_ms: 2.5,
                },
            ],
        };

        let result = execute_pipeline_tcp_loopback_with_config(
            &plan,
            24,
            3,
            TcpTransportConfig {
                max_inflight_batches: 4,
                reconnect_attempts: 4,
                reconnect_backoff_ms: 5,
                auth_token: Some("test-token".to_string()),
                ..Default::default()
            },
        )
        .expect("hardened loopback failed");

        assert_eq!(result.token_count, 24);
        assert_eq!(result.batch_count, 8);
        assert_eq!(result.stage_stats.len(), 2);
        assert!(result.total_time_ms > 0.0);
        assert!(result.throughput_tokens_per_sec > 0.0);
    }

    #[test]
    fn transport_rejects_auth_mismatch_frames() {
        let batch = TransportBatch {
            batch_id: 7,
            tokens_in_batch: 4,
            payload: vec![0.1, 0.2, 0.3, 0.4],
        };
        let mut encoded = Vec::new();
        let mut frame_buf = Vec::new();
        write_transport_batch(&mut encoded, &batch, 0, Some("token-a"), &mut frame_buf)
            .expect("encode frame");

        let mut cursor = Cursor::new(encoded);
        let mut payload_buf = Vec::new();
        let err = read_transport_batch(&mut cursor, 0, Some("token-b"), &mut payload_buf)
            .expect_err("mismatched token should fail");
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn transport_rejects_unexpected_source_stage() {
        let batch = TransportBatch {
            batch_id: 2,
            tokens_in_batch: 2,
            payload: vec![1.0, 2.0],
        };
        let mut encoded = Vec::new();
        let mut frame_buf = Vec::new();
        write_transport_batch(&mut encoded, &batch, 1, Some("token"), &mut frame_buf)
            .expect("encode frame");

        let mut cursor = Cursor::new(encoded);
        let mut payload_buf = Vec::new();
        let err = read_transport_batch(&mut cursor, 0, Some("token"), &mut payload_buf)
            .expect_err("source-stage mismatch should fail");
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[cfg(unix)]
    #[test]
    fn unix_socket_loopback_execution_reports_metrics() {
        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "node-a".to_string(),
                    start_layer: 0,
                    end_layer: 10,
                    device: DeviceKind::Gpu,
                    est_latency_ms: 1.0,
                },
                StagePlacement {
                    node_id: "node-b".to_string(),
                    start_layer: 10,
                    end_layer: 20,
                    device: DeviceKind::Cpu,
                    est_latency_ms: 2.0,
                },
            ],
        };

        let result = execute_pipeline_tcp_loopback_with_config(
            &plan,
            16,
            2,
            TcpTransportConfig {
                transport: TransportKind::Unix,
                ..Default::default()
            },
        )
        .expect("unix socket loopback failed");

        assert_eq!(result.token_count, 16);
        assert_eq!(result.batch_count, 8);
        assert_eq!(result.stage_stats.len(), 2);
        assert!(result.total_time_ms > 0.0);
        assert!(result.throughput_tokens_per_sec > 0.0);
    }

    #[test]
    fn stage_handshake_round_trips_layer_count_and_device() {
        let stage = StagePlacement {
            node_id: "coordinator-local".to_string(),
            start_layer: 5,
            end_layer: 37,
            device: DeviceKind::Gpu,
            est_latency_ms: 12.0,
        };
        let mut buf = Vec::new();
        write_stage_handshake(&mut buf, &stage).expect("handshake write failed");

        let mut cursor = Cursor::new(buf);
        let decoded = read_stage_handshake(&mut cursor).expect("handshake read failed");
        assert_eq!(decoded.num_layers(), stage.num_layers());
        assert_eq!(decoded.device, DeviceKind::Gpu);
    }

    #[test]
    fn execute_pipeline_with_remote_stage_rejects_when_remote_node_has_no_stages() {
        let plan = PipelinePlan {
            stages: vec![StagePlacement {
                node_id: "solo".to_string(),
                start_layer: 0,
                end_layer: 10,
                device: DeviceKind::Cpu,
                est_latency_ms: 1.0,
            }],
        };
        let err = execute_pipeline_with_remote_stage(
            &plan,
            "solo",
            "nowhere",
            8,
            2,
            TcpTransportConfig::default(),
            "127.0.0.1:1".parse().unwrap(),
        )
        .unwrap_err();
        assert!(
            err.contains("no pipeline stages assigned to node 'nowhere'"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn merge_stages_for_node_collapses_multiple_chunks_into_one_placement() {
        // Regression: real layer assignment splits one node's range into
        // several raw stages (verified manually: 60 layers / 2 nodes
        // produced 11 raw stages, not 2) — merging must still produce the
        // node's FULL range, not just its first chunk.
        let stages = vec![
            StagePlacement {
                node_id: "node-remote".to_string(),
                start_layer: 32,
                end_layer: 38,
                device: DeviceKind::Gpu,
                est_latency_ms: 1.0,
            },
            StagePlacement {
                node_id: "node-remote".to_string(),
                start_layer: 38,
                end_layer: 44,
                device: DeviceKind::Gpu,
                est_latency_ms: 1.5,
            },
            StagePlacement {
                node_id: "node-local".to_string(),
                start_layer: 0,
                end_layer: 32,
                device: DeviceKind::Cpu,
                est_latency_ms: 3.0,
            },
        ];
        let merged = merge_stages_for_node("node-remote", &stages).expect("merge failed");
        assert_eq!(merged.start_layer, 32);
        assert_eq!(merged.end_layer, 44);
        assert_eq!(merged.num_layers(), 12);
        assert_eq!(merged.device, DeviceKind::Gpu);
        assert!((merged.est_latency_ms - 2.5).abs() < 1e-6);
    }

    /// The real thing, end to end: a genuine second OS thread (standing in
    /// for a genuine second OS *process*, which the CLI-level integration
    /// test in the ghost-link crate covers) binds a real TCP listener on an
    /// OS-assigned port, and the coordinator function connects out to it —
    /// not the local-bind-on-both-ends simulation every other function in
    /// this file uses. Proves the round trip is real by asserting non-zero
    /// bridge time was actually measured, and that the worker's own
    /// independently-tracked batch count agrees with the coordinator's.
    #[test]
    fn execute_pipeline_with_remote_stage_does_a_real_tcp_round_trip() {
        let listener = TcpListener::bind("127.0.0.1:0").expect("failed to bind test listener");
        let worker_addr = listener.local_addr().expect("failed to read local_addr");

        let worker_handle =
            thread::spawn(move || run_stage_worker(listener, &TcpTransportConfig::default()));

        let plan = PipelinePlan {
            stages: vec![
                StagePlacement {
                    node_id: "local".to_string(),
                    start_layer: 0,
                    end_layer: 10,
                    device: DeviceKind::Cpu,
                    est_latency_ms: 1.0,
                },
                StagePlacement {
                    node_id: "remote".to_string(),
                    start_layer: 10,
                    end_layer: 20,
                    device: DeviceKind::Cpu,
                    est_latency_ms: 1.0,
                },
            ],
        };

        let result = execute_pipeline_with_remote_stage(
            &plan,
            "local",
            "remote",
            16,
            2,
            TcpTransportConfig::default(),
            worker_addr,
        )
        .expect("real remote-stage execution failed");

        let worker_summary = worker_handle
            .join()
            .expect("worker thread panicked")
            .expect("worker returned an error");

        assert_eq!(result.token_count, 16);
        assert_eq!(result.batch_count, 8);
        assert_eq!(result.stage_stats.len(), 2);
        assert_eq!(worker_summary.batches_processed, 8);
        assert_eq!(
            result.stage_stats[1].processed_batches,
            worker_summary.batches_processed,
            "coordinator's batch count must agree with what the independent worker process actually processed"
        );
        assert!(
            result.stage_stats[1].avg_bridge_write_ms > 0.0,
            "remote round-trip time must be non-zero for a real socket hop, got {}",
            result.stage_stats[1].avg_bridge_write_ms
        );
    }

    #[test]
    fn spawn_tcp_bridge_skips_connection_attempts_when_circuit_breaker_is_open() {
        use crate::circuit_breaker::{CircuitBreakerConfig, CircuitState};

        // Bind-then-drop reliably yields an address nothing is listening on.
        let dead_addr = {
            let l = TcpListener::bind("127.0.0.1:0").unwrap();
            l.local_addr().unwrap()
        };
        let dummy_listener = BridgeListener::Tcp(TcpListener::bind("127.0.0.1:0").unwrap());

        let breaker = CircuitBreaker::with_config(CircuitBreakerConfig {
            failure_threshold: 1,
            open_duration_secs: 3600,
            reset_timeout_secs: 3600,
        });
        breaker.record_failure();
        assert_eq!(breaker.current_state(), CircuitState::Open);

        let (in_tx, in_rx) = mpsc::sync_channel::<TransportBatch>(1);
        let (out_tx, _out_rx) = mpsc::sync_channel::<TransportBatch>(1);
        drop(in_tx); // closes the passthrough loop's input immediately

        let config = TcpTransportConfig {
            reconnect_attempts: 5,
            reconnect_backoff_ms: 2000,
            ..TcpTransportConfig::default()
        };

        let start = Instant::now();
        let handle = spawn_tcp_bridge(
            0,
            in_rx,
            out_tx,
            config,
            dummy_listener,
            BridgeAddr::Tcp(dead_addr),
            Some(breaker),
        );
        handle.join().expect("bridge thread panicked");

        // 5 attempts * 2s backoff would take >=10s if the open breaker weren't
        // honored; a healthy short-circuit finishes in well under a second.
        assert!(
            start.elapsed() < Duration::from_millis(1500),
            "open circuit breaker should skip connection attempts entirely, took {:?}",
            start.elapsed()
        );
    }
}