inferencelayer 0.2.1

//! P2P pipeline-parallel sharding: pool VRAM across devices (Macs, V100s, anything wgpu reaches)
//! by giving each device a contiguous layer range and streaming the residual stream between them.
//!
//! Why pipeline- and not tensor-parallel: TP needs 2 all-reduces per layer per token — NVLink
//! territory; PP ships ONE hidden-state per stage boundary per token (`hidden·4` bytes ≈ 4 KB for
//! a 1 B model) and survives ordinary Ethernet (see docs/ai-sota/serving-engines-vllm-sglang-p2p.md
//! §4). The last stage returns the greedy TOKEN (4 bytes), not logits (600 KB).
//!
//! Topology: the coordinator owns stage 0 (embedding + first layers) plus the tokenizer and drives
//! the chain; each worker (`lfm2-shard-worker`) owns one layer range. Every stage holds only its
//! own layers' weights and KV — the VRAM pooling that lets N small devices serve a model none of
//! them could hold alone. Two transports share one wire protocol:
//!
//! - **Hub-and-spoke** (default): every stage answers the coordinator, which relays the residual
//!   stream to the next stage — each boundary crosses the coordinator's network twice.
//! - **TRUE P2P forward mode** (`OP_CHAIN`): each worker ships its output DIRECTLY to the next
//!   peer and forwards the in-band control ops (zslot/btab/dn-restore) down the chain, so
//!   per-stage ordering is exactly the hub-spoke semantics; only the LAST stage's tokens return,
//!   to the coordinator's sink listener. Inter-stage traffic never touches the coordinator —
//!   the shape that lets Macs, Linux boxes, and mixed GPUs pool VRAM over ordinary Ethernet
//!   without the coordinator's NIC becoming the bottleneck. One coordinator session at a time:
//!   a worker's engine is a single shared stage across its connections.
//!
//! Wire protocol (length-prefixed, little-endian):
//! ```text
//! request:  op:u32  a:u32  n:u32  payload            op 0 = step, 1 = reset, 2 = shutdown, …
//! response: tag:u32 n:u32  payload[n]:f32            tag 0 = hidden, 1 = token, 2 = info bytes
//! sink msg: tag:u32(=3) group:u32 n:u32 payload[n]   last stage → coordinator sink (P2P only)
//! ```
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::path::Path;

use anyhow::{Context, Result, anyhow};

use crate::forward::{BatchCol, StageBatchOut, StageOut};
use crate::{GpuCtx, Lfm2Gpu, Weights};

const OP_STEP: u32 = 0;
const OP_RESET: u32 = 1;
const OP_SHUTDOWN: u32 = 2;
/// Build the stage's batch plan for k columns (micro-batched serving). Acked.
const OP_BINIT: u32 = 3;
/// Overwrite one KV block-table row (in-band, unacked — TCP order sequences it before the next
/// step on this connection). Header (row, n), payload n block ids as f32 bit patterns.
const OP_BTAB: u32 = 4;
/// One batched stage step: header (k_used, n_hidden), then 4·k_used u32 column meta
/// (pos, btrow, need_logit, token), then the residual payload (EMPTY toward a stage_first
/// worker — it embeds the tokens itself, so the hop shrinks to the meta).
const OP_BSTEP: u32 = 5;
/// Zero a sequence slot's recurrent state (fresh admission). Header (slot, 0); unacked, in-band.
const OP_ZSLOT: u32 = 6;
/// [`OP_BSTEP`] on the WIDE (KC=16) prefill plan — same wire shape.
const OP_BSTEPW: u32 = 7;
/// MTP draft chain on the LAST stage: header = (op, pos0, col_j << 8 | k_draft), payload = one
/// f32 carrying the first token's bits. Response: TAG_TOKEN with k_draft drafted ids.
const OP_MTP: u32 = 8;
/// Roll a slot's DN state back to after column `col` of the last spec micro-batch (in-band,
/// no response — ordered ahead of the next step like OP_ZSLOT): header = (op, slot, col).
const OP_DNRESTORE: u32 = 9;
/// Batched MTP draft on the last stage: header (op, ncols, k); payload = ncols × (slot,
/// first_tok, pos0) u32-triples (bits-in-f32). Response: TAG_TOKEN, ncols·k drafted ids.
const OP_MTP_BATCH: u32 = 10;
/// Seed slot's draft-chain slab from verify column `col`: header (op, slot, col), no response.
const OP_MTP_SEED: u32 = 11;
/// Configure TRUE-P2P forward mode: header (op, flags, n_bytes), payload = next-hop `host:port`
/// UTF-8 (`":port"` ⇒ substitute this connection's peer IP — how the last stage reaches the
/// coordinator's sink without the coordinator knowing its own routable address; empty ⇒ clear).
/// flags bit0: the next hop is the coordinator's RETURN SINK. Acked: payload[0] bits 0 = ok,
/// 1 = dial failed (the session keeps its previous — cleared — chain).
const OP_CHAIN: u32 = 12;
/// Stage self-description: responds (TAG_INFO, n_bytes) +
/// `start|end|total|subgroups|vram_mb|pinned|ckpt_fp|backend` (start==end==0 ⇒ not loaded).
const OP_INFO: u32 = 13;
/// Load (or RELOAD) the stage's layer range: header (op, start, end); acked (0 ok / 1 fail).
/// Idempotent — a matching resident range is a no-op, so re-running a coordinator against an
/// unchanged fleet never re-reads checkpoints. The action half of auto-split (see `plan_split`).
const OP_LOAD: u32 = 14;
/// Connection auth: header (op, 0, n_bytes) + fleet-token bytes; acked (0 ok / 1 reject).
/// When the worker holds a token (`--token` / env `OSFKB_SHARD_TOKEN`), EVERY connection must
/// open with a valid hello — coordinator control, worker→worker hops, and the sink dial alike.
const OP_HELLO: u32 = 15;
/// Weight shipping, file open: header (op, 0, n_meta) + `fp16hex|start|end|filename` — the
/// worker starts `<cache>/<fp>/<start>-<end>/<filename>.part`. Unacked (END acks the file).
const OP_SHIP_BEGIN: u32 = 16;
/// Weight shipping, payload: header (op, 0, n_bytes) + raw bytes appended to the open file.
const OP_SHIP_CHUNK: u32 = 17;
/// Weight shipping, commit: header (op, fnv_lo, fnv_hi); the worker verifies the FNV64 of the
/// received bytes, renames `.part` into place (safetensors also records the SOURCE checkpoint
/// fingerprint as provenance), and acks 0 ok / 1 fail.
const OP_SHIP_END: u32 = 18;
/// Reverse registration (worker → coordinator registry, after any [`OP_HELLO`]): header
/// (op, data_listen_port, n_bytes) + optional advertised `host:port` for worker→worker hops
/// (empty ⇒ the registry derives source-IP:port). The dialed socket then BECOMES the worker's
/// control connection — every NAT-crossing socket is dialed FROM the worker's side.
const OP_JOIN: u32 = 19;
/// Coordinator → donor: ship the `id→bytes` display vocab once (fire-and-forget, in-band before
/// serving). DISPLAY-ONLY — lets a donor render the text it processes; never affects serving.
const OP_VOCAB: u32 = 20;
/// WebRTC signaling: the coordinator relays a non-trickle SDP handshake between two fleet peers so
/// they can open a direct data channel (P2P across NAT / browser↔browser). The header's `pos` field
/// carries the sub-kind ([`SIG_OFFER`]/[`SIG_ANSWER`]/[`SIG_FINISH`]); the payload is an opaque SDP
/// blob the coordinator never parses. See `WEBRTC_TRANSPORT.md`.
const OP_SIGNAL: u32 = 21;
/// `OP_SIGNAL` sub-kind — coordinator asks this peer to MAKE an offer; response is its full SDP
/// (candidates already gathered — non-trickle, so nothing is pushed unsolicited later).
const SIG_OFFER: u32 = 0;
/// `OP_SIGNAL` sub-kind — payload is the remote offer; this peer accepts it and responds with its
/// full answer SDP.
const SIG_ANSWER: u32 = 1;
/// `OP_SIGNAL` sub-kind — payload is the remote answer; this peer applies it (handshake completes,
/// channel opens). Response is an empty ack.
const SIG_FINISH: u32 = 2;
/// P3 WebRTC forwarding (feature `webrtc`): after an [`OP_SIGNAL`] handshake opened this peer's data
/// channel, convert it into an INBOUND serving lane — open the byte pipe and spawn a `serve_conn` on
/// it, so forwarded frames (hidden states + control ops) are served exactly like any inbound socket.
/// The answerer of the hop. `[status,len]+body` ack (status 0 = ok). See `WEBRTC_TRANSPORT.md`.
const OP_RTC_SERVE: u32 = 22;
/// P3 WebRTC forwarding: convert this peer's opened data channel into the session's FORWARD chain —
/// open the byte pipe and install it as `chain` (header `pos` bit0 = the next hop is the return
/// sink). The offerer of the hop. `[status,len]+body` ack (status 0 = ok).
const OP_RTC_CHAIN: u32 = 23;
const TAG_HIDDEN: u32 = 0;
const TAG_TOKEN: u32 = 1;
/// Byte-payload response (OP_INFO).
const TAG_INFO: u32 = 2;
/// Sink message from the LAST stage in forward mode: (TAG_DONE, group, n) + n token bits.
const TAG_DONE: u32 = 3;

fn write_msg_bytes(s: &mut Conn, header: &[u32], payload: &[u8]) -> Result<()> {
    let mut buf = Vec::with_capacity(header.len() * 4 + payload.len());
    for h in header {
        buf.extend_from_slice(&h.to_le_bytes());
    }
    buf.extend_from_slice(payload);
    s.write_all(&buf)?;
    Ok(())
}

fn write_msg(s: &mut Conn, header: &[u32], payload: &[f32]) -> Result<()> {
    write_msg_bytes(s, header, bytemuck::cast_slice(payload))
}

fn read_bytes(s: &mut Conn, n: usize) -> Result<Vec<u8>> {
    let mut b = vec![0u8; n];
    s.read_exact(&mut b)?;
    Ok(b)
}

fn read_u32s(s: &mut Conn, n: usize) -> Result<Vec<u32>> {
    let mut b = vec![0u8; n * 4];
    s.read_exact(&mut b)?;
    Ok(b.chunks_exact(4)
        .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect())
}

fn read_f32s(s: &mut Conn, n: usize) -> Result<Vec<f32>> {
    if n == 0 {
        // bytemuck rejects the EMPTY Vec<u8>'s dangling pointer (align 1) for f32 casts.
        return Ok(Vec::new());
    }
    let mut b = vec![0u8; n * 4];
    s.read_exact(&mut b)?;
    Ok(bytemuck::cast_slice(&b).to_vec())
}

/// One fleet connection. The protocol is length-prefixed, so a WebSocket is just a BYTE PIPE
/// (binary frames concatenate into the stream) — every protocol path is transport-agnostic,
/// which is the browser-stage prerequisite: anything that can open a WS can be a peer.
pub(crate) enum Conn {
    Tcp(TcpStream),
    #[cfg(feature = "ws")]
    Ws(Box<WsPipe>),
    /// A WebRTC data channel as a byte pipe (feature `webrtc`) — the last-rung P2P transport for
    /// peers behind NAT. Same length-prefixed protocol; only the byte pipe differs (see the
    /// `webrtc_native` module docs).
    #[cfg(feature = "webrtc")]
    Webrtc(Box<crate::webrtc_native::WebrtcPipe>),
}

impl Conn {
    fn peer_addr(&self) -> Result<std::net::SocketAddr> {
        match self {
            Conn::Tcp(s) => Ok(s.peer_addr()?),
            #[cfg(feature = "ws")]
            Conn::Ws(w) => w.peer_addr(),
            // A data channel has no socket peer address; the sink-shorthand path that needs one
            // only runs on the (TCP/WS) control connection, never on a WebRTC data hop.
            #[cfg(feature = "webrtc")]
            Conn::Webrtc(_) => Err(anyhow!("webrtc data channel has no socket peer address")),
        }
    }
}

impl std::io::Read for Conn {
    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
        match self {
            Conn::Tcp(s) => s.read(out),
            #[cfg(feature = "ws")]
            Conn::Ws(w) => w.read(out),
            #[cfg(feature = "webrtc")]
            Conn::Webrtc(w) => w.read(out),
        }
    }
}

impl std::io::Write for Conn {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        match self {
            Conn::Tcp(s) => s.write(buf),
            #[cfg(feature = "ws")]
            Conn::Ws(w) => w.write(buf),
            #[cfg(feature = "webrtc")]
            Conn::Webrtc(w) => w.write(buf),
        }
    }
    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Conn::Tcp(s) => s.flush(),
            #[cfg(feature = "ws")]
            Conn::Ws(w) => w.flush(),
            #[cfg(feature = "webrtc")]
            Conn::Webrtc(w) => w.flush(),
        }
    }
}

/// Frame limits sized for the protocol's largest message (32 MB ship chunks + headroom) —
/// tungstenite's 16 MB frame default would reject them.
#[cfg(feature = "ws")]
fn ws_config() -> tungstenite::protocol::WebSocketConfig {
    tungstenite::protocol::WebSocketConfig::default()
        .max_message_size(Some(256 * 1024 * 1024))
        .max_frame_size(Some(64 * 1024 * 1024))
}

/// [`Conn::Ws`]: a blocking WebSocket presented as a byte stream. One protocol message is
/// written as one binary frame (write-through — the WS analogue of TCP_NODELAY); reads drain
/// incoming binary frames in order. Control frames are tungstenite's business.
#[cfg(feature = "ws")]
pub(crate) struct WsPipe {
    ws: tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<TcpStream>>,
    buf: Vec<u8>,
    pos: usize,
}

#[cfg(feature = "ws")]
impl WsPipe {
    fn new(ws: tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<TcpStream>>) -> Self {
        Self {
            ws,
            buf: Vec::new(),
            pos: 0,
        }
    }

    fn peer_addr(&self) -> Result<std::net::SocketAddr> {
        match self.ws.get_ref() {
            tungstenite::stream::MaybeTlsStream::Plain(s) => Ok(s.peer_addr()?),
            _ => Err(anyhow!("unexpected TLS stream")),
        }
    }
}

#[cfg(feature = "ws")]
impl std::io::Read for WsPipe {
    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
        while self.pos >= self.buf.len() {
            match self.ws.read() {
                Ok(tungstenite::Message::Binary(b)) => {
                    self.buf = b.as_ref().to_vec();
                    self.pos = 0;
                }
                Ok(tungstenite::Message::Close(_)) => return Ok(0),
                Ok(_) => continue, // ping/pong/text — replies are queued internally
                Err(tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed) => {
                    return Ok(0);
                }
                Err(tungstenite::Error::Io(e)) => return Err(e),
                Err(e) => return Err(std::io::Error::other(e)),
            }
        }
        let n = out.len().min(self.buf.len() - self.pos);
        out[..n].copy_from_slice(&self.buf[self.pos..self.pos + n]);
        self.pos += n;
        Ok(n)
    }
}

#[cfg(feature = "ws")]
impl std::io::Write for WsPipe {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.ws
            .send(tungstenite::Message::binary(buf.to_vec()))
            .map_err(std::io::Error::other)?;
        Ok(buf.len())
    }
    fn flush(&mut self) -> std::io::Result<()> {
        self.ws.flush().map_err(std::io::Error::other)
    }
}

/// Dial a fleet address: `host:port` = raw TCP; `ws://host:port[/path]` = WebSocket (feature
/// `ws`). ONE codepath for coordinator dials, `--join`, worker→worker hops, and sink dials —
/// which is why a `ws://` advertise string flows through the whole chain untouched.
pub(crate) fn dial(addr: &str, timeout: std::time::Duration) -> Result<Conn> {
    let tcp_target = |hostport: &str| -> Result<std::net::SocketAddr> {
        std::net::ToSocketAddrs::to_socket_addrs(hostport)
            .with_context(|| format!("resolve {hostport}"))?
            .next()
            .ok_or_else(|| anyhow!("{hostport} resolved to nothing"))
    };
    if let Some(rest) = addr.strip_prefix("ws://") {
        #[cfg(feature = "ws")]
        {
            let hostport = rest.split('/').next().unwrap_or(rest);
            let tcp = TcpStream::connect_timeout(&tcp_target(hostport)?, timeout)?;
            tcp.set_nodelay(true).ok();
            let (ws, _) = tungstenite::client::client_with_config(
                addr,
                tungstenite::stream::MaybeTlsStream::Plain(tcp),
                Some(ws_config()),
            )
            .map_err(|e| anyhow!("ws handshake with {addr}: {e}"))?;
            return Ok(Conn::Ws(Box::new(WsPipe::new(ws))));
        }
        #[cfg(not(feature = "ws"))]
        {
            let _ = rest;
            anyhow::bail!("{addr}: this build lacks the `ws` feature");
        }
    }
    let s = TcpStream::connect_timeout(&tcp_target(addr)?, timeout)?;
    s.set_nodelay(true).ok();
    Ok(Conn::Tcp(s))
}

/// Wrap an accepted socket. WS peers open with an HTTP request (first byte `G`); the raw
/// protocol's first byte is an op id (a small integer, far below `b'G'` = 71) — so ONE
/// listener serves both transports.
pub(crate) fn accept_conn(s: TcpStream) -> Result<Conn> {
    s.set_nodelay(true).ok();
    let mut first = [0u8; 1];
    let n = s.peek(&mut first)?;
    if n == 1 && first[0] == b'G' {
        #[cfg(feature = "ws")]
        {
            let ws = tungstenite::accept_with_config(
                tungstenite::stream::MaybeTlsStream::Plain(s),
                Some(ws_config()),
            )
            .map_err(|e| anyhow!("ws accept: {e}"))?;
            return Ok(Conn::Ws(Box::new(WsPipe::new(ws))));
        }
        #[cfg(not(feature = "ws"))]
        anyhow::bail!("WS peer, but this build lacks the `ws` feature");
    }
    Ok(Conn::Tcp(s))
}

/// The next-hop socket of a P2P-chained session (see [`OP_CHAIN`]).
struct ChainOut {
    s: Conn,
    /// True when the next hop is the coordinator's return sink (tokens-only endpoint).
    sink: bool,
}

/// One loaded layer range: the engine plus everything whose lifetime is tied to the shard
/// (batch plans, MTP head, stage geometry). Replaced wholesale on [`OP_LOAD`] reloads.
struct LoadedStage {
    gpu: Lfm2Gpu,
    mtp: Option<crate::forward::MtpEngine>,
    plans: Vec<crate::forward::BatchPlan>,
    bpw: Option<crate::forward::BatchPlan>,
    start: usize,
    end: usize,
    stage_first: bool,
    stage_last: bool,
    /// Full model depth (the local file may be a shipped mini-checkpoint of a bigger stack).
    total: usize,
    /// SOURCE-checkpoint fingerprint — provenance, not the local file's own hash.
    src_fp: u64,
}

/// Session-mutable stage state. ONE lock serializes GPU work and chain writes across the
/// worker's connections (coordinator control + predecessor data) — the engine is a single
/// shared stage; concurrent coordinator sessions are unsupported by design. `loaded` is None
/// for an AUTO worker (started without `--layers`) until the coordinator's [`OP_LOAD`].
struct StageEngine {
    loaded: Option<LoadedStage>,
    chain: Option<ChainOut>,
}

/// Immutable per-worker facts + the locked engine, shared by every connection thread.
struct WorkerShared {
    ctx: GpuCtx,
    /// Local full checkpoint (`--model`). None ⇒ MODELLESS: weights arrive via [`OP_SHIP_*`]
    /// into `cache_root` and load from there.
    model_dir: Option<std::path::PathBuf>,
    /// Shipped-weights cache: `<cache_root>/<src_fp>/<start>-<end>/{config.json,
    /// model.safetensors, source.fp}`.
    cache_root: std::path::PathBuf,
    /// Full depth of the LOCAL checkpoint (0 when modelless and nothing shipped yet).
    total_layers: usize,
    /// Usable-VRAM hint for the auto-split planner (0 = unknown ⇒ equal shares).
    vram_mb: u64,
    /// True when `--layers` fixed the range at startup (a fully pinned fleet skips planning).
    pinned: bool,
    /// Checkpoint identity (config.json ⊕ safetensors tensor table) — catches
    /// same-architecture-different-weights fleet mistakes without reading any weights.
    ckpt_fp: u64,
    /// Fleet token; when set, EVERY connection must open with a valid [`OP_HELLO`].
    token: Option<String>,
    /// The set of `.part` paths currently being written by SOME connection. The partial is now a
    /// deterministic name (so an interrupted ship RESUMES rather than restarts), which means two
    /// coordinators racing the same shipment would interleave appends into one inode — so a second
    /// concurrent shipper of the same target is refused until the first releases (on OP_SHIP_END or
    /// when its connection drops), at which point a reconnect resumes cleanly.
    shipping: std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>,
    eng: std::sync::Mutex<StageEngine>,
}

/// RAII claim on a `.part` target (see [`WorkerShared::shipping`]). Released — freeing the target for
/// a resuming reconnect — when the owning [`ShipFile`] is dropped (OP_SHIP_END or a dropped session).
struct ShipClaim {
    sh: std::sync::Arc<WorkerShared>,
    key: std::path::PathBuf,
}

impl Drop for ShipClaim {
    fn drop(&mut self) {
        self.sh
            .shipping
            .lock()
            .expect("shipping set poisoned")
            .remove(&self.key);
    }
}

/// Startup options for [`run_worker`] beyond address + checkpoint.
#[derive(Default)]
pub struct WorkerOptions {
    /// Fixed layer range (`--layers start:end`, requires `--model`). None ⇒ AUTO: wait for
    /// the coordinator's [`OP_LOAD`] assignment (the auto-split path).
    pub layers: Option<(usize, usize)>,
    /// Usable VRAM in GiB (`--vram-gb`) — feeds the proportional auto-split planner.
    pub vram_gb: Option<f64>,
    /// Fleet token (`--token`); falls back to env `OSFKB_SHARD_TOKEN`.
    pub token: Option<String>,
    /// Shipped-weights cache root (`--cache-dir`); falls back to env `OSFKB_SHARD_CACHE`,
    /// then `~/.cache/lfm2-shards`.
    pub cache_dir: Option<std::path::PathBuf>,
    /// Coordinator fleet-registry address to JOIN (`--join host:port`): the worker dials OUT,
    /// registers, and serves over that socket — NAT-friendly by construction. Re-joins after
    /// every coordinator session. The local listener stays up for worker→worker data hops.
    pub join: Option<String>,
    /// `host:port` to announce for worker→worker hops on join (`--advertise`); empty ⇒ the
    /// registry derives source-IP + our listen port (right on LANs, wrong across NAT).
    pub advertise: Option<String>,
}

fn fnv64_update(mut h: u64, bytes: &[u8]) -> u64 {
    for b in bytes {
        h ^= u64::from(*b);
        h = h.wrapping_mul(0x100_0000_01b3);
    }
    h
}

fn fnv64(bytes: &[u8]) -> u64 {
    fnv64_update(0xcbf2_9ce4_8422_2325, bytes)
}

/// Cheap checkpoint identity: config.json + the safetensors TENSOR TABLE (header JSON, or the
/// HF shard index for multi-file checkpoints) — never the weights themselves.
fn checkpoint_fingerprint(dir: &Path) -> Result<u64> {
    let mut h = fnv64(&std::fs::read(dir.join("config.json")).context("read config.json")?);
    let single = dir.join("model.safetensors");
    let header = if single.exists() {
        use std::io::Read as _;
        let mut f = std::fs::File::open(&single)?;
        let mut lenb = [0u8; 8];
        f.read_exact(&mut lenb)?;
        // Defensive cap: a corrupt header length must not allocate the whole file.
        let hn = u64::from_le_bytes(lenb).min(16 * 1024 * 1024) as usize;
        let mut hb = vec![0u8; hn];
        f.read_exact(&mut hb)?;
        hb
    } else {
        std::fs::read(dir.join("model.safetensors.index.json")).context("read safetensors index")?
    };
    h ^= fnv64(&header).rotate_left(1);
    Ok(h)
}

/// The byte-stream plan for a MINIMAL single-file safetensors covering layers `start..end`:
/// the range's `…layers.{i}.…` tensors, plus EVERY non-layer tensor when the range touches an
/// end of the stack — a provably-sufficient superset of `load_shard`'s selection (embeddings,
/// final norm, lm_head, `mtp.*` all live outside the layer namespace).
pub(crate) struct MiniCkpt {
    /// Ready-to-write file prefix: u64 header length + padded header JSON.
    pub header: Vec<u8>,
    /// Absolute `(offset, len)` byte slices into the SOURCE file, in emission order.
    pub slices: Vec<(u64, u64)>,
    /// Total file size (header + data) — progress reporting.
    pub total_bytes: u64,
}

/// Layer index of a tensor name (`model.layers.7.…` / `model.language_model.layers.7.…`), or
/// None for the non-layer tensors (embeddings, norms, head, mtp).
fn tensor_layer(name: &str) -> Option<usize> {
    let idx = name.find(".layers.")?;
    // The MTP draft head (`mtp.layers.0.*`) shares the `.layers.N.` infix with the transformer stack
    // but is NOT a transformer layer — it's a whole-model head the LAST pipeline stage needs. Treat
    // it as a non-layer (end) tensor so shipping routes it to the final shard, not layer 0's. (The
    // transformer prefixes — `model.layers.` / `model.language_model.layers.` — never contain `mtp`.)
    if name[..idx].contains("mtp") {
        return None;
    }
    let rest = &name[idx + ".layers.".len()..];
    rest[..rest.find('.')?].parse().ok()
}

/// Plan the mini-checkpoint for `start..end` of `total` layers from a SINGLE-FILE source
/// (`model.safetensors`; HF-sharded sources are not shippable yet — copy those manually).
pub(crate) fn plan_mini_ckpt(
    src: &Path,
    start: usize,
    end: usize,
    total: usize,
) -> Result<MiniCkpt> {
    use std::io::Read as _;
    let mut f = std::fs::File::open(src).with_context(|| format!("open {}", src.display()))?;
    let mut lenb = [0u8; 8];
    f.read_exact(&mut lenb)?;
    let hlen = u64::from_le_bytes(lenb);
    anyhow::ensure!(hlen < 256 * 1024 * 1024, "implausible safetensors header");
    let mut hb = vec![0u8; hlen as usize];
    f.read_exact(&mut hb)?;
    let table: serde_json::Value = serde_json::from_slice(&hb).context("safetensors header")?;
    let table = table.as_object().context("header is not an object")?;
    let data_base = 8 + hlen;
    let touches_ends = start == 0 || end == total;
    // Deterministic emission order = ascending source offset (sequential reads while shipping).
    let mut picked: Vec<(&String, u64, u64, &serde_json::Value)> = Vec::new();
    for (name, meta) in table {
        if name == "__metadata__" {
            continue;
        }
        let keep = match tensor_layer(name) {
            Some(i) => start <= i && i < end,
            None => touches_ends,
        };
        if !keep {
            continue;
        }
        let offs = meta["data_offsets"]
            .as_array()
            .context("tensor data_offsets")?;
        let (a, b) = (
            offs[0].as_u64().context("offset")?,
            offs[1].as_u64().context("offset")?,
        );
        picked.push((name, a, b, meta));
    }
    anyhow::ensure!(!picked.is_empty(), "no tensors selected for {start}..{end}");
    picked.sort_by_key(|(_, a, _, _)| *a);
    let mut out = serde_json::Map::new();
    let mut cursor = 0u64;
    let mut slices = Vec::with_capacity(picked.len());
    for (name, a, b, meta) in picked {
        let len = b - a;
        let mut m = meta.clone();
        m["data_offsets"] = serde_json::json!([cursor, cursor + len]);
        out.insert(name.clone(), m);
        slices.push((data_base + a, len));
        cursor += len;
    }
    let mut hjson = serde_json::to_vec(&serde_json::Value::Object(out))?;
    // Pad the header to 8 bytes with spaces (the reference writer's convention).
    while hjson.len() % 8 != 0 {
        hjson.push(b' ');
    }
    let mut header = Vec::with_capacity(8 + hjson.len());
    header.extend_from_slice(&(hjson.len() as u64).to_le_bytes());
    header.extend_from_slice(&hjson);
    let total_bytes = header.len() as u64 + cursor;
    Ok(MiniCkpt {
        header,
        slices,
        total_bytes,
    })
}

/// Load one layer range into a fresh stage engine (`--layers` startup and [`OP_LOAD`] share
/// it). `total`/`src_fp` are the SOURCE checkpoint's identity — for shipped mini-checkpoints
/// they come from the shipment's provenance, never from hashing the local file.
fn load_stage(
    ctx: &GpuCtx,
    model_dir: &Path,
    start: usize,
    end: usize,
    total: usize,
    src_fp: u64,
) -> Result<LoadedStage> {
    let w = Weights::load_shard(ctx, model_dir, start, end)?;
    // MTP draft engine (stage_last only — load_shard loads the head there when present).
    let mtp = crate::forward::MtpEngine::new(ctx, &w);
    let (stage_first, stage_last) = (w.cfg.stage_first, w.cfg.stage_last);
    let gpu = Lfm2Gpu::new(ctx, w);
    Ok(LoadedStage {
        gpu,
        mtp,
        plans: Vec::new(),
        bpw: None,
        start,
        end,
        stage_first,
        stage_last,
        total,
        src_fp,
    })
}

/// A shipped-cache entry that is COMPLETE (committed by [`OP_SHIP_END`]) — returns its
/// `(full_depth, source_fingerprint)` provenance.
fn cache_entry_meta(entry: &Path) -> Option<(usize, u64)> {
    let fp = u64::from_str_radix(
        std::fs::read_to_string(entry.join("source.fp"))
            .ok()?
            .trim(),
        16,
    )
    .ok()?;
    if !entry.join("model.safetensors").exists() {
        return None;
    }
    let cfg =
        crate::weights::Lfm2Config::from_json(&std::fs::read(entry.join("config.json")).ok()?)
            .ok()?;
    Some((cfg.n_layers, fp))
}

/// Serve on `addr` until killed. With `--layers` the range loads at startup (PINNED); without,
/// the worker starts empty (AUTO) and loads whatever the coordinator assigns via [`OP_LOAD`].
/// Connections are served concurrently (a P2P chain needs the coordinator control connection
/// AND the predecessor data connection open at once); the engine lock serializes stage work.
pub fn run_worker(addr: &str, model_dir: Option<&Path>, opts: WorkerOptions) -> Result<()> {
    let ctx = GpuCtx::new()?;
    // The checkpoint config carries the FULL depth (load_shard rewrites cfg.n_layers to the
    // shard's) — OP_INFO reports it so the coordinator can plan + validate coverage. A
    // MODELLESS worker (no --model) reports 0|0 until weights are shipped to it.
    let (total_layers, ckpt_fp) = match model_dir {
        Some(d) => (
            crate::weights::Lfm2Config::from_json(&std::fs::read(d.join("config.json"))?)?.n_layers,
            checkpoint_fingerprint(d)?,
        ),
        None => (0, 0),
    };
    let cache_root = opts
        .cache_dir
        .or_else(|| std::env::var_os("OSFKB_SHARD_CACHE").map(std::path::PathBuf::from))
        .or_else(|| {
            std::env::var_os("HOME").map(|h| std::path::PathBuf::from(h).join(".cache/lfm2-shards"))
        })
        .unwrap_or_else(|| std::env::temp_dir().join("lfm2-shards"));
    let token = opts
        .token
        .or_else(|| std::env::var("OSFKB_SHARD_TOKEN").ok())
        .filter(|t| !t.is_empty());
    let vram_mb = opts.vram_gb.map_or(0, |g| (g * 1024.0) as u64);
    let auth = if token.is_some() { ", token auth" } else { "" };
    match (opts.layers, model_dir) {
        (Some((s0, e0)), _) => eprintln!(
            "shard worker [{s0}..{e0}) on {addr} backend {} (pinned{auth})",
            ctx.backend
        ),
        (None, Some(_)) => eprintln!(
            "shard worker AUTO on {addr} backend {} (awaiting OP_LOAD; vram {vram_mb} MB{auth})",
            ctx.backend
        ),
        (None, None) => eprintln!(
            "shard worker MODELLESS on {addr} backend {} (weights arrive by shipping; \
             cache {}; vram {vram_mb} MB{auth})",
            ctx.backend,
            cache_root.display()
        ),
    }
    let loaded = match opts.layers {
        Some((s0, e0)) => {
            let d = model_dir.context("--layers requires --model")?;
            anyhow::ensure!(
                e0 > s0 && e0 <= total_layers,
                "--layers {s0}:{e0} invalid for depth {total_layers}"
            );
            Some(load_stage(&ctx, d, s0, e0, total_layers, ckpt_fp)?)
        }
        None => None,
    };
    let listener = TcpListener::bind(addr)?;
    eprintln!("shard worker ready");
    let shared = std::sync::Arc::new(WorkerShared {
        ctx,
        model_dir: model_dir.map(Path::to_path_buf),
        cache_root,
        total_layers,
        vram_mb,
        pinned: opts.layers.is_some(),
        ckpt_fp,
        token,
        shipping: std::sync::Mutex::new(std::collections::HashSet::new()),
        eng: std::sync::Mutex::new(StageEngine {
            loaded,
            chain: None,
        }),
    });
    if let Some(join_addr) = opts.join.clone() {
        // JOIN loop: dial the coordinator's registry, authenticate, register (announcing our
        // data-hop address), then serve THAT socket as the control connection. On session end
        // (or a dead registry) retry — the worker is a persistent fleet member across
        // coordinator runs, exactly like a listener-mode worker's "awaiting next".
        let data_port: u32 = addr
            .rsplit_once(':')
            .and_then(|(_, p)| p.parse().ok())
            .context("--listen must be host:port")?;
        let advertise = opts.advertise.clone().unwrap_or_default();
        let sh = std::sync::Arc::clone(&shared);
        std::thread::spawn(move || {
            loop {
                match dial(&join_addr, std::time::Duration::from_secs(5)) {
                    Ok(mut s) => {
                        let hello_ok = match sh.token.as_deref() {
                            Some(tok) => hop_hello(&mut s, tok)
                                .map_err(|e| eprintln!("shard worker: join auth failed: {e:#}"))
                                .is_ok(),
                            None => true,
                        };
                        if hello_ok
                            && write_msg_bytes(
                                &mut s,
                                &[OP_JOIN, data_port, advertise.len() as u32],
                                advertise.as_bytes(),
                            )
                            .is_ok()
                        {
                            eprintln!("shard worker: joined coordinator at {join_addr}");
                            // The dial target is operator-configured — the socket is pre-authed.
                            match serve_conn(&sh, s, true) {
                                Ok(()) => eprintln!("shard worker: session over; re-joining"),
                                Err(e) => eprintln!("shard worker: session error: {e:#}"),
                            }
                        }
                    }
                    Err(e) => {
                        eprintln!("shard worker: join {join_addr} unreachable ({e}); retrying")
                    }
                }
                std::thread::sleep(std::time::Duration::from_secs(2));
            }
        });
    }
    for conn in listener.incoming() {
        let s = conn?;
        let sh = std::sync::Arc::clone(&shared);
        std::thread::spawn(move || {
            let c = match accept_conn(s) {
                Ok(c) => c,
                Err(e) => {
                    eprintln!("shard worker: handshake failed: {e:#}");
                    return;
                }
            };
            match serve_conn(&sh, c, false) {
                Ok(()) => eprintln!("shard worker: session ended; awaiting next"),
                Err(e) => eprintln!("shard worker: connection error: {e:#}"),
            }
        });
    }
    Ok(())
}

/// One in-flight shipped file on a connection (between [`OP_SHIP_BEGIN`] and [`OP_SHIP_END`]).
struct ShipFile {
    entry: std::path::PathBuf,
    fname: String,
    /// Source-checkpoint fingerprint (hex) — written as provenance on commit.
    fp: String,
    /// The DETERMINISTIC partial `<entry>/<fname>.part`, appended to across resumes and renamed to
    /// the final name on a clean [`OP_SHIP_END`].
    part: std::path::PathBuf,
    w: std::io::BufWriter<std::fs::File>,
    fnv: u64,
    /// Held for the ship's lifetime; releasing it (on drop) frees the target for a resuming reconnect.
    _claim: ShipClaim,
}

/// FNV-1a seed; the shipped-file integrity hash accumulated across every chunk.
const SHIP_FNV_INIT: u64 = 0xcbf2_9ce4_8422_2325;

/// Ship chunk size. Small enough that a single in-flight chunk (the stop-and-wait window) never
/// overwhelms a browser receiver, large enough that per-chunk ack RTT is negligible next to the
/// bytes it carries. Under tungstenite's 64 MB frame cap.
const SHIP_CHUNK_BYTES: usize = 8 * 1024 * 1024;

/// Reconstruct `(bytes_present, running_fnv)` from a partial ship file — one sequential pass so the
/// resume point is always consistent with the bytes actually on disk (crash- and reset-safe). A
/// missing partial resumes at zero.
fn resume_state(part: &Path) -> Result<(u64, u64)> {
    use std::io::Read as _;
    let mut f = match std::fs::File::open(part) {
        Ok(f) => f,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok((0, SHIP_FNV_INIT)),
        Err(e) => return Err(e.into()),
    };
    let mut fnv = SHIP_FNV_INIT;
    let mut have = 0u64;
    let mut buf = vec![0u8; 1 << 20];
    loop {
        let k = f.read(&mut buf)?;
        if k == 0 {
            break;
        }
        fnv = fnv64_update(fnv, &buf[..k]);
        have += k as u64;
    }
    Ok((have, fnv))
}

/// The stage engine, or the actionable error for ops that arrived before any [`OP_LOAD`].
fn need_loaded(loaded: &mut Option<LoadedStage>) -> Result<&mut LoadedStage> {
    loaded.as_mut().ok_or_else(|| {
        anyhow!(
            "stage not loaded — pin --layers on the worker or connect an auto-split coordinator"
        )
    })
}

/// Authenticate an OUTBOUND hop (next worker or sink) before trusting it with the chain.
fn hop_hello(s: &mut Conn, token: &str) -> Result<()> {
    write_msg_bytes(s, &[OP_HELLO, 0, token.len() as u32], token.as_bytes())?;
    let hdr = read_u32s(s, 2)?;
    let p = read_f32s(s, hdr[1] as usize)?;
    anyhow::ensure!(
        p.first().map(|x| x.to_bits()) == Some(0),
        "next hop rejected the fleet token"
    );
    Ok(())
}

/// One connection's op loop. Payloads are read lock-free (each connection owns its stream);
/// the engine lock is held per op, covering GPU work AND the forward write so chained output
/// order matches compute order.
fn serve_conn(sh: &std::sync::Arc<WorkerShared>, mut s: Conn, pre_authed: bool) -> Result<()> {
    let ctx = &sh.ctx;
    let trace = std::env::var_os("OSFKB_STAGE_TRACE").is_some();
    // MTP-training data tap (see training/README.md): APPEND so successive coordinator
    // sessions accumulate one corpus — the loader segments a stream when its position
    // resets (a new session reuses slot ids and restarts positions at the prompt).
    let mut dump_hidden = std::env::var("OSFKB_DUMP_HIDDEN").ok().map(|path| {
        std::io::BufWriter::new(
            std::fs::OpenOptions::new()
                .append(true)
                .create(true)
                .open(path)
                .expect("dump file"),
        )
    });
    let (mut bstep_n, mut bstep_recv, mut bstep_gpu, mut bstep_send) = (0u64, 0f64, 0f64, 0f64);
    let mut bstep_idle = 0f64;
    let mut last_done = std::time::Instant::now();
    // Always-on donor self-report (independent of OSFKB_STAGE_TRACE): a periodic stderr line so a
    // native GPU donor shows what layer range it's serving and at what rate. The window resets on
    // an idle gap between turns, so tok/s reflects ACTIVE serving rather than wall-clock.
    let (mut donor_steps, mut donor_win_toks) = (0u64, 0u64);
    let (mut donor_win_start, mut donor_last_step) =
        (std::time::Instant::now(), std::time::Instant::now());
    let mut last_stat = std::time::Instant::now();
    // Donor token→text readout (armed by OP_VOCAB): a bounded ring of the recent ids this stage
    // sees, decoded ONLY on the throttled stat tick. Display-only; never touches serving.
    let mut vocab: Option<crate::vocab::ByteVocab> = None;
    let mut recent: std::collections::VecDeque<u32> = std::collections::VecDeque::new();
    // Token gate: with a fleet token configured, the FIRST op must be a valid OP_HELLO.
    // JOIN sockets arrive pre-authed: the worker dialed an operator-configured address and
    // already proved ITSELF to the registry.
    let mut authed = pre_authed || sh.token.is_none();
    // In-flight shipped file (transfer state is connection-local, not engine state).
    let mut ship: Option<ShipFile> = None;
    // WebRTC peer for this session (armed by OP_SIGNAL when the `webrtc` feature is built). Kept
    // alive here so the data channel's async open/message events still fire after the signaling ops
    // return — mirrors the browser worker's session-local peer.
    #[cfg(feature = "webrtc")]
    let mut peer: Option<crate::webrtc_native::NativeWebrtcPeer> = None;
    loop {
        let hdr = match read_u32s(&mut s, 3) {
            Ok(h) => h,
            Err(_) => break, // peer gone → this session lane ends; keep listening
        };
        let (op, pos, n) = (hdr[0], hdr[1] as usize, hdr[2] as usize);
        if op == OP_HELLO {
            let raw = read_bytes(&mut s, n)?;
            let ok = sh
                .token
                .as_deref()
                .is_none_or(|t| t.as_bytes() == raw.as_slice());
            write_msg(&mut s, &[TAG_TOKEN, 1], &[f32::from_bits(u32::from(!ok))])?;
            if !ok {
                return Err(anyhow!("fleet token rejected"));
            }
            authed = true;
            continue;
        }
        anyhow::ensure!(
            authed,
            "unauthenticated connection: this worker requires the fleet token (OSFKB_SHARD_TOKEN)"
        );
        // WebRTC signaling (feature `webrtc`) is served BEFORE taking the engine lock: it does no
        // engine work, and each step blocks on ICE gathering (~100 ms+), which must never stall the
        // shared stage. The sub-kind rides the header's `pos` field. Without the feature, the
        // `OP_SIGNAL` match arm below answers "unsupported" and the coordinator falls back to TCP/WS.
        #[cfg(feature = "webrtc")]
        if op == OP_SIGNAL {
            let payload = read_bytes(&mut s, n)?;
            let (status, body) = match pos as u32 {
                SIG_OFFER => crate::webrtc_native::do_offer(&mut peer),
                SIG_ANSWER => crate::webrtc_native::do_answer(&mut peer, &payload),
                SIG_FINISH => crate::webrtc_native::do_finish(&peer, &payload),
                other => (1u32, format!("bad webrtc sub-kind {other}").into_bytes()),
            };
            write_msg_bytes(&mut s, &[status, body.len() as u32], &body)?;
            continue;
        }
        // P3 (feature `webrtc`): promote the handshaken data channel into a forwarding lane. Both
        // ops open the byte pipe (a brief `block_on`, off the engine lock) and are then transport-
        // identical to their TCP/WS equivalents — the answerer serves the channel as an inbound
        // session; the offerer installs it as this session's forward chain.
        #[cfg(feature = "webrtc")]
        if op == OP_RTC_SERVE {
            let _ = read_bytes(&mut s, n)?;
            let (status, body) = match crate::webrtc_native::open_pipe_blocking(peer.take()) {
                Ok(pipe) => {
                    // Serve the inbound channel on its own thread, exactly like an accepted socket
                    // (mirrors how a TCP predecessor's dial is accepted into `serve_conn`).
                    let sh2 = std::sync::Arc::clone(sh);
                    std::thread::spawn(move || {
                        if let Err(e) = serve_conn(&sh2, Conn::Webrtc(Box::new(pipe)), true) {
                            eprintln!("shard worker: webrtc inbound session ended: {e:#}");
                        }
                    });
                    (0u32, Vec::new())
                }
                Err(e) => (1u32, e.into_bytes()),
            };
            write_msg_bytes(&mut s, &[status, body.len() as u32], &body)?;
            continue;
        }
        #[cfg(feature = "webrtc")]
        if op == OP_RTC_CHAIN {
            let _ = read_bytes(&mut s, n)?;
            let sink = pos & 1 == 1;
            let (status, body) = match crate::webrtc_native::open_pipe_blocking(peer.take()) {
                Ok(pipe) => {
                    sh.eng.lock().expect("engine lock poisoned").chain = Some(ChainOut {
                        s: Conn::Webrtc(Box::new(pipe)),
                        sink,
                    });
                    (0u32, Vec::new())
                }
                Err(e) => (1u32, e.into_bytes()),
            };
            write_msg_bytes(&mut s, &[status, body.len() as u32], &body)?;
            continue;
        }
        // Engine lock for the WHOLE op (payload reads included): compute order equals
        // chained-write order, and concurrent connections (coordinator control + P2P
        // predecessor data) stay serialized on the single shared stage.
        let mut guard = sh.eng.lock().expect("engine lock poisoned");
        let StageEngine { loaded, chain } = &mut *guard;
        match op {
            OP_BINIT => {
                let LoadedStage {
                    gpu, plans, bpw, ..
                } = need_loaded(loaded)?;
                // Header a = k | (n_groups << 16). PER-GROUP decode plans: each in-flight
                // micro-batch group owns its plan so its DN snapshots survive the other
                // groups' interleaved µbatches until ITS restore arrives (spec rollback).
                let (kcols, groups) = (pos & 0xFFFF, (pos >> 16).max(1));
                // OSFKB_MTP_SPEC=1 → the decode plan snapshots DN state per column, enabling
                // draft-verify rollback (dn_restore); ~2 MB/col/DN-layer of extra traffic.
                let spec = std::env::var("OSFKB_MTP_SPEC").ok().as_deref() == Some("1");
                *plans = (0..groups)
                    .map(|_| {
                        if spec {
                            gpu.make_batch_plan_spec(ctx, kcols)
                        } else {
                            gpu.make_batch_plan(ctx, kcols)
                        }
                    })
                    .collect();
                // Header field `n` = wide-prefill width (0 = disabled).
                *bpw = if n > 0 {
                    anyhow::ensure!(
                        ctx.subgroups,
                        "coordinator requested wide prefill; this stage's adapter lacks subgroups"
                    );
                    Some(gpu.make_batch_plan_wide(ctx, n))
                } else {
                    None
                };
                write_msg(&mut s, &[TAG_TOKEN, 1], &[f32::from_bits(0)])?;
            }
            OP_ZSLOT => {
                let LoadedStage { gpu, .. } = need_loaded(loaded)?;
                gpu.zero_dn_slot(ctx, pos);
                // P2P: control ops ride the chain ahead of the µbatch they precede.
                if let Some(c) = chain.as_mut().filter(|c| !c.sink) {
                    write_msg(&mut c.s, &hdr, &[])?;
                }
            }
            OP_BTAB => {
                let LoadedStage { gpu, mtp, .. } = need_loaded(loaded)?;
                let raw = read_f32s(&mut s, n)?;
                let blocks: Vec<u32> = raw.iter().map(|x| x.to_bits()).collect();
                gpu.write_btab_row(ctx, pos as u32, &blocks);
                // The MTP draft model keeps per-slot KV of its own — mirror the mapping.
                if let Some(m) = mtp.as_ref() {
                    m.gpu.write_btab_row(ctx, pos as u32, &blocks);
                }
                if let Some(c) = chain.as_mut().filter(|c| !c.sink) {
                    write_msg(&mut c.s, &hdr, &raw)?;
                }
            }
            OP_BSTEP | OP_BSTEPW => {
                let LoadedStage {
                    gpu,
                    mtp,
                    plans,
                    bpw,
                    stage_first,
                    stage_last,
                    start,
                    end,
                    ..
                } = need_loaded(loaded)?;
                let group = (pos >> 16) & 0xFF;
                let spec_k = (pos >> 24) & 0x3F;
                let spec_on = (pos >> 30) & 1 == 1;
                let spec_verify = (pos >> 31) & 1 == 1;
                let pos = pos & 0xFFFF;
                // Inter-arrival idle: time since the previous bstep FINISHED — the pipeline
                // starvation share (0 = fully fed stage).
                if trace && last_done.elapsed().as_secs_f64() < 3600.0 {
                    bstep_idle += last_done.elapsed().as_secs_f64();
                }
                let t0 = std::time::Instant::now();
                let colsu = read_u32s(&mut s, 4 * pos)?;
                let hidden = read_f32s(&mut s, n)?;
                let t_recv = t0.elapsed();
                let plan = if op == OP_BSTEPW {
                    bpw.as_ref()
                        .ok_or_else(|| anyhow!("OP_BSTEPW without a wide plan"))?
                } else {
                    plans
                        .get(group)
                        .ok_or_else(|| anyhow!("OP_BSTEP group {group} out of range"))?
                };
                let cols: Vec<BatchCol> = colsu
                    .chunks(4)
                    .map(|c| BatchCol::text(c[3], c[0], c[1], c[2] != 0))
                    .collect();
                // A stage_first worker embeds its own tokens (headless coordinator).
                let hin = (!*stage_first).then_some(hidden.as_slice());
                let t1 = std::time::Instant::now();
                let out = gpu.batch_stage_step(ctx, plan, &cols, hin)?;
                let t_gpu = t1.elapsed();
                if *stage_last && let Some(dump) = &mut dump_hidden {
                    // MTP-training data: the PRE-final-norm residual per column — the exact
                    // tensor the draft chain seeds from. Record: pos u32, token u32, h f32s.
                    let hbuf = gpu.read_cur_for_tests(ctx, plan, cols.len())?;
                    let hsz = hbuf.len() / cols.len();
                    for (p, col) in cols.iter().enumerate() {
                        use std::io::Write;
                        dump.write_all(&col.pos.to_le_bytes())?;
                        dump.write_all(&col.btrow.to_le_bytes())?;
                        dump.write_all(&col.token.to_le_bytes())?;
                        dump.write_all(bytemuck::cast_slice(&hbuf[p * hsz..(p + 1) * hsz]))?;
                    }
                }
                if trace && op == OP_BSTEPW {
                    // WIDE (prefill) µbatches are few and long — log each one.
                    let br = gpu.take_step_breakdown();
                    eprintln!(
                        "bstepw ncols={pos} recv {:.1}ms gpu {:.1}ms [prep {:.1} enc {:.1} poll {:.1}] idle-before {:.1}ms",
                        t_recv.as_secs_f64() * 1e3,
                        t_gpu.as_secs_f64() * 1e3,
                        br.0 * 1e3,
                        br.1 * 1e3,
                        br.2 * 1e3,
                        last_done.elapsed().as_secs_f64() * 1e3
                            - t_gpu.as_secs_f64() * 1e3
                            - t_recv.as_secs_f64() * 1e3,
                    );
                }
                if trace && pos <= 8 {
                    // Small-width breakdown (mono verify): what the wall is made of.
                    let br = gpu.take_step_breakdown();
                    eprintln!(
                        "bstep small ncols={pos}: total {:.2}ms [prep {:.2} encode {:.2} poll+read {:.2}]",
                        t_gpu.as_secs_f64() * 1e3,
                        br.0 * 1e3,
                        br.1 * 1e3,
                        br.2 * 1e3
                    );
                }
                let t2 = std::time::Instant::now();
                match out {
                    StageBatchOut::Hidden(h) => match chain.as_mut() {
                        None => write_msg(&mut s, &[TAG_HIDDEN, h.len() as u32], &h)?,
                        Some(c) if !c.sink => {
                            // TRUE-P2P hop: same op + column meta verbatim, THIS stage's
                            // residual as the payload — the coordinator never sees it.
                            let mut fh = Vec::with_capacity(3 + colsu.len());
                            fh.extend_from_slice(&[op, hdr[1], h.len() as u32]);
                            fh.extend_from_slice(&colsu);
                            write_msg(&mut c.s, &fh, &h)?;
                        }
                        Some(_) => {
                            return Err(anyhow!("non-final stage chained to the return sink"));
                        }
                    },
                    StageBatchOut::Tokens(t) => {
                        let mut resp: Vec<u32> = t[..pos].to_vec();
                        // Serving-MTP SPONTANEOUS DRAFT: the last stage has everything —
                        // the drafts (cols meta tokens), the verify tokens, and this
                        // group's plan `cur` — so it accepts, seeds, and drafts here,
                        // appending the new drafts to the SAME response (zero extra trips).
                        if (spec_verify || spec_on) && mtp.is_some() {
                            let m = mtp.as_ref().expect("checked");
                            let plan = plans
                                .get(group)
                                .ok_or_else(|| anyhow!("spec group out of range"))?;
                            let mut reqs: Vec<crate::forward::DraftReq> = Vec::new();
                            // Seed copies (slot ← main col), folded into the chain's first
                            // submission (`draft_chain_batch_seeded`) — 1:1 with `reqs`.
                            let mut seeds: Vec<(u32, usize)> = Vec::new();
                            let mut pairs: Vec<crate::forward::DraftPair> = Vec::new();
                            // Maximal runs of equal btrow = spans (spec µbatch) / chunks.
                            let mut i = 0usize;
                            while i < cols.len() {
                                let mut j = i + 1;
                                while j < cols.len() && cols[j].btrow == cols[i].btrow {
                                    j += 1;
                                }
                                let slot = cols[i].btrow;
                                if spec_verify {
                                    // Span [tok, d1..dk]: accept prefix, seed at the last
                                    // valid column, request the next chain.
                                    let k = j - i - 1;
                                    let mut a = 0usize;
                                    while a < k && t[i + a] == cols[i + 1 + a].token {
                                        a += 1;
                                    }
                                    seeds.push((slot, i + a));
                                    reqs.push(crate::forward::DraftReq {
                                        btrow: slot,
                                        first_token: t[i + a],
                                        pos0: cols[i + a].pos as usize,
                                        mpos0: cols[i + a].mpos,
                                    });
                                } else {
                                    // Plain under spec serving: draft-KV prefill pairs for
                                    // the chunk body, a fresh chain for the emitting tail.
                                    for p in i..j.saturating_sub(1) {
                                        pairs.push(crate::forward::DraftPair {
                                            btrow: slot,
                                            cur_col: p,
                                            token: cols[p + 1].token,
                                            pos: cols[p].pos as usize,
                                            mpos: cols[p + 1].mpos,
                                        });
                                    }
                                    let lastc = j - 1;
                                    if cols[lastc].need_logit {
                                        seeds.push((slot, lastc));
                                        reqs.push(crate::forward::DraftReq {
                                            btrow: slot,
                                            first_token: t[lastc],
                                            pos0: cols[lastc].pos as usize,
                                            mpos0: cols[lastc].mpos,
                                        });
                                    }
                                }
                                i = j;
                            }
                            for chunk in pairs.chunks(crate::forward::MtpEngine::K_DRAFT_BATCH) {
                                m.draft_pairs(ctx, plan, chunk)?;
                            }
                            if spec_k > 0 {
                                for (chunk, seed_chunk) in reqs
                                    .chunks(crate::forward::MtpEngine::K_DRAFT_BATCH)
                                    .zip(seeds.chunks(crate::forward::MtpEngine::K_DRAFT_BATCH))
                                {
                                    let drafts = m.draft_chain_batch_seeded(
                                        ctx,
                                        Some(plan),
                                        seed_chunk,
                                        chunk,
                                        spec_k,
                                    )?;
                                    resp.extend(drafts.into_iter().flatten());
                                }
                            } else {
                                // Depth-0 (occupancy drain): accepts were seeded for a FUTURE
                                // re-draft round, so the copies must still land.
                                for (slot, col) in &seeds {
                                    m.seed_slot_from_col(ctx, plan, *col, *slot as usize);
                                }
                            }
                        }
                        let bits: Vec<f32> = resp.iter().map(|x| f32::from_bits(*x)).collect();
                        match chain.as_mut() {
                            None => {
                                write_msg(&mut s, &[TAG_TOKEN, bits.len() as u32], &bits)?;
                            }
                            Some(c) if c.sink => {
                                // Only the tokens travel back — straight to the sink.
                                write_msg(
                                    &mut c.s,
                                    &[TAG_DONE, group as u32, bits.len() as u32],
                                    &bits,
                                )?;
                            }
                            Some(_) => {
                                return Err(anyhow!("final stage must chain to the return sink"));
                            }
                        }
                    }
                }
                last_done = std::time::Instant::now();
                if donor_last_step.elapsed().as_secs_f64() > 0.5 {
                    // Idle gap between turns → fresh rate window (don't count idle as serving).
                    donor_win_toks = 0;
                    donor_win_start = std::time::Instant::now();
                }
                donor_last_step = std::time::Instant::now();
                donor_steps += 1;
                donor_win_toks += pos as u64;
                // Buffer the ids this stage sees (O(1)/step); on the last stage these are the
                // fed-back generated tokens, so the decoded ring reads as the text being generated.
                for c in &cols {
                    recent.push_back(c.token);
                }
                while recent.len() > 64 {
                    recent.pop_front();
                }
                if last_stat.elapsed().as_secs_f64() >= 2.0
                    && donor_win_start.elapsed().as_secs_f64() >= 0.25
                {
                    let rate = donor_win_toks as f64 / donor_win_start.elapsed().as_secs_f64();
                    let tail = vocab.as_ref().map_or(String::new(), |v| {
                        let text = v.decode(&recent.iter().copied().collect::<Vec<_>>());
                        let chars: Vec<char> = text.chars().collect();
                        chars[chars.len().saturating_sub(48)..]
                            .iter()
                            .collect::<String>()
                            .replace('\n', " ")
                    });
                    eprintln!(
                        "shard worker: serving {}..{} · {donor_steps} steps · {rate:.1} tok/s · \"{tail}\"",
                        *start, *end
                    );
                    donor_win_toks = 0;
                    donor_win_start = std::time::Instant::now();
                    last_stat = std::time::Instant::now();
                }
                if trace {
                    bstep_n += 1;
                    bstep_recv += t_recv.as_secs_f64();
                    bstep_gpu += t_gpu.as_secs_f64();
                    bstep_send += t2.elapsed().as_secs_f64();
                    if bstep_n % 50 == 0 {
                        eprintln!(
                            "stage trace [{bstep_n}]: recv {:.1}ms gpu+sync {:.1}ms send {:.1}ms idle {:.1}ms (avg/step, ncols {})",
                            bstep_recv * 20.0,
                            bstep_gpu * 20.0,
                            bstep_send * 20.0,
                            bstep_idle * 20.0,
                            pos
                        );
                        bstep_recv = 0.0;
                        bstep_gpu = 0.0;
                        bstep_send = 0.0;
                        bstep_idle = 0.0;
                    }
                }
            }
            OP_MTP_BATCH => {
                let LoadedStage { mtp, .. } = need_loaded(loaded)?;
                let (ncols, kdraft) = (pos, n);
                let payload = read_f32s(&mut s, ncols * 3)?;
                let m = mtp
                    .as_ref()
                    .ok_or_else(|| anyhow!("OP_MTP_BATCH without an MTP head"))?;
                let reqs: Vec<(u32, u32, usize)> = payload
                    .chunks(3)
                    .map(|c| (c[0].to_bits(), c[1].to_bits(), c[2].to_bits() as usize))
                    .collect();
                let drafts = m.draft_chain_batch(ctx, &reqs, kdraft)?;
                let flat: Vec<f32> = drafts
                    .iter()
                    .flat_map(|d| d.iter().map(|x| f32::from_bits(*x)))
                    .collect();
                write_msg(&mut s, &[TAG_TOKEN, flat.len() as u32], &flat)?;
            }
            OP_MTP_SEED => {
                let LoadedStage { mtp, plans, .. } = need_loaded(loaded)?;
                // header: (op, slot, col | group << 16) — seed the draft chain slab from
                // that GROUP's plan verify column (in-band, no response).
                let (col, group) = (n & 0xFFFF, (n >> 16) & 0xFFFF);
                let m = mtp
                    .as_ref()
                    .ok_or_else(|| anyhow!("OP_MTP_SEED without an MTP head"))?;
                let plan = plans
                    .get(group)
                    .ok_or_else(|| anyhow!("OP_MTP_SEED group {group} out of range"))?;
                m.seed_slot_from_col(ctx, plan, col, pos);
            }
            OP_MTP => {
                let LoadedStage { mtp, plans, .. } = need_loaded(loaded)?;
                let payload = read_f32s(&mut s, 1)?;
                let first_tok = payload[0].to_bits();
                let (col_j, kdraft) = (n >> 8, n & 0xff);
                let m = mtp
                    .as_ref()
                    .ok_or_else(|| anyhow!("OP_MTP without an MTP head on this stage"))?;
                let plan = plans
                    .first()
                    .ok_or_else(|| anyhow!("OP_MTP before OP_BINIT"))?;
                // Seed copy + chain in ONE submit (the seed used to be its own submit).
                let drafts = m.draft_chain_seeded(ctx, plan, col_j, first_tok, pos, kdraft)?;
                let mut resp = drafts.clone();
                if std::env::var("OSFKB_MTP_BETA_PROBE").is_ok() {
                    for t3 in m.last_top3() {
                        resp.extend_from_slice(&t3);
                    }
                }
                let bits: Vec<f32> = resp.iter().map(|x| f32::from_bits(*x)).collect();
                write_msg(&mut s, &[TAG_TOKEN, bits.len() as u32], &bits)?;
            }
            OP_DNRESTORE => {
                let LoadedStage { gpu, plans, .. } = need_loaded(loaded)?;
                // n = col | (group << 16).
                let (col, group) = (n & 0xFFFF, (n >> 16) & 0xFFFF);
                let plan = plans
                    .get(group)
                    .ok_or_else(|| anyhow!("OP_DNRESTORE group {group} out of range"))?;
                gpu.dn_restore(ctx, plan, pos as u32, col);
                if let Some(c) = chain.as_mut().filter(|c| !c.sink) {
                    write_msg(&mut c.s, &hdr, &[])?;
                }
            }
            OP_STEP => {
                let LoadedStage { gpu, .. } = need_loaded(loaded)?;
                let hidden = read_f32s(&mut s, n)?;
                match gpu.stage_forward(ctx, pos, 0, Some(&hidden))? {
                    StageOut::Hidden(h) => match chain.as_mut() {
                        None => write_msg(&mut s, &[TAG_HIDDEN, h.len() as u32], &h)?,
                        Some(c) if !c.sink => {
                            write_msg(&mut c.s, &[OP_STEP, pos as u32, h.len() as u32], &h)?;
                        }
                        Some(_) => {
                            return Err(anyhow!("non-final stage chained to the return sink"));
                        }
                    },
                    StageOut::Token(t) => match chain.as_mut() {
                        None => write_msg(&mut s, &[TAG_TOKEN, 1], &[f32::from_bits(t)])?,
                        Some(c) if c.sink => {
                            write_msg(&mut c.s, &[TAG_DONE, 0, 1], &[f32::from_bits(t)])?;
                        }
                        Some(_) => {
                            return Err(anyhow!("final stage must chain to the return sink"));
                        }
                    },
                }
            }
            OP_RESET => {
                let LoadedStage { gpu, .. } = need_loaded(loaded)?;
                gpu.reset(ctx);
                write_msg(&mut s, &[TAG_TOKEN, 1], &[f32::from_bits(0)])?;
            }
            OP_CHAIN => {
                let raw = read_bytes(&mut s, n)?;
                *chain = None; // any previous hop is dropped before dialing the new one
                let ok = if raw.is_empty() {
                    true
                } else {
                    let want = String::from_utf8(raw).context("chain addr utf8")?;
                    let addr = if let Some(port) = want.strip_prefix(':') {
                        // Sink shorthand: the coordinator's routable address is wherever
                        // THIS control connection came from.
                        std::net::SocketAddr::new(s.peer_addr()?.ip(), port.parse()?).to_string()
                    } else {
                        want
                    };
                    match dial(&addr, std::time::Duration::from_secs(5)) {
                        Ok(mut t) => {
                            // Authenticate to the next hop BEFORE acking our own chain setup —
                            // a tokened fleet never ships activations to an unverified peer.
                            // The SINK hello is write-only: the coordinator is still blocked on
                            // THIS op's ack, so it verifies on accept (an acked sink hello
                            // would deadlock setup); the sink address's provenance is the
                            // already-authenticated control connection.
                            let hop_ok = match (sh.token.as_deref(), pos & 1 == 1) {
                                (Some(tok), true) => write_msg_bytes(
                                    &mut t,
                                    &[OP_HELLO, 0, tok.len() as u32],
                                    tok.as_bytes(),
                                )
                                .is_ok(),
                                (Some(tok), false) => hop_hello(&mut t, tok)
                                    .map_err(|e| {
                                        eprintln!("shard worker: hop auth {addr} failed: {e:#}");
                                    })
                                    .is_ok(),
                                (None, _) => true,
                            };
                            if hop_ok {
                                *chain = Some(ChainOut {
                                    s: t,
                                    sink: pos & 1 == 1,
                                });
                            }
                            hop_ok
                        }
                        Err(err) => {
                            eprintln!("shard worker: chain dial {addr} failed: {err}");
                            false
                        }
                    }
                };
                write_msg(&mut s, &[TAG_TOKEN, 1], &[f32::from_bits(u32::from(!ok))])?;
            }
            OP_LOAD => {
                // Payload: 16 hex bytes — the SOURCE checkpoint the coordinator wants. Acks:
                // 0 ok, 1 fail, 2 need-weights (ship then retry), 3 wrong-checkpoint.
                let (start, end) = (pos, n);
                let fpx = String::from_utf8(read_bytes(&mut s, 16)?).context("load fp utf8")?;
                let want_fp = u64::from_str_radix(&fpx, 16).context("load fp hex")?;
                let ack: u32 = if loaded
                    .as_ref()
                    .is_some_and(|l| l.start == start && l.end == end && l.src_fp == want_fp)
                {
                    eprintln!("shard worker: layers {start}..{end} already resident — no reload");
                    0
                } else {
                    // Resolve a source for this (checkpoint, range): the local full model
                    // when its identity matches, else a COMPLETE shipped-cache entry.
                    let src: std::result::Result<(std::path::PathBuf, usize, u64), u32> =
                        if let Some(md) = &sh.model_dir {
                            if sh.ckpt_fp == want_fp {
                                Ok((md.clone(), sh.total_layers, sh.ckpt_fp))
                            } else {
                                eprintln!(
                                    "shard worker: OP_LOAD wants {fpx} but --model is \
                                     {:016x} — wrong checkpoint",
                                    sh.ckpt_fp
                                );
                                Err(3)
                            }
                        } else {
                            let entry = sh.cache_root.join(&fpx).join(format!("{start}-{end}"));
                            match cache_entry_meta(&entry) {
                                Some((total, fp)) if fp == want_fp => Ok((entry, total, fp)),
                                _ => Err(2), // not cached (or incomplete) → ship me the weights
                            }
                        };
                    match src {
                        Err(code) => code,
                        Ok((dir, total, fp)) if end > start && end <= total => {
                            *loaded = None; // release the old shard's VRAM before the new load
                            match load_stage(ctx, &dir, start, end, total, fp) {
                                Ok(l) => {
                                    eprintln!("shard worker: loaded layers {start}..{end}");
                                    *loaded = Some(l);
                                    0
                                }
                                Err(e) => {
                                    eprintln!("shard worker: OP_LOAD {start}..{end} failed: {e:#}");
                                    1
                                }
                            }
                        }
                        Ok((_, total, _)) => {
                            eprintln!(
                                "shard worker: OP_LOAD {start}..{end} invalid for depth {total}"
                            );
                            1
                        }
                    }
                };
                write_msg(&mut s, &[TAG_TOKEN, 1], &[f32::from_bits(ack)])?;
            }
            OP_SHIP_BEGIN => {
                // meta: `fp16hex|start|end|filename` → a DETERMINISTIC partial `<cache>/<fp>/<s>-<e>/
                // <file>.part`. RESUMABLE: if a prior (interrupted) ship left a partial, we compute
                // how many bytes it already holds and its running FNV, reply them, and the sender
                // skips that prefix — so a WAN reset resumes instead of restarting from zero.
                let meta = String::from_utf8(read_bytes(&mut s, n)?).context("ship meta utf8")?;
                // 4 fields (fp|start|end|fname) or 5 with a trailing total-size hint (a browser
                // receiver uses it to pre-size; this disk-backed native receiver ignores it).
                let parts: Vec<&str> = meta.splitn(5, '|').collect();
                anyhow::ensure!(parts.len() >= 4, "malformed ship meta: {meta}");
                let entry = sh.cache_root.join(parts[0]).join(format!(
                    "{}-{}",
                    parts[1].parse::<usize>()?,
                    parts[2].parse::<usize>()?
                ));
                std::fs::create_dir_all(&entry)?;
                let fname = parts[3].to_string();
                anyhow::ensure!(
                    fname == "config.json" || fname == "model.safetensors",
                    "unexpected ship filename {fname}"
                );
                let part = entry.join(format!("{fname}.part"));
                // Claim the target so a concurrent shipper can't interleave appends into it. A
                // busy target replies the sentinel `have = u64::MAX`, and the sender backs off.
                let claimed = sh
                    .shipping
                    .lock()
                    .expect("shipping set poisoned")
                    .insert(part.clone());
                if !claimed {
                    eprintln!("shard worker: ship {fname} busy (another shipper holds {part:?})");
                    write_msg(&mut s, &[u32::MAX, u32::MAX, 0, 0], &[])?;
                    continue;
                }
                let claim = ShipClaim {
                    sh: std::sync::Arc::clone(sh),
                    key: part.clone(),
                };
                // Reconstruct (have, fnv) from any existing partial — one sequential read, so the
                // resume point is always consistent with the bytes actually on disk (crash-safe).
                let (have, fnv) = resume_state(&part)?;
                let f = std::fs::OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(&part)?;
                eprintln!(
                    "shard worker: ship begin {} → {} (resume at {:.2} GB)",
                    fname,
                    entry.display(),
                    have as f64 / 1e9
                );
                ship = Some(ShipFile {
                    entry,
                    fname,
                    fp: parts[0].to_string(),
                    part,
                    w: std::io::BufWriter::new(f),
                    fnv,
                    _claim: claim,
                });
                // Reply the resume point: [have_lo, have_hi, fnv_lo, fnv_hi].
                write_msg(
                    &mut s,
                    &[
                        have as u32,
                        (have >> 32) as u32,
                        fnv as u32,
                        (fnv >> 32) as u32,
                    ],
                    &[],
                )?;
            }
            OP_SHIP_CHUNK => {
                let sf = ship
                    .as_mut()
                    .ok_or_else(|| anyhow!("chunk without begin"))?;
                let raw = read_bytes(&mut s, n)?;
                sf.fnv = fnv64_update(sf.fnv, &raw);
                std::io::Write::write_all(&mut sf.w, &raw)?;
                // Durably land the bytes so a later resume sees them, then ACK — the sender's
                // back-pressure window only advances on this ack, bounding in-flight/buffered data
                // (a slow browser receiver can never be flooded into a stall).
                std::io::Write::flush(&mut sf.w)?;
                write_msg(&mut s, &[TAG_TOKEN, 0], &[])?;
            }
            OP_SHIP_END => {
                // header: (op, fnv_lo, fnv_hi) — integrity over every chunk since begin.
                let sf = ship.take().ok_or_else(|| anyhow!("end without begin"))?;
                let want = (pos as u64) | ((n as u64) << 32);
                let ok = sf.fnv == want;
                if ok {
                    let mut w = sf.w;
                    std::io::Write::flush(&mut w)?;
                    drop(w);
                    std::fs::rename(&sf.part, sf.entry.join(&sf.fname))?;
                    if sf.fname == "model.safetensors" {
                        // Commit marker + provenance: the SOURCE checkpoint's fingerprint.
                        std::fs::write(sf.entry.join("source.fp"), &sf.fp)?;
                    }
                    eprintln!("shard worker: ship done {}", sf.fname);
                } else {
                    eprintln!(
                        "shard worker: ship {} FAILED integrity ({:016x} != {:016x})",
                        sf.fname, sf.fnv, want
                    );
                    // Drop the corrupt partial so the next attempt re-ships from a clean slate.
                    let _ = std::fs::remove_file(&sf.part);
                }
                write_msg(&mut s, &[TAG_TOKEN, 1], &[f32::from_bits(u32::from(!ok))])?;
            }
            OP_INFO => {
                // A loaded stage reports its shard's PROVENANCE (source fp + full depth) —
                // a modelless worker knows neither until its first load.
                let (start, end, total, fp) = loaded
                    .as_ref()
                    .map_or((0, 0, sh.total_layers, sh.ckpt_fp), |l| {
                        (l.start, l.end, l.total, l.src_fp)
                    });
                let desc = format!(
                    "{start}|{end}|{total}|{}|{}|{}|{fp:016x}|{}",
                    u8::from(ctx.subgroups),
                    sh.vram_mb,
                    u8::from(sh.pinned),
                    ctx.backend
                );
                write_msg_bytes(&mut s, &[TAG_INFO, desc.len() as u32], desc.as_bytes())?;
            }
            OP_VOCAB => {
                // Display vocab (fire-and-forget from the coordinator): arm the token→text
                // readout. A bad/truncated blob just leaves it disarmed — never a serving error.
                let raw = read_bytes(&mut s, n)?;
                vocab = crate::vocab::ByteVocab::from_blob(&raw);
            }
            OP_SIGNAL => {
                // Feature-OFF fallback: a native worker built without the `webrtc` feature has no
                // data-channel endpoint, so it answers "unsupported" (status != 0) and the
                // coordinator falls back to TCP/WS for this hop. With the feature on, the pre-lock
                // branch above handled this op (and `continue`d), so this arm is never reached.
                // Consume the payload regardless so the stream stays framed.
                let _payload = read_bytes(&mut s, n)?;
                let msg: &[u8] = b"webrtc endpoint not built on this worker";
                write_msg_bytes(&mut s, &[1, msg.len() as u32], msg)?;
            }
            OP_SHUTDOWN => {
                // Drop the next-hop socket so a P2P chain unwinds stage by stage (each
                // successor's data lane EOFs), then end this session lane; keep listening.
                *chain = None;
                break;
            }
            other => return Err(anyhow!("bad op {other}")),
        }
    }
    Ok(())
}

/// Coordinator-side handle to one worker stage.
pub struct ShardClient {
    s: Conn,
}

impl ShardClient {
    /// Wrap an already-established, already-authenticated connection — the JOIN path, where
    /// the WORKER dialed the coordinator's registry.
    fn from_stream(s: Conn) -> Self {
        Self { s }
    }

    /// Connect to a stage; with env `OSFKB_SHARD_TOKEN` set, authenticates immediately (the
    /// fleet token — same value on every worker and the coordinator).
    pub fn connect(addr: &str) -> Result<Self> {
        let s = dial(addr, std::time::Duration::from_secs(5))
            .with_context(|| format!("connect {addr}"))?;
        let mut c = Self { s };
        if let Ok(t) = std::env::var("OSFKB_SHARD_TOKEN")
            && !t.is_empty()
        {
            hop_hello(&mut c.s, &t).with_context(|| format!("authenticate to {addr}"))?;
        }
        Ok(c)
    }

    /// Assign (or re-assign) the stage's layer range OF the checkpoint `fp` — write half;
    /// loads run in parallel across the fleet when every send precedes the first
    /// [`Self::load_wait`].
    pub fn load_send(&mut self, start: usize, end: usize, fp: u64) -> Result<()> {
        write_msg_bytes(
            &mut self.s,
            &[OP_LOAD, start as u32, end as u32],
            format!("{fp:016x}").as_bytes(),
        )
    }

    /// Block until the stage answered its [`Self::load_send`] assignment.
    pub fn load_wait(&mut self) -> Result<LoadAck> {
        let hdr = read_u32s(&mut self.s, 2)?;
        let p = read_f32s(&mut self.s, hdr[1] as usize)?;
        Ok(match p.first().map(|x| x.to_bits()) {
            Some(0) => LoadAck::Ok,
            Some(2) => LoadAck::NeedWeights,
            Some(3) => LoadAck::WrongCheckpoint,
            _ => LoadAck::Failed,
        })
    }

    /// Open a (possibly resumed) ship. Returns `(bytes_already_present, running_fnv)` the receiver
    /// reconstructed from its partial — the sender skips that prefix and seeds its FNV with it.
    fn ship_begin(
        &mut self,
        fpx: &str,
        range: (usize, usize),
        fname: &str,
        total: u64,
    ) -> Result<(u64, u64)> {
        // 5th field = total file size, so the receiver can PRE-SIZE its accumulation buffer once
        // instead of growing it by doubling — essential in a wasm32 browser worker, whose single
        // `Vec` caps at isize::MAX (~2 GB): a 1.7 GB layer would otherwise overflow when the doubling
        // jumps 1 GB → 2 GB. (Older 4-field metas parse fine; total just defaults to 0 = no pre-size.)
        let meta = format!("{fpx}|{}|{}|{fname}|{total}", range.0, range.1);
        write_msg_bytes(
            &mut self.s,
            &[OP_SHIP_BEGIN, 0, meta.len() as u32],
            meta.as_bytes(),
        )?;
        let r = read_u32s(&mut self.s, 4)?;
        let have = r[0] as u64 | ((r[1] as u64) << 32);
        let fnv = r[2] as u64 | ((r[3] as u64) << 32);
        anyhow::ensure!(
            have != u64::MAX,
            "ship target {fname} busy on the stage (another coordinator is shipping it); retry"
        );
        Ok((have, fnv))
    }

    /// Send one chunk and BLOCK on its ack. Stop-and-wait bounds the receiver's in-flight buffer to a
    /// single chunk — a slow (e.g. browser) receiver applies back-pressure instead of being flooded.
    fn ship_chunk(&mut self, fnv: &mut u64, raw: &[u8]) -> Result<()> {
        *fnv = fnv64_update(*fnv, raw);
        write_msg_bytes(&mut self.s, &[OP_SHIP_CHUNK, 0, raw.len() as u32], raw)?;
        let _ack = read_u32s(&mut self.s, 2)?; // wait until the byte is durably landed
        Ok(())
    }

    fn ship_end(&mut self, fnv: u64) -> Result<()> {
        write_msg(
            &mut self.s,
            &[OP_SHIP_END, fnv as u32, (fnv >> 32) as u32],
            &[],
        )?;
        let hdr = read_u32s(&mut self.s, 2)?;
        let p = read_f32s(&mut self.s, hdr[1] as usize)?;
        anyhow::ensure!(
            p.first().map(|x| x.to_bits()) == Some(0),
            "stage rejected the shipped file (integrity)"
        );
        Ok(())
    }

    /// Ship a small in-memory file (config.json) into the stage's cache for `(fp, range)`. Small
    /// enough that resume is a no-op in practice (whole thing fits one chunk); the (have, fnv)
    /// negotiation still runs so the wire stays uniform.
    pub(crate) fn ship_bytes(
        &mut self,
        fpx: &str,
        range: (usize, usize),
        fname: &str,
        bytes: &[u8],
    ) -> Result<()> {
        let (have, mut fnv) = self.ship_begin(fpx, range, fname, bytes.len() as u64)?;
        let have = (have as usize).min(bytes.len());
        for c in bytes[have..].chunks(SHIP_CHUNK_BYTES) {
            self.ship_chunk(&mut fnv, c)?;
        }
        self.ship_end(fnv)
    }

    /// Stream a planned mini-checkpoint (see [`plan_mini_ckpt`]) out of `src` into the stage's cache.
    /// The wire stream is `header ++ concat(slices)`; RESUMABLE — the receiver reports how many bytes
    /// it already holds and we skip that prefix (seeding the FNV with the receiver's running hash), so
    /// a mid-transfer WAN reset continues instead of restarting. Back-pressured one chunk at a time.
    pub(crate) fn ship_mini_ckpt(
        &mut self,
        fpx: &str,
        range: (usize, usize),
        src: &Path,
        plan: &MiniCkpt,
    ) -> Result<()> {
        use std::io::{Read as _, Seek as _};
        let (have0, mut fnv) =
            self.ship_begin(fpx, range, "model.safetensors", plan.total_bytes)?;
        // A partial that overruns the plan is stale/corrupt (e.g. a different split cached here once):
        // ignore it and re-ship from scratch rather than skip past real bytes.
        let have = if have0 > plan.total_bytes { 0 } else { have0 };
        if have > 0 {
            eprintln!(
                "ship {:?}: resuming at {:.2} / {:.2} GB",
                range,
                have as f64 / 1e9,
                plan.total_bytes as f64 / 1e9
            );
        }
        let mut f = std::fs::File::open(src)?;
        let mut buf = vec![0u8; SHIP_CHUNK_BYTES];
        let mut vpos = 0u64; // absolute position in the virtual header++slices stream
        let (mut sent, mut mark) = (have, have);
        // Header region [0, hlen): in-memory, ship the post-`have` tail (chunked; header can be large).
        let hlen = plan.header.len() as u64;
        if have < hlen {
            for c in plan.header[have as usize..].chunks(SHIP_CHUNK_BYTES) {
                self.ship_chunk(&mut fnv, c)?;
                sent += c.len() as u64;
            }
        }
        vpos += hlen;
        // Slice regions: each maps virtual [vpos, vpos+len) to source bytes [off, off+len).
        for (off, len) in &plan.slices {
            let vstart = vpos;
            vpos += len;
            if have >= vpos {
                continue; // entirely below the resume point
            }
            let skip = have.saturating_sub(vstart); // bytes to skip inside this slice
            f.seek(std::io::SeekFrom::Start(off + skip))?;
            let mut left = len - skip;
            while left > 0 {
                let take = left.min(buf.len() as u64) as usize;
                f.read_exact(&mut buf[..take])?;
                self.ship_chunk(&mut fnv, &buf[..take])?;
                left -= take as u64;
                sent += take as u64;
                if sent - mark >= 512 * 1024 * 1024 {
                    mark = sent;
                    eprintln!(
                        "ship {:?}: {:.1} / {:.1} GB",
                        range,
                        sent as f64 / 1e9,
                        plan.total_bytes as f64 / 1e9
                    );
                }
            }
        }
        self.ship_end(fnv)
    }

    /// Ship the residual stream; get back the next stage's output.
    pub fn step(&mut self, pos: usize, hidden: &[f32]) -> Result<StageOut> {
        write_msg(
            &mut self.s,
            &[OP_STEP, pos as u32, hidden.len() as u32],
            hidden,
        )?;
        let hdr = read_u32s(&mut self.s, 2)?;
        let (tag, n) = (hdr[0], hdr[1] as usize);
        let payload = read_f32s(&mut self.s, n)?;
        Ok(match tag {
            TAG_TOKEN => StageOut::Token(payload[0].to_bits()),
            _ => StageOut::Hidden(payload),
        })
    }

    pub fn reset(&mut self) -> Result<()> {
        write_msg(&mut self.s, &[OP_RESET, 0, 0], &[])?;
        let _ = read_u32s(&mut self.s, 2)?;
        let _ = read_f32s(&mut self.s, 1)?;
        Ok(())
    }

    pub fn shutdown(&mut self) {
        let _ = write_msg(&mut self.s, &[OP_SHUTDOWN, 0, 0], &[]);
    }

    /// Build the stage's batch plans: `k` decode columns + (if `kw > 0`) a `kw`-column WIDE
    /// prefill plan (blocks until built + acked).
    pub fn binit(&mut self, k: usize, kw: usize) -> Result<()> {
        self.binit_grouped(k, kw, 1)
    }

    /// [`Self::binit`] with per-group plans (spec-rollback isolation across in-flight groups).
    pub fn binit_grouped(&mut self, k: usize, kw: usize, n_groups: usize) -> Result<()> {
        assert!(k <= 0xFFFF && n_groups <= 0xFFFF);
        write_msg(
            &mut self.s,
            &[OP_BINIT, (k | (n_groups << 16)) as u32, kw as u32],
            &[],
        )?;
        let _ = read_u32s(&mut self.s, 2)?;
        let _ = read_f32s(&mut self.s, 1)?;
        Ok(())
    }

    /// Zero a sequence slot's recurrent state on the stage (unacked, in-band).
    pub fn zslot(&mut self, slot: u32) -> Result<()> {
        write_msg(&mut self.s, &[OP_ZSLOT, slot, 0], &[])
    }

    /// Ship the donor-display vocab (`id→bytes`) once, fire-and-forget: DISPLAY-ONLY, so a lost or
    /// unsupported blob just means no live text on that donor — never a serving error. Sent in-band
    /// before the first step; the receiver reads it (TCP-ordered) before any `OP_BSTEP`, no ack.
    pub fn send_vocab(&mut self, blob: &[u8]) -> Result<()> {
        write_msg_bytes(&mut self.s, &[OP_VOCAB, 0, blob.len() as u32], blob)
    }

    /// Ask this peer to MAKE a WebRTC offer. Returns its full SDP (ICE candidates already gathered).
    /// Errors if the worker has no WebRTC endpoint (native without the `webrtc` feature).
    pub fn webrtc_offer(&mut self) -> Result<Vec<u8>> {
        write_msg_bytes(&mut self.s, &[OP_SIGNAL, SIG_OFFER, 0], &[])?;
        self.read_signal_reply("offer")
    }

    /// Hand this peer a remote `offer`; returns its full answer SDP.
    pub fn webrtc_answer(&mut self, offer: &[u8]) -> Result<Vec<u8>> {
        write_msg_bytes(
            &mut self.s,
            &[OP_SIGNAL, SIG_ANSWER, offer.len() as u32],
            offer,
        )?;
        self.read_signal_reply("answer")
    }

    /// Hand this peer the remote `answer`, completing the handshake (its channel opens).
    pub fn webrtc_finish(&mut self, answer: &[u8]) -> Result<()> {
        write_msg_bytes(
            &mut self.s,
            &[OP_SIGNAL, SIG_FINISH, answer.len() as u32],
            answer,
        )?;
        self.read_signal_reply("finish").map(|_| ())
    }

    /// Tell this stage (the ANSWERER of a WebRTC hop) to serve its opened data channel as an inbound
    /// lane — a `serve_conn` on the channel, receiving the predecessor's forwarded frames. Acked.
    pub fn rtc_serve(&mut self) -> Result<()> {
        write_msg_bytes(&mut self.s, &[OP_RTC_SERVE, 0, 0], &[])?;
        self.read_signal_reply("rtc-serve").map(|_| ())
    }

    /// Tell this stage (the OFFERER of a WebRTC hop) to install its opened data channel as this
    /// session's forward chain — its STEP/BSTEP outputs then ship over the channel. `sink` marks the
    /// next hop as the coordinator's return sink. Acked.
    pub fn rtc_chain(&mut self, sink: bool) -> Result<()> {
        write_msg_bytes(&mut self.s, &[OP_RTC_CHAIN, u32::from(sink), 0], &[])?;
        self.read_signal_reply("rtc-chain").map(|_| ())
    }

    /// Read a `[status, len] + body` signaling reply; a non-zero status means the worker rejected the
    /// step (e.g. no endpoint), with `body` as the reason.
    fn read_signal_reply(&mut self, step: &str) -> Result<Vec<u8>> {
        let hdr = read_u32s(&mut self.s, 2)?;
        let body = read_bytes(&mut self.s, hdr[1] as usize)?;
        anyhow::ensure!(
            hdr[0] == 0,
            "webrtc {step} rejected by peer: {}",
            String::from_utf8_lossy(&body)
        );
        Ok(body)
    }

    /// Configure TRUE-P2P forward mode on this stage: its STEP/BSTEP outputs ship straight to
    /// `next` (a peer worker, or the coordinator's return sink when `sink`), and in-band
    /// control ops are forwarded down the chain. `":port"` makes the stage derive the host
    /// from this control connection's peer address. Acked; fails if the stage could not dial.
    pub fn chain(&mut self, next: &str, sink: bool) -> Result<()> {
        write_msg_bytes(
            &mut self.s,
            &[OP_CHAIN, u32::from(sink), next.len() as u32],
            next.as_bytes(),
        )?;
        let hdr = read_u32s(&mut self.s, 2)?;
        let payload = read_f32s(&mut self.s, hdr[1] as usize)?;
        anyhow::ensure!(
            payload.first().map(|x| x.to_bits()) == Some(0),
            "stage failed to dial its next hop {next}"
        );
        Ok(())
    }

    /// Clear any forward-mode chain left behind by a previous coordinator session (hub-spoke
    /// sessions send this at start — a crashed P2P coordinator must not leave routing around).
    pub fn chain_clear(&mut self) -> Result<()> {
        write_msg_bytes(&mut self.s, &[OP_CHAIN, 0, 0], &[])?;
        let hdr = read_u32s(&mut self.s, 2)?;
        let _ = read_f32s(&mut self.s, hdr[1] as usize)?;
        Ok(())
    }

    /// The stage's self-description: layer range (0..0 = not loaded), full model depth,
    /// subgroup support, VRAM hint, pinned flag, checkpoint fingerprint, backend.
    pub fn info(&mut self) -> Result<StageInfo> {
        write_msg(&mut self.s, &[OP_INFO, 0, 0], &[])?;
        let hdr = read_u32s(&mut self.s, 2)?;
        anyhow::ensure!(hdr[0] == TAG_INFO, "info expects TAG_INFO");
        let text = String::from_utf8(read_bytes(&mut self.s, hdr[1] as usize)?)
            .context("stage info utf8")?;
        let parts: Vec<&str> = text.splitn(8, '|').collect();
        anyhow::ensure!(parts.len() == 8, "malformed stage info: {text}");
        Ok(StageInfo {
            start: parts[0].parse()?,
            end: parts[1].parse()?,
            total_layers: parts[2].parse()?,
            subgroups: parts[3] == "1",
            vram_mb: parts[4].parse()?,
            pinned: parts[5] == "1",
            ckpt_fp: u64::from_str_radix(parts[6], 16).context("ckpt fingerprint hex")?,
            backend: parts[7].to_string(),
        })
    }

    /// [`Self::step`]'s write half only — P2P mode, where the token returns via the sink.
    pub fn step_send(&mut self, pos: usize, hidden: &[f32]) -> Result<()> {
        write_msg(
            &mut self.s,
            &[OP_STEP, pos as u32, hidden.len() as u32],
            hidden,
        )
    }

    /// Draft `k` tokens with the stage's MTP head, seeded from micro-batch column `col_j`'s
    /// hidden and `first_tok` at absolute position `pos0` (see [`OP_MTP`]).
    pub fn mtp_draft(
        &mut self,
        col_j: usize,
        first_tok: u32,
        pos0: usize,
        k: usize,
    ) -> Result<Vec<u32>> {
        write_msg(
            &mut self.s,
            &[OP_MTP, pos0 as u32, ((col_j << 8) | k) as u32],
            &[f32::from_bits(first_tok)],
        )?;
        let hdr = read_u32s(&mut self.s, 2)?;
        anyhow::ensure!(hdr[0] == TAG_TOKEN, "mtp_draft expects tokens");
        Ok(read_f32s(&mut self.s, hdr[1] as usize)?
            .iter()
            .map(|x| x.to_bits())
            .collect())
    }

    /// Batched draft chains (serving): one entry per slot; returns ncols×k drafted ids.
    pub fn mtp_draft_batch(
        &mut self,
        reqs: &[(u32, u32, usize)],
        k: usize,
    ) -> Result<Vec<Vec<u32>>> {
        let payload: Vec<f32> = reqs
            .iter()
            .flat_map(|(s, t, p)| {
                [
                    f32::from_bits(*s),
                    f32::from_bits(*t),
                    f32::from_bits(*p as u32),
                ]
            })
            .collect();
        write_msg(
            &mut self.s,
            &[OP_MTP_BATCH, reqs.len() as u32, k as u32],
            &payload,
        )?;
        let hdr = read_u32s(&mut self.s, 2)?;
        anyhow::ensure!(hdr[0] == TAG_TOKEN, "mtp_draft_batch expects tokens");
        let flat = read_f32s(&mut self.s, hdr[1] as usize)?;
        Ok(flat
            .chunks(k)
            .map(|c| c.iter().map(|x| x.to_bits()).collect())
            .collect())
    }

    /// Seed a slot's draft slab from verify column `col` (in-band, ordered, no response).
    pub fn mtp_seed(&mut self, slot: u32, col: usize) -> Result<()> {
        self.mtp_seed_grouped(slot, col, 0)
    }

    /// Group-aware seed (reads that group's plan `cur`).
    pub fn mtp_seed_grouped(&mut self, slot: u32, col: usize, group: usize) -> Result<()> {
        write_msg(
            &mut self.s,
            &[OP_MTP_SEED, slot, (col | (group << 16)) as u32],
            &[],
        )
    }

    /// In-band DN-state rollback of `slot` to after column `col` (no response; ordered ahead
    /// of the next step on this connection).
    pub fn dn_restore(&mut self, slot: u32, col: usize) -> Result<()> {
        self.dn_restore_grouped(slot, col, 0)
    }

    /// Group-aware rollback (restores from that group's plan snapshots).
    pub fn dn_restore_grouped(&mut self, slot: u32, col: usize, group: usize) -> Result<()> {
        write_msg(
            &mut self.s,
            &[OP_DNRESTORE, slot, (col | (group << 16)) as u32],
            &[],
        )
    }

    /// Overwrite one block-table row on the stage (unacked; ordered before the next bstep).
    pub fn btab(&mut self, row: u32, blocks: &[u32]) -> Result<()> {
        let bits: Vec<f32> = blocks.iter().map(|b| f32::from_bits(*b)).collect();
        write_msg(&mut self.s, &[OP_BTAB, row, blocks.len() as u32], &bits)
    }

    /// One batched stage step over `cols` (`(pos, btrow, need_logit)`) with the residual payload;
    /// `wide` selects the KC=16 prefill plan.
    pub fn bstep(
        &mut self,
        cols: &[(u32, u32, u32, u32)],
        hidden: &[f32],
        wide: bool,
    ) -> Result<StageBatchOut> {
        self.bstep_grouped(cols, hidden, wide, 0, (false, false, 0))
    }

    /// [`Self::bstep`] against a specific per-group plan (group packed into the header).
    /// Serving-MTP flags pack into header[1]: `ncols(16) | group(8)<<16 | k(6)<<24 |
    /// spec_on<<30 | spec_verify<<31`. `spec_verify` marks a SPAN µbatch (the last stage
    /// accepts + seeds + drafts spontaneously, appending drafts to its token response);
    /// `spec_on` marks a plain µbatch under spec serving (drafts for emitting columns +
    /// draft-KV prefill pairs). Both false ⇒ legacy wire format, bit-identical.
    pub fn bstep_grouped(
        &mut self,
        cols: &[(u32, u32, u32, u32)],
        hidden: &[f32],
        wide: bool,
        group: usize,
        spec: (bool, bool, usize),
    ) -> Result<StageBatchOut> {
        self.bstep_send_grouped(cols, hidden, wide, group, spec)?;
        let hdr = read_u32s(&mut self.s, 2)?;
        let (tag, n) = (hdr[0], hdr[1] as usize);
        let payload = read_f32s(&mut self.s, n)?;
        Ok(match tag {
            TAG_TOKEN => StageBatchOut::Tokens(payload.iter().map(|x| x.to_bits()).collect()),
            _ => StageBatchOut::Hidden(payload),
        })
    }

    /// Write half of [`Self::bstep_grouped`] — P2P mode, where the µbatch's tokens return via
    /// the chain's sink instead of this connection.
    pub fn bstep_send_grouped(
        &mut self,
        cols: &[(u32, u32, u32, u32)],
        hidden: &[f32],
        wide: bool,
        group: usize,
        spec: (bool, bool, usize),
    ) -> Result<()> {
        let (spec_on, spec_verify, spec_k) = spec;
        assert!(cols.len() <= 0xFFFF && group <= 0xFF && spec_k <= 0x3F);
        let op = if wide { OP_BSTEPW } else { OP_BSTEP };
        let mut header = Vec::with_capacity(3 + cols.len() * 4);
        header.extend_from_slice(&[
            op,
            (cols.len()
                | (group << 16)
                | (spec_k << 24)
                | (usize::from(spec_on) << 30)
                | (usize::from(spec_verify) << 31)) as u32,
            hidden.len() as u32,
        ]);
        for &(pos, row, nl, tok) in cols {
            header.extend_from_slice(&[pos, row, nl, tok]);
        }
        write_msg(&mut self.s, &header, hidden)
    }
}

/// Establish a direct WebRTC data channel between two fleet peers by relaying a non-trickle SDP
/// handshake through the coordinator: pull a complete offer from `a`, get `b`'s complete answer,
/// hand it back to `a`. The coordinator moves opaque blobs only — it never parses SDP/ICE. After
/// this returns `Ok`, `a` and `b` share a peer-to-peer channel (subject to ICE succeeding between
/// them). Errors if either peer lacks a WebRTC endpoint, so the caller can fall back to TCP/WS.
pub fn webrtc_pair(a: &mut ShardClient, b: &mut ShardClient) -> Result<()> {
    let offer = a.webrtc_offer()?;
    let answer = b.webrtc_answer(&offer)?;
    a.webrtc_finish(&answer)?;
    Ok(())
}

/// Wire a P3 forward hop `from` → `to` over a **WebRTC data channel** (the last-rung transport for
/// peers that can't reach each other by TCP/WS). Brokers the channel ([`webrtc_pair`]: `from` offers,
/// `to` answers), has `to` serve it as an inbound lane, then installs it as `from`'s forward chain —
/// after which the hop is transport-identical to a dialed TCP/WS chain (bitwise-equal serving).
/// `sink` marks `to` as the coordinator's return sink. Errors (so the caller can fall back to TCP/WS)
/// if either stage lacks a WebRTC endpoint or the channel doesn't come up.
pub fn webrtc_chain(from: &mut ShardClient, to: &mut ShardClient, sink: bool) -> Result<()> {
    webrtc_pair(from, to)?; // `from` = offerer, `to` = answerer
    to.rtc_serve()?; // the answerer serves inbound forwarded frames
    from.rtc_chain(sink)?; // the offerer forwards its output over the channel
    Ok(())
}

/// Bring up the worker→worker hop `i`→`i+1` (both in `clients`) over a WebRTC data channel — the
/// `split_at_mut` needed to borrow the two adjacent stages mutably, gated so a non-`webrtc` build
/// fails loudly instead of silently.
fn webrtc_hop(clients: &mut [ShardClient], i: usize) -> Result<()> {
    #[cfg(feature = "webrtc")]
    {
        let (head, tail) = clients.split_at_mut(i + 1);
        webrtc_chain(&mut head[i], &mut tail[0], false)
    }
    #[cfg(not(feature = "webrtc"))]
    {
        let _ = (clients, i);
        anyhow::bail!("a WebRTC hop was requested but this build lacks the `webrtc` feature")
    }
}

/// A stage's answer to [`ShardClient::load_send`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LoadAck {
    /// Range resident (freshly loaded or already there).
    Ok,
    /// The stage has NO source for this (checkpoint, range) — ship it weights, then retry.
    NeedWeights,
    /// The stage's local `--model` is a DIFFERENT checkpoint than requested.
    WrongCheckpoint,
    /// Load attempted and failed (see the worker's log).
    Failed,
}

/// One stage's self-description, as reported by [`ShardClient::info`].
#[derive(Debug, Clone)]
pub struct StageInfo {
    /// First layer this stage owns (inclusive). `start == end == 0` ⇒ nothing loaded yet.
    pub start: usize,
    /// One past the last layer this stage owns.
    pub end: usize,
    /// The checkpoint's full depth — every stage must agree.
    pub total_layers: usize,
    /// Whether the stage's adapter guarantees subgroups (wide-prefill capability).
    pub subgroups: bool,
    /// Usable-VRAM hint for the auto-split planner (0 = unknown).
    pub vram_mb: u64,
    /// True when the worker fixed its range at startup (`--layers`).
    pub pinned: bool,
    /// Checkpoint identity (config + tensor table) — every stage must agree.
    pub ckpt_fp: u64,
    /// `"<Backend>/<adapter>"`, e.g. `"Metal/Apple M3 Max"` or `"Vulkan/Tesla V100S..."`.
    pub backend: String,
}

/// Check that the coordinator split + worker ranges tile `0..n_layers` in chain order; returns
/// a printable topology map — the heterogeneous-fleet sanity gate before any tokens flow.
pub fn validate_chain(split: usize, infos: &[StageInfo]) -> Result<String> {
    anyhow::ensure!(!infos.is_empty(), "need at least one worker stage");
    let total = infos[0].total_layers;
    let mut cursor = split;
    let mut map = if split > 0 {
        format!("stage 0: coordinator, layers 0..{split} (local)")
    } else {
        "stage 0: coordinator, headless (tokenizer + scheduler only)".to_string()
    };
    for (i, inf) in infos.iter().enumerate() {
        anyhow::ensure!(
            inf.total_layers == total,
            "stage {}: model depth {} != {total} — different checkpoints across the fleet?",
            i + 1,
            inf.total_layers
        );
        anyhow::ensure!(
            inf.ckpt_fp == infos[0].ckpt_fp,
            "stage {}: checkpoint fingerprint {:016x} != stage 1's {:016x} — same architecture, \
             DIFFERENT weights (wrong model dir on that host?)",
            i + 1,
            inf.ckpt_fp,
            infos[0].ckpt_fp
        );
        anyhow::ensure!(
            inf.start == cursor,
            "stage {} covers layers {}..{} but the chain is at layer {cursor} — gap or overlap",
            i + 1,
            inf.start,
            inf.end
        );
        anyhow::ensure!(
            inf.end > inf.start && inf.end <= total,
            "stage {} range {}..{} is invalid for depth {total}",
            i + 1,
            inf.start,
            inf.end
        );
        map.push_str(&format!(
            "\nstage {}: layers {}..{} on {}",
            i + 1,
            inf.start,
            inf.end,
            inf.backend
        ));
        cursor = inf.end;
    }
    anyhow::ensure!(
        cursor == total,
        "chain ends at layer {cursor} but the model has {total} — the last stage must own the head"
    );
    Ok(map)
}

/// Apportion layers `split..total` over the worker stages, in chain order: proportional to
/// reported VRAM when EVERY stage reports one (largest-remainder rounding), equal shares
/// otherwise; every stage gets ≥ 1 layer. This is the auto-split planner — the coordinator
/// runs it whenever the fleet contains a worker that started without `--layers`.
pub fn plan_split(split: usize, total: usize, vram_mb: &[u64]) -> Result<Vec<(usize, usize)>> {
    anyhow::ensure!(!vram_mb.is_empty(), "no worker stages to plan");
    anyhow::ensure!(
        split < total,
        "coordinator owns 0..{split} of {total} layers — nothing left for the workers"
    );
    let (remaining, n) = (total - split, vram_mb.len());
    anyhow::ensure!(
        remaining >= n,
        "{remaining} remaining layers cannot cover {n} stages at ≥1 layer each"
    );
    let weights: Vec<f64> = if vram_mb.iter().all(|v| *v > 0) {
        vram_mb.iter().map(|v| *v as f64).collect()
    } else {
        vec![1.0; n] // any stage with unknown VRAM ⇒ equal shares
    };
    let wsum: f64 = weights.iter().sum();
    let mut counts = Vec::with_capacity(n);
    let mut rems = Vec::with_capacity(n);
    for (i, w) in weights.iter().enumerate() {
        let share = remaining as f64 * w / wsum;
        counts.push(share as usize);
        rems.push((i, share - share.floor()));
    }
    // Leftover layers go to the largest fractional remainders (index-ascending tie-break),
    // cycling defensively against float edge cases.
    rems.sort_by(|a, b| b.1.partial_cmp(&a.1).expect("finite").then(a.0.cmp(&b.0)));
    let mut leftover = remaining - counts.iter().sum::<usize>();
    let mut ri = 0usize;
    while leftover > 0 {
        counts[rems[ri % n].0] += 1;
        ri += 1;
        leftover -= 1;
    }
    // Floor of 1 layer per stage: zeros take from the fattest stage (lowest index on ties).
    while let Some(z) = counts.iter().position(|c| *c == 0) {
        let fat = (0..n).max_by_key(|i| counts[*i]).expect("non-empty counts");
        counts[fat] -= 1;
        counts[z] += 1;
    }
    let mut cursor = split;
    Ok(counts
        .into_iter()
        .map(|c| {
            let r = (cursor, cursor + c);
            cursor += c;
            r
        })
        .collect())
}

/// Fleet negotiation, run by every coordinator connect: a fully PINNED fleet passes through
/// untouched; any AUTO stage (started without `--layers`) triggers a fresh whole-fleet
/// [`plan_split`] assignment — [`OP_LOAD`]s go out to every stage first (loads run in
/// parallel), then all acks are awaited, then infos are re-queried for validation. Reloads of
/// an unchanged range are worker-side no-ops, so re-running against a stable fleet is instant.
pub(crate) fn negotiate_fleet(
    clients: &mut [ShardClient],
    split: usize,
    dir: &Path,
) -> Result<Vec<StageInfo>> {
    let infos: Vec<StageInfo> = clients
        .iter_mut()
        .map(ShardClient::info)
        .collect::<Result<_>>()?;
    // The fleet's checkpoint identity: any stage that HOLDS the checkpoint speaks for it,
    // else the coordinator's own model dir. The coordinator's dir (when hashable) must agree —
    // catches a wrong --model on either side before anything loads.
    let local_fp = checkpoint_fingerprint(dir).ok();
    let (fleet_fp, fleet_total) = match infos.iter().find(|i| i.ckpt_fp != 0) {
        Some(i) => (i.ckpt_fp, i.total_layers),
        None => (
            local_fp.context(
                "no stage holds the checkpoint and the coordinator's --model has no \
                 model.safetensors to ship from",
            )?,
            crate::weights::Lfm2Config::from_json(&std::fs::read(dir.join("config.json"))?)?
                .n_layers,
        ),
    };
    if let Some(lfp) = local_fp {
        anyhow::ensure!(
            lfp == fleet_fp,
            "coordinator --model fingerprint {lfp:016x} != the fleet's {fleet_fp:016x} — \
             different checkpoints"
        );
    }
    if infos.iter().all(|i| i.pinned) {
        return Ok(infos);
    }
    let vram: Vec<u64> = infos.iter().map(|i| i.vram_mb).collect();
    let ranges = plan_split(split, fleet_total, &vram)?;
    eprintln!(
        "auto-split: layers {split}..{fleet_total} over {} stages → {ranges:?}",
        clients.len()
    );
    for (c, (s0, e0)) in clients.iter_mut().zip(&ranges) {
        c.load_send(*s0, *e0, fleet_fp)?;
    }
    let mut needs_ship = Vec::new();
    for (i, c) in clients.iter_mut().enumerate() {
        match c
            .load_wait()
            .with_context(|| format!("stage {} loading layers {:?}", i + 1, ranges[i]))?
        {
            LoadAck::Ok => {}
            LoadAck::NeedWeights => needs_ship.push(i),
            LoadAck::WrongCheckpoint => anyhow::bail!(
                "stage {} holds a DIFFERENT checkpoint than the fleet's {fleet_fp:016x}",
                i + 1
            ),
            LoadAck::Failed => anyhow::bail!(
                "stage {} failed to load layers {:?} (see its log)",
                i + 1,
                ranges[i]
            ),
        }
    }
    // Weight shipping for stages with no local source: config + a synthesized minimal
    // safetensors of exactly their range, then a load retry that must succeed.
    if !needs_ship.is_empty() {
        let src = dir.join("model.safetensors");
        anyhow::ensure!(
            src.exists(),
            "{} stage(s) need weights shipped, but the coordinator's --model has no single-file \
             model.safetensors (HF-sharded sources are not shippable yet — copy those manually)",
            needs_ship.len()
        );
        let cfg_bytes = std::fs::read(dir.join("config.json"))?;
        let fpx = format!("{fleet_fp:016x}");
        for i in needs_ship {
            let range = ranges[i];
            let plan = plan_mini_ckpt(&src, range.0, range.1, fleet_total)?;
            eprintln!(
                "shipping layers {range:?} to stage {} ({:.2} GB)…",
                i + 1,
                plan.total_bytes as f64 / 1e9
            );
            let c = &mut clients[i];
            c.ship_bytes(&fpx, range, "config.json", &cfg_bytes)?;
            c.ship_mini_ckpt(&fpx, range, &src, &plan)?;
            c.load_send(range.0, range.1, fleet_fp)?;
            anyhow::ensure!(
                c.load_wait()? == LoadAck::Ok,
                "stage {} still cannot load {range:?} after shipping (see its log)",
                i + 1
            );
        }
    }
    clients.iter_mut().map(ShardClient::info).collect()
}

/// A resolved worker fleet in chain order: control connections + each stage's data-hop
/// address (what its PREDECESSOR dials for TRUE-P2P forwarding).
pub struct Fleet {
    pub clients: Vec<ShardClient>,
    pub data_addrs: Vec<String>,
}

/// Admit one registry connection: optional [`OP_HELLO`] (mandatory on a tokened fleet), then
/// [`OP_JOIN`]. Returns the worker's data-hop address (announced, or source-IP:port derived).
fn admit_join(s: &mut Conn, peer: std::net::SocketAddr, tok: Option<&str>) -> Result<String> {
    let mut hdr = read_u32s(s, 3)?;
    if hdr[0] == OP_HELLO {
        let raw = read_bytes(s, hdr[2] as usize)?;
        let ok = tok.is_none_or(|t| t.as_bytes() == raw.as_slice());
        write_msg(s, &[TAG_TOKEN, 1], &[f32::from_bits(u32::from(!ok))])?;
        anyhow::ensure!(ok, "bad fleet token");
        hdr = read_u32s(s, 3)?;
    } else {
        anyhow::ensure!(tok.is_none(), "tokenless join on a tokened fleet");
    }
    anyhow::ensure!(hdr[0] == OP_JOIN, "expected OP_JOIN, got op {}", hdr[0]);
    let adv = String::from_utf8(read_bytes(s, hdr[2] as usize)?).context("advertise utf8")?;
    Ok(if adv.is_empty() {
        std::net::SocketAddr::new(peer.ip(), hdr[1] as u16).to_string()
    } else {
        adv
    })
}

/// Accept `n` reverse-registered workers ([`OP_JOIN`]) on `listen` — the fleet REGISTRY.
/// Malformed or unauthenticated dials are logged and skipped; chain order = join order. Each
/// worker's dialed-out socket becomes its control connection, so a NAT'd worker needs NO
/// inbound reachability for control (and none at all under hub-spoke, or as stage 1 in P2P).
pub fn accept_fleet(listen: &str, n: usize) -> Result<Fleet> {
    anyhow::ensure!(n >= 1, "need at least one worker");
    let l = TcpListener::bind(listen)?;
    let tok = std::env::var("OSFKB_SHARD_TOKEN")
        .ok()
        .filter(|t| !t.is_empty());
    eprintln!("fleet registry on {listen}: waiting for {n} worker(s)…");
    let mut clients = Vec::with_capacity(n);
    let mut data_addrs = Vec::with_capacity(n);
    while clients.len() < n {
        let (raw, peer) = l.accept()?;
        let mut s = match accept_conn(raw) {
            Ok(c) => c,
            Err(e) => {
                eprintln!("fleet registry: handshake from {peer} failed: {e:#}");
                continue;
            }
        };
        match admit_join(&mut s, peer, tok.as_deref()) {
            Ok(data) => {
                eprintln!(
                    "fleet registry: worker {} joined from {peer} (data hops via {data})",
                    clients.len() + 1
                );
                data_addrs.push(data);
                clients.push(ShardClient::from_stream(s));
            }
            Err(e) => eprintln!("fleet registry: rejected {peer}: {e:#}"),
        }
    }
    Ok(Fleet {
        clients,
        data_addrs,
    })
}

/// A PERSISTENT fleet registry. Unlike [`accept_fleet`] (bind → take exactly `n` → close), the
/// listener stays open for the coordinator's whole life and a background thread admits EVERY join
/// into a shared pool. This removes the startup RACE — [`Self::form`] waits for a minimum then a
/// settling window and snapshots ALL pooled workers, so an over-subscribed startup forms with
/// everyone instead of first-come-first-served — and enables DYNAMIC GROWTH: workers that join
/// later land in the pool, and [`Self::pending_len`] signals a re-form should absorb them.
pub struct FleetRegistry {
    pending: std::sync::Arc<std::sync::Mutex<std::collections::VecDeque<(ShardClient, String)>>>,
}

impl FleetRegistry {
    /// Bind `listen` and spawn the admit loop (both live until the process exits). Malformed or
    /// unauthenticated dials are logged and skipped, exactly as [`accept_fleet`] does.
    pub fn bind(listen: &str) -> Result<Self> {
        let l = TcpListener::bind(listen)?;
        let tok = std::env::var("OSFKB_SHARD_TOKEN")
            .ok()
            .filter(|t| !t.is_empty());
        eprintln!("fleet registry on {listen}: open — workers may join at any time");
        let pending = std::sync::Arc::new(std::sync::Mutex::new(std::collections::VecDeque::new()));
        let p = pending.clone();
        std::thread::spawn(move || {
            loop {
                let (raw, peer) = match l.accept() {
                    Ok(x) => x,
                    Err(_) => continue,
                };
                let mut s = match accept_conn(raw) {
                    Ok(c) => c,
                    Err(e) => {
                        eprintln!("fleet registry: handshake from {peer} failed: {e:#}");
                        continue;
                    }
                };
                match admit_join(&mut s, peer, tok.as_deref()) {
                    Ok(data) => {
                        eprintln!(
                            "fleet registry: worker joined from {peer} (data hops via {data})"
                        );
                        if let Ok(mut q) = p.lock() {
                            q.push_back((ShardClient::from_stream(s), data));
                        }
                    }
                    Err(e) => eprintln!("fleet registry: rejected {peer}: {e:#}"),
                }
            }
        });
        Ok(Self { pending })
    }

    /// Workers waiting in the pool (joined since the last [`Self::form`]) — the dynamic-grow signal.
    pub fn pending_len(&self) -> usize {
        self.pending.lock().map(|q| q.len()).unwrap_or(0)
    }

    /// Block until at least `min_n` workers are pooled, then a SETTLING window absorbs any still
    /// arriving (so a persistently-reconnecting straggler within one retry interval is caught),
    /// then snapshot ALL pooled workers into a [`Fleet`] (≥ `min_n`). `OSFKB_FLEET_GRACE_MS` tunes
    /// the settle (default 2500 ms, > the worker's ~2 s re-join retry); `0` forms as soon as
    /// `min_n` are present (used by tests for determinism).
    pub fn form(&self, min_n: usize) -> Result<Fleet> {
        anyhow::ensure!(min_n >= 1, "need at least one worker");
        let grace = std::time::Duration::from_millis(
            std::env::var("OSFKB_FLEET_GRACE_MS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(2500),
        );
        eprintln!("fleet registry: waiting for ≥{min_n} worker(s)…");
        while self.pending_len() < min_n {
            std::thread::sleep(std::time::Duration::from_millis(50));
        }
        if !grace.is_zero() {
            // Keep absorbing until `grace` elapses with no new arrival (reset on each join).
            let mut last = self.pending_len();
            let mut quiet = std::time::Instant::now();
            loop {
                std::thread::sleep(std::time::Duration::from_millis(100));
                let now = self.pending_len();
                if now > last {
                    last = now;
                    quiet = std::time::Instant::now();
                } else if quiet.elapsed() >= grace {
                    break;
                }
            }
        }
        let drained: Vec<(ShardClient, String)> = self
            .pending
            .lock()
            .map_err(|_| anyhow!("registry pool poisoned"))?
            .drain(..)
            .collect();
        let (clients, data_addrs): (Vec<_>, Vec<_>) = drained.into_iter().unzip();
        eprintln!(
            "fleet registry: forming a fleet with {} worker(s)",
            clients.len()
        );
        Ok(Fleet {
            clients,
            data_addrs,
        })
    }
}

/// Bind the coordinator's return sink and wire the chain: worker i forwards to worker i+1 and
/// the LAST worker dials the sink (host derived from its control connection's peer address
/// unless `ret` names an explicit `host:port` the last worker can reach). With a fleet token
/// configured, the sink accepts ONLY a dial that opens with a valid [`OP_HELLO`] — a forged
/// connection cannot feed the coordinator tokens. Returns the accepted sink stream.
pub(crate) fn chain_workers(
    clients: &mut [ShardClient],
    workers: &[&str],
    ret: Option<&str>,
    webrtc_hops: &[usize],
    auto_webrtc: bool,
) -> Result<Conn> {
    anyhow::ensure!(clients.len() == workers.len() && !clients.is_empty());
    let (listener, advertise) = match ret {
        Some(a) => {
            let (_, port) = a
                .rsplit_once(':')
                .ok_or_else(|| anyhow!("return address must be host:port"))?;
            let port: u16 = port.parse().context("return port")?;
            (TcpListener::bind(("0.0.0.0", port))?, a.to_string())
        }
        None => {
            let l = TcpListener::bind("0.0.0.0:0")?;
            let port = l.local_addr()?.port();
            (l, format!(":{port}"))
        }
    };
    let n = clients.len();
    // Transport ladder per worker→worker hop `i`→`i+1`:
    //   • `webrtc_hops` forces the hop onto a WebRTC data channel (skips the direct dial);
    //   • otherwise a direct TCP/WS dial is tried, and — with `auto_webrtc` — if that dial fails
    //     (the peers can't reach each other, e.g. both behind NAT) it falls back to WebRTC.
    // The last stage's hop is to the coordinator's return sink (not a WebRTC peer), always TCP/WS.
    for i in 0..n {
        if i + 1 == n {
            clients[i].chain(&advertise, true)?;
        } else if webrtc_hops.contains(&i) {
            webrtc_hop(clients, i)?;
        } else {
            match clients[i].chain(workers[i + 1], false) {
                Ok(()) => {}
                Err(direct_err) if auto_webrtc => {
                    eprintln!(
                        "chain: direct hop {i}→{} unreachable ({direct_err:#}); falling back to WebRTC",
                        i + 1
                    );
                    // If WebRTC also can't bring the hop up, surface the ORIGINAL dial error so a
                    // non-WebRTC fleet's failure mode is unchanged (auto-fallback only ADDS a path).
                    webrtc_hop(clients, i).map_err(|werr| {
                        direct_err.context(format!("webrtc fallback also failed: {werr:#}"))
                    })?;
                }
                Err(direct_err) => return Err(direct_err),
            }
        }
    }
    // The last stage dialed before acking its OP_CHAIN — the connection is already pending.
    // The sink transport is KNOWN from the advertise scheme, so accept it explicitly rather
    // than peek-detecting: a TCP sink (the common, token-less case) sends NOTHING until the
    // first token at inference time, so a peek here would deadlock. A `ws://` sink is safe to
    // peek — the worker's WS client handshake sends its HTTP `GET` on connect — but we route
    // it through the same known-transport branch for symmetry.
    let (raw, _) = listener.accept()?;
    let mut s = if advertise.starts_with("ws://") {
        accept_conn(raw)? // ws handshake byte is already inbound; no deadlock
    } else {
        raw.set_nodelay(true).ok();
        Conn::Tcp(raw)
    };
    if let Ok(tok) = std::env::var("OSFKB_SHARD_TOKEN")
        && !tok.is_empty()
    {
        // Tokened fleet: the first sink message must be the last stage's (write-only) hello.
        let hdr = read_u32s(&mut s, 3)?;
        anyhow::ensure!(
            hdr[0] == OP_HELLO,
            "sink: expected the fleet hello, got op {}",
            hdr[0]
        );
        let raw = read_bytes(&mut s, hdr[2] as usize)?;
        anyhow::ensure!(
            raw == tok.as_bytes(),
            "sink: fleet token rejected — a foreign connection reached the return port"
        );
    }
    Ok(s)
}

/// Read one sink message — `(group, tokens)` from the last stage of a P2P chain.
pub(crate) fn read_done(s: &mut Conn) -> Result<(usize, Vec<u32>)> {
    let hdr = read_u32s(s, 3)?;
    anyhow::ensure!(hdr[0] == TAG_DONE, "sink expects TAG_DONE, got {}", hdr[0]);
    let toks = read_f32s(s, hdr[2] as usize)?
        .iter()
        .map(|x| x.to_bits())
        .collect();
    Ok((hdr[1] as usize, toks))
}

/// DUO-DEVICE single-stream self-speculative decode: the 2-stage MTP round run in ONE process
/// over TWO wgpu devices (no TCP, no worker round-trips — the measured 34.5 GB Q4 footprint
/// rules out a single 32 GB V100, so this is the portable mono architecture). Per round: draft
/// chain on device 1 (owns the `mtp.*` head + final norm), verify µbatch dev0 → 40 KB host
/// bounce → dev1, CPU accept, per-device DN rollback. Greedy outputs are BITWISE the plain
/// greedy outputs (same acceptance argument as [`decode_greedy_mtp`]). The on-GPU
/// accept/restore/seed bricks (see the plan) will replace the remaining per-round CPU work
/// incrementally. Returns (tokens, decode tok/s, acceptance).
pub fn decode_greedy_mtp_duo(
    model_dir: &Path,
    adapters: (usize, usize),
    split: usize,
    prompt: &[u32],
    ngen: usize,
    k_draft: usize,
) -> Result<(Vec<u32>, f64, f64)> {
    anyhow::ensure!(!prompt.is_empty() && k_draft >= 1);
    let kb = k_draft + 1;
    let ctx0 = crate::GpuCtx::new_at(adapters.0)?;
    let ctx1 = crate::GpuCtx::new_at(adapters.1)?;
    let w0 = crate::Weights::load_shard(&ctx0, model_dir, 0, split)?;
    let cfg_nl = {
        // load_shard stores the SHARD's layer count in cfg.n_layers; re-read the checkpoint
        // config for the full depth (cheap: config.json only).
        let cfg =
            crate::weights::Lfm2Config::from_json(&std::fs::read(model_dir.join("config.json"))?)?;
        cfg.n_layers
    };
    let w1 = crate::Weights::load_shard(&ctx1, model_dir, split, cfg_nl)?;
    let gpu0 = crate::Lfm2Gpu::new(&ctx0, w0);
    let gpu1 = crate::Lfm2Gpu::new(&ctx1, w1);
    // (The fast LM head that took this path from 71.5 → 88.7 tok/s is now the engine-wide default
    // and needs no opt-in here — it is batch-invariant at every width, so serving takes it too.
    // See `Q4LmHead::pipeline_for`.)
    let mtp = crate::forward::MtpEngine::new(&ctx1, &gpu1.w)
        .ok_or_else(|| anyhow!("checkpoint has no mtp.* head"))?;
    let slot = 1u32;
    let total = prompt.len() + ngen + k_draft + 1;
    let blocks: Vec<u32> = (0..(total as u32).div_ceil(16)).collect();
    let bp0 = gpu0.make_batch_plan_spec(&ctx0, kb);
    let bp1 = gpu1.make_batch_plan_spec(&ctx1, kb);
    for (ctx, gpu) in [(&ctx0, &gpu0), (&ctx1, &gpu1)] {
        gpu.reset(ctx);
        gpu.zero_dn_slot(ctx, slot as usize);
        gpu.write_btab_row(ctx, slot, &blocks);
    }
    mtp.gpu.write_btab_row(&ctx1, slot, &blocks);
    let brk = std::cell::RefCell::new((0.0f64, 0.0f64, 0.0f64, 0.0f64, 0.0f64, 0.0f64, 0usize));
    let dbrk = std::env::var("OSFKB_DUO_BREAK").is_ok();
    let run_mb = |cols: &[BatchCol]| -> Result<Vec<u32>> {
        let s0 = std::time::Instant::now();
        let h = match gpu0.batch_stage_step(&ctx0, &bp0, cols, None)? {
            crate::forward::StageBatchOut::Hidden(h) => h,
            _ => return Err(anyhow!("stage 0 must produce hidden")),
        };
        let w0 = s0.elapsed().as_secs_f64();
        let b0 = gpu0.take_step_breakdown();
        let s1 = std::time::Instant::now();
        let out = match gpu1.batch_stage_step(&ctx1, &bp1, cols, Some(&h))? {
            crate::forward::StageBatchOut::Tokens(t) => t,
            _ => return Err(anyhow!("stage 1 must produce tokens")),
        };
        let w1 = s1.elapsed().as_secs_f64();
        let b1 = gpu1.take_step_breakdown();
        if dbrk {
            let mut a = brk.borrow_mut();
            a.0 += w0;
            a.1 += b0.2; // dev0 poll+read (GPU wall)
            a.2 += b0.0 + b0.1; // dev0 prep+enc (CPU)
            a.3 += w1;
            a.4 += b1.2;
            a.5 += b1.0 + b1.1;
            a.6 += 1;
        }
        Ok(out)
    };
    // Prefill in ≤kb-column chunks; draft-KV prefill pairs ride dev1 directly.
    let mut col_seed = 0usize;
    let mut t1 = 0u32;
    for chunk in prompt.chunks(kb) {
        let base = (chunk.as_ptr() as usize - prompt.as_ptr() as usize) / 4;
        let cols: Vec<BatchCol> = chunk
            .iter()
            .enumerate()
            .map(|(i, t)| BatchCol::text(*t, (base + i) as u32, slot, base + i + 1 == prompt.len()))
            .collect();
        let out = run_mb(&cols)?;
        col_seed = cols.len() - 1;
        t1 = out[col_seed];
        for (i, _) in chunk.iter().enumerate() {
            let q = base + i;
            if q + 1 < prompt.len() {
                let _ = mtp.draft_chain_seeded(&ctx1, &bp1, i, prompt[q + 1], q, 1)?;
            }
        }
    }
    let mut emitted: Vec<u32> = vec![t1];
    let mut tok = t1;
    let mut pos0 = prompt.len() - 1;
    let (mut drafted, mut accepted) = (0usize, 0usize);
    let t_start = std::time::Instant::now();
    let (mut t_draft, mut t_verify, mut rounds) = (0f64, 0f64, 0usize);
    while emitted.len() < ngen {
        let td = std::time::Instant::now();
        let drafts = mtp.draft_chain_seeded(&ctx1, &bp1, col_seed, tok, pos0, k_draft)?;
        t_draft += td.elapsed().as_secs_f64();
        let cols: Vec<BatchCol> = std::iter::once(tok)
            .chain(drafts.iter().copied())
            .enumerate()
            .map(|(i, t)| BatchCol::text(t, (pos0 + 1 + i) as u32, slot, true))
            .collect();
        let tv = std::time::Instant::now();
        let out = run_mb(&cols)?;
        t_verify += tv.elapsed().as_secs_f64();
        let mut j = 0usize;
        while j < k_draft && out[j] == drafts[j] {
            j += 1;
        }
        drafted += k_draft;
        accepted += j;
        for t in out.iter().take(j + 1) {
            emitted.push(*t);
            if emitted.len() == ngen {
                break;
            }
        }
        if j < k_draft {
            gpu0.dn_restore(&ctx0, &bp0, slot, j);
            gpu1.dn_restore(&ctx1, &bp1, slot, j);
        }
        rounds += 1;
        pos0 = pos0 + 1 + j;
        tok = out[j];
        col_seed = j;
    }
    let secs = t_start.elapsed().as_secs_f64();
    eprintln!(
        "mtp duo round breakdown: {} rounds, draft {:.1}ms verify {:.1}ms per round",
        rounds,
        t_draft * 1000.0 / rounds.max(1) as f64,
        t_verify * 1000.0 / rounds.max(1) as f64
    );
    if dbrk {
        let a = *brk.borrow();
        let n = a.6.max(1) as f64;
        eprintln!(
            "duo verify stage split (per run_mb, avg over {} calls):\n  dev0: wall {:.2}ms [gpu+poll {:.2} cpu-prep+enc {:.2}]\n  dev1: wall {:.2}ms [gpu+poll {:.2} cpu-prep+enc {:.2}]",
            a.6,
            a.0 * 1e3 / n,
            a.1 * 1e3 / n,
            a.2 * 1e3 / n,
            a.3 * 1e3 / n,
            a.4 * 1e3 / n,
            a.5 * 1e3 / n,
        );
    }
    Ok((
        emitted,
        ngen as f64 / secs,
        accepted as f64 / drafted.max(1) as f64,
    ))
}

/// HEADLESS single-stream self-speculative greedy decode: the checkpoint's MTP head drafts
/// `k_draft` tokens on the LAST stage, one (k_draft+1)-column micro-batch verifies them across
/// the chain, the accepted prefix is emitted, and the DeltaNet state rolls back to the last
/// valid column on partial acceptance ([`ShardClient::dn_restore`] — KV rolls back positionally
/// for free). Greedy outputs are BITWISE the plain greedy outputs by construction: a draft is
/// only kept when it EQUALS the model's own argmax at its position. Workers must run with
/// `OSFKB_MTP_SPEC=1` (snapshot plans) and the last stage must hold an `mtp.*` head.
/// Returns (tokens, decode tok/s, draft acceptance rate).
pub fn decode_greedy_mtp(
    workers: &[&str],
    prompt: &[u32],
    ngen: usize,
    k_draft: usize,
) -> Result<(Vec<u32>, f64, f64)> {
    anyhow::ensure!(!workers.is_empty() && !prompt.is_empty() && k_draft >= 1);
    let kb = k_draft + 1; // verify micro-batch width
    let mut clients = workers
        .iter()
        .map(|a| ShardClient::connect(a))
        .collect::<Result<Vec<_>>>()?;
    let slot = 1u32;
    let total = prompt.len() + ngen + k_draft + 1;
    let blocks: Vec<u32> = (0..(total as u32).div_ceil(16)).collect();
    for c in &mut clients {
        c.chain_clear()?; // coordinator-driven MTP is hub-spoke: shed stale P2P routing
        c.reset()?;
        c.binit(kb, 0)?;
        c.zslot(slot)?;
        c.btab(slot, &blocks)?;
    }
    // One micro-batch through the whole chain; the first hop carries token ids (headless).
    let trace = std::env::var("OSFKB_MTP_TRACE").is_ok();
    let run_mb = |clients: &mut [ShardClient], cols: &[(u32, u32, u32, u32)]| -> Result<Vec<u32>> {
        let mut hidden: Vec<f32> = Vec::new();
        let mut spans = Vec::with_capacity(clients.len());
        for (i, c) in clients.iter_mut().enumerate() {
            let t0 = std::time::Instant::now();
            let out = c.bstep(cols, &hidden, false)?;
            spans.push(t0.elapsed().as_secs_f64() * 1e3);
            match out {
                StageBatchOut::Hidden(h) => hidden = h,
                StageBatchOut::Tokens(t) => {
                    anyhow::ensure!(i + 1 == workers.len(), "tokens before the last stage");
                    if trace {
                        eprintln!(
                            "verify stage spans: {}",
                            spans
                                .iter()
                                .map(|v| format!("{v:.1}ms"))
                                .collect::<Vec<_>>()
                                .join(" ")
                        );
                    }
                    return Ok(t);
                }
            }
        }
        Err(anyhow!("last stage returned no tokens"))
    };
    // Prefill in ≤kb-column same-slot chunks; only the last prompt position needs logits.
    let mut col_seed = 0usize;
    let mut t1 = 0u32;
    for chunk in prompt.chunks(kb) {
        let base = (chunk.as_ptr() as usize - prompt.as_ptr() as usize) / 4;
        let cols: Vec<(u32, u32, u32, u32)> = chunk
            .iter()
            .enumerate()
            .map(|(i, t)| {
                let pos = (base + i) as u32;
                let last = base + i + 1 == prompt.len();
                (pos, slot, u32::from(last), *t)
            })
            .collect();
        let out = run_mb(&mut clients, &cols)?;
        col_seed = cols.len() - 1;
        t1 = out[col_seed];
        if std::env::var("OSFKB_MTP_DEBUG_ROUNDS").is_ok() {
            eprintln!("prefill chunk base {base} ncols {} out {out:?}", cols.len());
        }
        // MTP PREFILL: feed the draft layer every (hidden_q, token_{q+1}) prompt pair so its KV
        // covers the prompt (an empty draft context floors acceptance — measured 0.40 without).
        // The chunk's residual streams are resident in the last stage's `cur`; the returned
        // draft is discarded — only the KV write matters. One k=1 chain per pair (v1; a batched
        // no-head prefill op can amortize this later).
        let last = clients.len() - 1;
        for (i, _) in chunk.iter().enumerate() {
            let q = base + i;
            if q + 1 < prompt.len() {
                let _ = clients[last].mtp_draft(i, prompt[q + 1], q, 1)?;
            }
        }
    }
    // Decode rounds: draft on the last stage, verify across the chain, accept the prefix.
    // The prefill's own pick t1 is the FIRST generated token.
    let mut emitted: Vec<u32> = vec![t1];
    let mut tok = t1;
    let mut pos0 = prompt.len() - 1; // position of the last VALID processed input
    let (mut drafted, mut accepted) = (0usize, 0usize);
    // Probe: OSFKB_MTP_SHIFT=-1 drafts at the HIDDEN's position instead of the token's (rope
    // convention check); per-depth acceptance stats separate chain decay from uniform offsets.
    let shift: i64 = std::env::var("OSFKB_MTP_SHIFT")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0);
    let mut depth_hit = vec![0usize; k_draft];
    let mut depth_seen = vec![0usize; k_draft];
    // β-probe: per-depth top-k coverage of the VERIFY argmax (k = 1..3) among the draft's
    // top-3 candidates — conditional on the prefix having been accepted (the tree-DP input).
    let beta_probe = std::env::var("OSFKB_MTP_BETA_PROBE").is_ok();
    let mut cover = vec![[0usize; 3]; k_draft];
    let mut cover_seen = vec![0usize; k_draft];
    let t_start = std::time::Instant::now();
    let last = clients.len() - 1;
    let (mut t_draft, mut t_verify, mut t_restore, mut rounds) = (0f64, 0f64, 0f64, 0usize);
    while emitted.len() < ngen {
        let td = std::time::Instant::now();
        let raw =
            clients[last].mtp_draft(col_seed, tok, (pos0 as i64 + shift) as usize, k_draft)?;
        let (drafts, top3) = raw.split_at(k_draft);
        let drafts = drafts.to_vec();
        t_draft += td.elapsed().as_secs_f64();
        let cols: Vec<(u32, u32, u32, u32)> = std::iter::once(tok)
            .chain(drafts.iter().copied())
            .enumerate()
            .map(|(i, t)| ((pos0 + 1 + i) as u32, slot, 1u32, t))
            .collect();
        let tv = std::time::Instant::now();
        let out = run_mb(&mut clients, &cols)?;
        t_verify += tv.elapsed().as_secs_f64();
        let mut j = 0usize;
        while j < k_draft && out[j] == drafts[j] {
            j += 1;
        }
        // `OSFKB_MTP_DEBUG_ROUNDS=<n>`: dump the first n rounds' raw (drafts, verify picks,
        // accept prefix) — id-0 provenance triage (draft chain vs verify columns).
        if rounds
            < std::env::var("OSFKB_MTP_DEBUG_ROUNDS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0)
        {
            eprintln!(
                "round {rounds}: pos0 {pos0} tok {tok} col_seed {col_seed} drafts {drafts:?} out {:?} j {j}",
                &out[..=k_draft.min(out.len() - 1)]
            );
        }
        for i in 0..k_draft {
            if i <= j {
                depth_seen[i] += 1;
            }
            if i < j {
                depth_hit[i] += 1;
            }
        }
        if beta_probe {
            // Coverage at depth i is meaningful while the prefix (depths < i) was accepted.
            for i in 0..=j.min(k_draft - 1) {
                cover_seen[i] += 1;
                let target = out[i];
                for r in 0..3 {
                    if top3.get(i * 3 + r) == Some(&target) {
                        for c in cover[i].iter_mut().skip(r) {
                            *c += 1;
                        }
                        break;
                    }
                }
            }
        }
        drafted += k_draft;
        accepted += j;
        for t in out.iter().take(j + 1) {
            emitted.push(*t);
            if emitted.len() == ngen {
                break;
            }
        }
        if j < k_draft {
            let tr = std::time::Instant::now();
            for c in &mut clients {
                c.dn_restore(slot, j)?;
            }
            t_restore += tr.elapsed().as_secs_f64();
        }
        rounds += 1;
        pos0 = pos0 + 1 + j;
        tok = out[j];
        col_seed = j;
    }
    let secs = t_start.elapsed().as_secs_f64();
    let per_depth: Vec<String> = (0..k_draft)
        .map(|i| format!("d{}:{}/{}", i + 1, depth_hit[i], depth_seen[i]))
        .collect();
    eprintln!("mtp per-depth acceptance: {}", per_depth.join(" "));
    if beta_probe {
        for d in 0..k_draft {
            let n = cover_seen[d].max(1);
            eprintln!(
                "beta depth {}: top1 {:.2} top2 {:.2} top3 {:.2} (n={})",
                d + 1,
                cover[d][0] as f64 / n as f64,
                cover[d][1] as f64 / n as f64,
                cover[d][2] as f64 / n as f64,
                cover_seen[d]
            );
        }
    }
    eprintln!(
        "mtp round breakdown: {} rounds, draft {:.1}ms verify {:.1}ms restore {:.1}ms per round",
        rounds,
        t_draft * 1000.0 / rounds.max(1) as f64,
        t_verify * 1000.0 / rounds.max(1) as f64,
        t_restore * 1000.0 / rounds.max(1) as f64
    );
    for c in &mut clients {
        c.shutdown();
    }
    Ok((
        emitted,
        ngen as f64 / secs,
        accepted as f64 / drafted.max(1) as f64,
    ))
}

/// MULTI-STREAM self-speculative greedy decode — the serving-MTP round structure at a fixed
/// concurrency: every round runs ONE verify µbatch carrying each stream's (1 + k_draft)
/// same-slot columns, accepts per stream, rolls DN back per slot (per-column snapshots), seeds
/// each slot's draft slab from its accepted column (on-GPU), then drafts ALL streams in one
/// batched chain on the last stage. Greedy outputs stay BITWISE plain-greedy per stream.
/// Returns (per-stream tokens, aggregate tok/s, acceptance).
pub fn decode_greedy_mtp_multi(
    workers: &[&str],
    prompt: &[u32],
    ngen: usize,
    k_draft: usize,
    n_streams: usize,
) -> Result<(Vec<Vec<u32>>, f64, f64)> {
    anyhow::ensure!(!workers.is_empty() && !prompt.is_empty() && k_draft >= 1 && n_streams >= 1);
    let span = 1 + k_draft;
    let kb = span * n_streams;
    let mut clients = workers
        .iter()
        .map(|a| ShardClient::connect(a))
        .collect::<Result<Vec<_>>>()?;
    let blocks_per = (prompt.len() + ngen + k_draft + 2).div_ceil(16);
    for c in &mut clients {
        c.chain_clear()?; // coordinator-driven MTP is hub-spoke: shed stale P2P routing
        c.reset()?;
        c.binit(kb, 0)?;
        for st in 0..n_streams {
            let slot = (st + 1) as u32;
            c.zslot(slot)?;
            let blocks: Vec<u32> = (0..blocks_per as u32)
                .map(|b| st as u32 * blocks_per as u32 + b)
                .collect();
            c.btab(slot, &blocks)?;
        }
    }
    let run_mb = |clients: &mut [ShardClient], cols: &[(u32, u32, u32, u32)]| -> Result<Vec<u32>> {
        let mut hidden: Vec<f32> = Vec::new();
        for (i, c) in clients.iter_mut().enumerate() {
            match c.bstep(cols, &hidden, false)? {
                StageBatchOut::Hidden(h) => hidden = h,
                StageBatchOut::Tokens(t) => {
                    anyhow::ensure!(i + 1 == clients.len(), "tokens before the last stage");
                    return Ok(t);
                }
            }
        }
        Err(anyhow!("last stage returned no tokens"))
    };
    let last = clients.len() - 1;
    // Prefill each stream (same prompt ⇒ same tokens; distinct slots/KV), chunks of ≤ kb cols,
    // then MTP-prefill its draft KV pair-by-pair (cheap 1-layer op).
    let mut t1 = vec![0u32; n_streams];
    #[allow(clippy::needless_range_loop)] // st derives the slot id; t1[st] is one of two uses
    for st in 0..n_streams {
        let slot = (st + 1) as u32;
        for chunk in prompt.chunks(kb) {
            let base = (chunk.as_ptr() as usize - prompt.as_ptr() as usize) / 4;
            let cols: Vec<(u32, u32, u32, u32)> = chunk
                .iter()
                .enumerate()
                .map(|(i, t)| {
                    let pos = (base + i) as u32;
                    let is_last = base + i + 1 == prompt.len();
                    (pos, slot, u32::from(is_last), *t)
                })
                .collect();
            let out = run_mb(&mut clients, &cols)?;
            for (i, _) in chunk.iter().enumerate() {
                let q = base + i;
                if q + 1 < prompt.len() {
                    let _ = clients[last].mtp_draft(i, prompt[q + 1], q, 1)?;
                } else {
                    t1[st] = out[i];
                    clients[last].mtp_seed(slot, i)?;
                }
            }
        }
    }
    let mut emitted: Vec<Vec<u32>> = t1.iter().map(|t| vec![*t]).collect();
    let mut tok: Vec<u32> = t1.clone();
    let mut pos0: Vec<usize> = vec![prompt.len() - 1; n_streams];
    let (mut drafted, mut accepted) = (0usize, 0usize);
    let t_start = std::time::Instant::now();
    while emitted.iter().any(|e| e.len() < ngen) {
        // Phase A: batched draft for every live stream.
        let live: Vec<usize> = (0..n_streams)
            .filter(|s| emitted[*s].len() < ngen)
            .collect();
        let reqs: Vec<(u32, u32, usize)> = live
            .iter()
            .map(|s| ((s + 1) as u32, tok[*s], pos0[*s]))
            .collect();
        let drafts = clients[last].mtp_draft_batch(&reqs, k_draft)?;
        // Phase B: ONE verify µbatch, span columns per live stream.
        let mut cols: Vec<(u32, u32, u32, u32)> = Vec::with_capacity(span * live.len());
        for (li, s) in live.iter().enumerate() {
            let slot = (s + 1) as u32;
            cols.push((pos0[*s] as u32 + 1, slot, 1, tok[*s]));
            for (d, t) in drafts[li].iter().enumerate() {
                cols.push(((pos0[*s] + 2 + d) as u32, slot, 1, *t));
            }
        }
        let out = run_mb(&mut clients, &cols)?;
        // Phase C: accept per stream; rollback + seed at the stream's own column base.
        for (li, s) in live.iter().enumerate() {
            let base = li * span;
            let slot = (s + 1) as u32;
            let mut j = 0usize;
            while j < k_draft && out[base + j] == drafts[li][j] {
                j += 1;
            }
            drafted += k_draft;
            accepted += j;
            for t in out[base..=base + j].iter() {
                if emitted[*s].len() < ngen {
                    emitted[*s].push(*t);
                }
            }
            if j < k_draft {
                for c in &mut clients {
                    c.dn_restore(slot, base + j)?;
                }
            }
            clients[last].mtp_seed(slot, base + j)?;
            pos0[*s] += 1 + j;
            tok[*s] = out[base + j];
        }
    }
    let secs = t_start.elapsed().as_secs_f64();
    let total: usize = emitted.iter().map(Vec::len).sum();
    for c in &mut clients {
        c.shutdown();
    }
    Ok((
        emitted,
        total as f64 / secs,
        accepted as f64 / drafted.max(1) as f64,
    ))
}

/// The sharded pipeline: local stage 0 + remote stages in chain order (last one returns tokens).
pub struct ShardedPipeline {
    ctx: GpuCtx,
    first: Lfm2Gpu,
    remotes: Vec<ShardClient>,
    /// `Some` = TRUE-P2P chain: hops travel worker→worker and tokens return on this sink.
    sink: Option<Conn>,
}

impl ShardedPipeline {
    /// Load stage 0 (`0..split`) locally and connect the remote chain (already-listening workers).
    pub fn connect(model_dir: &Path, split: usize, workers: &[&str]) -> Result<Self> {
        anyhow::ensure!(!workers.is_empty(), "need at least one remote stage");
        let ctx = GpuCtx::new()?;
        let w = Weights::load_shard(&ctx, model_dir, 0, split)?;
        let first = Lfm2Gpu::new(&ctx, w);
        let mut remotes = workers
            .iter()
            .map(|a| ShardClient::connect(a))
            .collect::<Result<Vec<_>>>()?;
        for r in &mut remotes {
            r.chain_clear()?; // hub-spoke session: shed any stale P2P routing
        }
        let infos = negotiate_fleet(&mut remotes, split, model_dir)?;
        eprintln!("{}", validate_chain(split, &infos)?);
        Ok(Self {
            ctx,
            first,
            remotes,
            sink: None,
        })
    }

    /// [`Self::connect`] in TRUE-P2P forward mode: each stage ships its residual straight to
    /// the next peer and only tokens return (to `ret`, or to an auto-bound sink whose address
    /// the last worker derives from its control connection). Validates that the worker ranges
    /// tile the model and prints the (possibly heterogeneous) topology.
    pub fn connect_p2p(
        model_dir: &Path,
        split: usize,
        workers: &[&str],
        ret: Option<&str>,
    ) -> Result<Self> {
        Self::connect_p2p_impl(model_dir, split, workers, workers, ret, &[], false)
    }

    /// [`Self::connect_p2p`] with the worker→worker hops in `webrtc_hops` (indices `i` = hop
    /// `i`→`i+1`) brought up over a WebRTC data channel instead of a dialed TCP/WS socket — the
    /// bitwise-equal proof that hidden states forward correctly over the last-rung P2P transport.
    #[cfg(feature = "webrtc")]
    pub fn connect_p2p_webrtc(
        model_dir: &Path,
        split: usize,
        workers: &[&str],
        ret: Option<&str>,
        webrtc_hops: &[usize],
    ) -> Result<Self> {
        Self::connect_p2p_impl(model_dir, split, workers, workers, ret, webrtc_hops, false)
    }

    /// [`Self::connect_p2p`] with the automatic transport ladder: each worker dials its successor at
    /// `dial_addrs[i+1]` (the successor's *advertised* address, which may differ from the coordinator's
    /// control address and be unreachable across NAT); if that direct dial fails, the hop falls back
    /// to a WebRTC data channel. This is how WebRTC actually gets used in production — as the last rung
    /// when peers can't reach each other directly.
    #[cfg(feature = "webrtc")]
    pub fn connect_p2p_auto(
        model_dir: &Path,
        split: usize,
        control_addrs: &[&str],
        dial_addrs: &[&str],
        ret: Option<&str>,
    ) -> Result<Self> {
        Self::connect_p2p_impl(model_dir, split, control_addrs, dial_addrs, ret, &[], true)
    }

    #[allow(clippy::too_many_arguments)]
    fn connect_p2p_impl(
        model_dir: &Path,
        split: usize,
        control_addrs: &[&str],
        dial_addrs: &[&str],
        ret: Option<&str>,
        webrtc_hops: &[usize],
        auto_webrtc: bool,
    ) -> Result<Self> {
        anyhow::ensure!(!control_addrs.is_empty(), "need at least one remote stage");
        anyhow::ensure!(
            control_addrs.len() == dial_addrs.len(),
            "control and dial address lists must have the same length"
        );
        anyhow::ensure!(split > 0, "the M=1 coordinator owns stage 0 (split ≥ 1)");
        let ctx = GpuCtx::new()?;
        let w = Weights::load_shard(&ctx, model_dir, 0, split)?;
        let first = Lfm2Gpu::new(&ctx, w);
        let mut remotes = control_addrs
            .iter()
            .map(|a| ShardClient::connect(a))
            .collect::<Result<Vec<_>>>()?;
        let infos = negotiate_fleet(&mut remotes, split, model_dir)?;
        eprintln!("{}", validate_chain(split, &infos)?);
        let sink = chain_workers(&mut remotes, dial_addrs, ret, webrtc_hops, auto_webrtc)?;
        Ok(Self {
            ctx,
            first,
            remotes,
            sink: Some(sink),
        })
    }

    /// Greedy decode across the chain. M=1 (each token traverses every stage in turn); the gate
    /// this enables is bitwise equality with the unsharded engine — multi-user pipelining (the
    /// throughput mode: ≥ #stages concurrent requests keep every device busy) drives these same
    /// hops from the scheduler instead.
    pub fn decode_greedy(&mut self, prompt: &[u32], ngen: usize) -> Result<(Vec<u32>, f64)> {
        self.first.reset(&self.ctx);
        for r in &mut self.remotes {
            r.reset()?;
        }
        let total = prompt.len() + ngen;
        let mut out = Vec::with_capacity(ngen);
        let mut tok = prompt[0];
        let t0 = std::time::Instant::now();
        for pos in 0..total - 1 {
            let mut cur = match self.first.stage_forward(&self.ctx, pos, tok, None)? {
                StageOut::Hidden(h) => h,
                StageOut::Token(_) => return Err(anyhow!("stage 0 must not be last")),
            };
            let t = if self.sink.is_some() {
                // TRUE P2P: ONE write to worker 0; the chain forwards stage to stage and the
                // token comes back on the sink — the residual never revisits the coordinator.
                self.remotes[0].step_send(pos, &cur)?;
                let sink = self.sink.as_mut().expect("checked");
                let (_, toks) = read_done(sink)?;
                *toks.first().ok_or_else(|| anyhow!("empty sink message"))?
            } else {
                let mut produced = None;
                for r in &mut self.remotes {
                    match r.step(pos, &cur)? {
                        StageOut::Hidden(h) => cur = h,
                        StageOut::Token(t) => produced = Some(t),
                    }
                }
                produced.ok_or_else(|| anyhow!("last stage returned no token"))?
            };
            tok = if pos + 1 < prompt.len() {
                prompt[pos + 1]
            } else {
                out.push(t);
                t
            };
        }
        let secs = t0.elapsed().as_secs_f64();
        Ok((out, total as f64 / secs))
    }

    /// End the session on every worker (they keep listening for the next coordinator).
    pub fn shutdown_workers(&mut self) {
        for r in &mut self.remotes {
            r.shutdown();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn inf(start: usize, end: usize, total: usize, backend: &str) -> StageInfo {
        StageInfo {
            start,
            end,
            total_layers: total,
            subgroups: true,
            vram_mb: 0,
            pinned: true,
            ckpt_fp: 0xF00D,
            backend: backend.to_string(),
        }
    }

    #[test]
    fn should_reject_checkpoint_fingerprint_mismatch() {
        let mut b = inf(3, 6, 6, "Vulkan/B");
        b.ckpt_fp = 0xBEEF;
        let err = validate_chain(0, &[inf(0, 3, 6, "Metal/A"), b]).unwrap_err();
        assert!(err.to_string().contains("DIFFERENT weights"), "{err}");
    }

    #[test]
    fn should_accept_contiguous_heterogeneous_chain() {
        let map = validate_chain(
            2,
            &[inf(2, 4, 6, "Metal/Apple M3"), inf(4, 6, 6, "Vulkan/V100S")],
        )
        .expect("valid chain");
        assert!(map.contains("Metal/Apple M3") && map.contains("Vulkan/V100S"));
        assert!(map.contains("layers 0..2 (local)"));
    }

    #[test]
    fn should_accept_headless_chain() {
        let map =
            validate_chain(0, &[inf(0, 3, 6, "Metal/A"), inf(3, 6, 6, "Metal/B")]).expect("valid");
        assert!(map.contains("headless"));
    }

    #[test]
    fn should_reject_gap_between_stages() {
        let err = validate_chain(2, &[inf(3, 6, 6, "Metal/A")]).unwrap_err();
        assert!(err.to_string().contains("gap or overlap"), "{err}");
    }

    #[test]
    fn should_reject_uncovered_tail() {
        let err = validate_chain(2, &[inf(2, 5, 6, "Metal/A")]).unwrap_err();
        assert!(err.to_string().contains("must own the head"), "{err}");
    }

    #[test]
    fn should_reject_depth_mismatch_across_fleet() {
        let err =
            validate_chain(0, &[inf(0, 3, 6, "Metal/A"), inf(3, 8, 8, "Vulkan/B")]).unwrap_err();
        assert!(err.to_string().contains("different checkpoints"), "{err}");
    }

    #[test]
    fn should_plan_split_proportional_to_vram() {
        let r = plan_split(0, 6, &[1024, 2048]).expect("plan");
        assert_eq!(r, vec![(0, 2), (2, 6)]);
    }

    #[test]
    fn should_plan_split_equally_when_any_vram_unknown() {
        let r = plan_split(0, 6, &[0, 1024]).expect("plan");
        assert_eq!(r, vec![(0, 3), (3, 6)]);
    }

    #[test]
    fn should_plan_split_start_after_coordinator_share() {
        let r = plan_split(2, 6, &[0, 0]).expect("plan");
        assert_eq!(r, vec![(2, 4), (4, 6)]);
    }

    #[test]
    fn should_plan_split_floor_every_stage_at_one_layer() {
        let r = plan_split(0, 5, &[10 * 1024, 1]).expect("plan");
        assert_eq!(r, vec![(0, 4), (4, 5)]);
    }

    #[test]
    fn should_plan_split_cover_exactly_with_remainders() {
        let r = plan_split(0, 7, &[1, 1, 1]).expect("plan");
        assert_eq!(*r.last().map(|(_, e)| e).expect("stages"), 7);
        assert!(r.windows(2).all(|w| w[0].1 == w[1].0), "contiguous: {r:?}");
        assert!(r.iter().all(|(s, e)| e > s), "≥1 layer each: {r:?}");
    }

    #[test]
    fn should_plan_split_reject_more_stages_than_layers() {
        let err = plan_split(2, 3, &[0, 0]).unwrap_err();
        assert!(err.to_string().contains("cannot cover"), "{err}");
    }

    /// A toy 4-layer safetensors: per-layer tensors + non-layer tensors (embed/norm/mtp).
    fn write_toy_st(path: &Path) -> Vec<(String, Vec<u8>)> {
        use safetensors::tensor::TensorView;
        let mut tensors: Vec<(String, Vec<f32>)> = vec![
            ("model.embed_tokens.weight".into(), vec![1.0, 2.0]),
            ("model.norm.weight".into(), vec![3.0]),
            ("mtp.fc.weight".into(), vec![9.0, 8.0]),
        ];
        for i in 0..4 {
            tensors.push((
                format!("model.layers.{i}.w"),
                vec![i as f32, 10.0 + i as f32],
            ));
        }
        let bytes: Vec<(String, Vec<u8>)> = tensors
            .iter()
            .map(|(k, v)| (k.clone(), bytemuck::cast_slice(v).to_vec()))
            .collect();
        let views: Vec<(String, TensorView<'_>)> = bytes
            .iter()
            .map(|(k, b)| {
                (
                    k.clone(),
                    TensorView::new(safetensors::Dtype::F32, vec![b.len() / 4], b).unwrap(),
                )
            })
            .collect();
        std::fs::write(path, safetensors::tensor::serialize(views, &None).unwrap()).unwrap();
        bytes
    }

    /// Materialize a planned mini-checkpoint into bytes (header + sliced source data).
    fn materialize(src: &Path, plan: &MiniCkpt) -> Vec<u8> {
        use std::io::{Read as _, Seek as _};
        let mut out = plan.header.clone();
        let mut f = std::fs::File::open(src).unwrap();
        for (off, len) in &plan.slices {
            let mut b = vec![0u8; *len as usize];
            f.seek(std::io::SeekFrom::Start(*off)).unwrap();
            f.read_exact(&mut b).unwrap();
            out.extend_from_slice(&b);
        }
        out
    }

    #[test]
    fn should_plan_mini_ckpt_with_exact_middle_range_subset() {
        let dir = std::env::temp_dir().join(format!("lfm2-mini-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let src = dir.join("model.safetensors");
        let source = write_toy_st(&src);
        // MIDDLE range: layers only — no embeddings, no head, no mtp.
        let plan = plan_mini_ckpt(&src, 1, 3, 4).expect("plan");
        let mini = materialize(&src, &plan);
        assert_eq!(mini.len() as u64, plan.total_bytes);
        let st = safetensors::SafeTensors::deserialize(&mini).expect("valid safetensors");
        let mut names: Vec<&str> = st.names().into_iter().map(String::as_str).collect();
        names.sort_unstable();
        assert_eq!(names, vec!["model.layers.1.w", "model.layers.2.w"]);
        for n in names {
            let want = &source.iter().find(|(k, _)| k == n).unwrap().1;
            assert_eq!(st.tensor(n).unwrap().data(), want.as_slice(), "{n} data");
        }
    }

    #[test]
    fn should_plan_mini_ckpt_include_nonlayer_tensors_at_stack_ends() {
        let dir = std::env::temp_dir().join(format!("lfm2-mini2-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let src = dir.join("model.safetensors");
        write_toy_st(&src);
        let first = plan_mini_ckpt(&src, 0, 2, 4).expect("plan");
        let last = plan_mini_ckpt(&src, 2, 4, 4).expect("plan");
        for (plan, layer) in [(&first, "model.layers.0.w"), (&last, "model.layers.3.w")] {
            let mini = materialize(&src, plan);
            let st = safetensors::SafeTensors::deserialize(&mini).expect("valid");
            let names: Vec<String> = st.names().into_iter().cloned().collect();
            assert!(names.iter().any(|n| n == "model.embed_tokens.weight"));
            assert!(names.iter().any(|n| n == "model.norm.weight"));
            assert!(names.iter().any(|n| n == "mtp.fc.weight"));
            assert!(names.iter().any(|n| n == layer));
        }
    }

    #[test]
    fn should_classify_mtp_layers_as_a_non_layer_end_tensor() {
        // Transformer layers parse their index; the MTP head (`mtp.layers.0.*`) shares the
        // `.layers.N.` infix but must NOT be mistaken for transformer layer 0 (that mis-routes it
        // away from the last pipeline stage that needs it — the real 35B-MoE shipping bug).
        assert_eq!(
            tensor_layer("model.layers.7.self_attn.q_proj.weight"),
            Some(7)
        );
        assert_eq!(
            tensor_layer("model.language_model.layers.12.mlp.gate.weight"),
            Some(12)
        );
        assert_eq!(tensor_layer("mtp.layers.0.self_attn.q_proj.weight"), None);
        assert_eq!(tensor_layer("mtp.fc.weight"), None);
        assert_eq!(tensor_layer("model.embed_tokens.weight"), None);
    }

    #[test]
    fn should_ship_mtp_layer_tensors_to_the_last_stage_not_layer_zero() {
        use safetensors::tensor::TensorView;
        // A toy 4-layer model WITH an MTP head whose tensors use the `mtp.layers.0.*` naming.
        let dir = std::env::temp_dir().join(format!("lfm2-mtp-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let src = dir.join("model.safetensors");
        let mut t: Vec<(String, Vec<f32>)> = vec![(
            "mtp.layers.0.self_attn.q_proj.weight".into(),
            vec![1.0, 2.0],
        )];
        for i in 0..4 {
            t.push((format!("model.layers.{i}.w"), vec![i as f32]));
        }
        let bytes: Vec<(String, Vec<u8>)> = t
            .iter()
            .map(|(k, v)| (k.clone(), bytemuck::cast_slice(v).to_vec()))
            .collect();
        let views: Vec<(String, TensorView<'_>)> = bytes
            .iter()
            .map(|(k, b)| {
                (
                    k.clone(),
                    TensorView::new(safetensors::Dtype::F32, vec![b.len() / 4], b).unwrap(),
                )
            })
            .collect();
        std::fs::write(&src, safetensors::tensor::serialize(views, &None).unwrap()).unwrap();
        // The LAST stage (2..4 of 4) must carry the MTP head (the fix). A MIDDLE stage (1..3, neither
        // end) must NOT — proving MTP is treated as an end-tensor, not transformer layer 0 (the bug
        // put it in whatever stage covered layer 0). Non-end duplication onto stage 0 via
        // `touches_ends` is existing, harmless behavior for all non-layer tensors, so not asserted.
        let last = materialize(&src, &plan_mini_ckpt(&src, 2, 4, 4).expect("plan last"));
        let middle = materialize(&src, &plan_mini_ckpt(&src, 1, 3, 4).expect("plan middle"));
        let names = |m: &[u8]| -> Vec<String> {
            safetensors::SafeTensors::deserialize(m)
                .unwrap()
                .names()
                .into_iter()
                .cloned()
                .collect()
        };
        assert!(
            names(&last)
                .iter()
                .any(|n| n == "mtp.layers.0.self_attn.q_proj.weight"),
            "last stage must carry the MTP head"
        );
        assert!(
            !names(&middle).iter().any(|n| n.starts_with("mtp.")),
            "a middle stage must NOT carry the MTP head (it's an end-tensor, not layer 0)"
        );
    }

    #[test]
    fn should_reconstruct_resume_state_from_a_partial() {
        let dir = std::env::temp_dir().join(format!("lfm2-resume-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        // Missing partial → resume at zero.
        let missing = dir.join("absent.part");
        assert_eq!(resume_state(&missing).unwrap(), (0, SHIP_FNV_INIT));
        // Present partial → (byte count, FNV over exactly those bytes).
        let part = dir.join("present.part");
        let bytes: Vec<u8> = (0..5000u32).map(|i| (i * 7 + 1) as u8).collect();
        std::fs::write(&part, &bytes).unwrap();
        let (have, fnv) = resume_state(&part).unwrap();
        assert_eq!(have, bytes.len() as u64);
        assert_eq!(fnv, fnv64_update(SHIP_FNV_INIT, &bytes));
    }

    /// Drives the REAL sender (`ship_mini_ckpt`) against a mock receiver that reports a non-zero
    /// resume offset and acks every chunk. Asserts the sender (1) ships ONLY the bytes past the
    /// resume point, (2) waits for each ack (back-pressure — a desynced ack stream would hang/panic),
    /// and (3) the reassembled prefix+remainder equals the full mini-checkpoint with a matching FNV.
    #[test]
    fn should_resume_shipping_from_a_partial_and_backpressure() {
        use std::io::{Read as _, Write as _};
        use std::net::{TcpListener, TcpStream};
        let dir = std::env::temp_dir().join(format!("lfm2-shipresume-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let src = dir.join("model.safetensors");
        write_toy_st(&src);
        let plan = plan_mini_ckpt(&src, 0, 4, 4).expect("plan");
        let full = materialize(&src, &plan);
        let skip = (full.len() / 3) as u64; // pretend the receiver already holds this prefix
        let fnv_prefix = fnv64_update(SHIP_FNV_INIT, &full[..skip as usize]);

        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let full_c = full.clone();
        let recv = std::thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut hdr = [0u8; 12];
            let mut got: Vec<u8> = Vec::new();
            let mut acks = 0u32;
            loop {
                s.read_exact(&mut hdr).unwrap();
                let op = u32::from_le_bytes(hdr[0..4].try_into().unwrap());
                let a = u32::from_le_bytes(hdr[4..8].try_into().unwrap());
                let b = u32::from_le_bytes(hdr[8..12].try_into().unwrap());
                if op == OP_SHIP_BEGIN {
                    let mut meta = vec![0u8; b as usize];
                    s.read_exact(&mut meta).unwrap();
                    // Report the resume point: (have, fnv-of-have).
                    let reply = [
                        skip as u32,
                        (skip >> 32) as u32,
                        fnv_prefix as u32,
                        (fnv_prefix >> 32) as u32,
                    ];
                    for w in reply {
                        s.write_all(&w.to_le_bytes()).unwrap();
                    }
                } else if op == OP_SHIP_CHUNK {
                    let mut c = vec![0u8; b as usize];
                    s.read_exact(&mut c).unwrap();
                    got.extend_from_slice(&c);
                    // ack (TAG_TOKEN, 0) — the sender blocks on this (back-pressure).
                    s.write_all(&TAG_TOKEN.to_le_bytes()).unwrap();
                    s.write_all(&0u32.to_le_bytes()).unwrap();
                    acks += 1;
                } else if op == OP_SHIP_END {
                    let fnv = (a as u64) | ((b as u64) << 32);
                    // Full reassembly = the pretended prefix + what we received.
                    let reassembled: Vec<u8> = full_c[..skip as usize]
                        .iter()
                        .chain(got.iter())
                        .copied()
                        .collect();
                    let ok =
                        fnv == fnv64_update(SHIP_FNV_INIT, &reassembled) && reassembled == full_c;
                    s.write_all(&TAG_TOKEN.to_le_bytes()).unwrap();
                    s.write_all(&1u32.to_le_bytes()).unwrap();
                    s.write_all(&f32::from_bits(u32::from(!ok)).to_le_bytes())
                        .unwrap();
                    return (got, ok, acks);
                }
            }
        });

        let mut client = ShardClient::from_stream(Conn::Tcp(TcpStream::connect(addr).unwrap()));
        client
            .ship_mini_ckpt("dead00", (0, 4), &src, &plan)
            .expect("ship");
        let (got, ok, acks) = recv.join().unwrap();
        assert!(ok, "receiver rejected integrity");
        assert_eq!(
            got,
            full[skip as usize..],
            "shipped exactly the post-resume remainder"
        );
        assert!(acks >= 1, "chunks were acked (back-pressure handshake ran)");
    }

    /// A 20 MB message crosses tungstenite's 16 MB default frame cap — proves the byte-pipe
    /// framing AND that `ws_config` is applied on BOTH the accept and dial sides.
    #[cfg(feature = "ws")]
    #[test]
    fn should_roundtrip_large_messages_over_the_ws_byte_pipe() {
        let l = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
        let addr = format!("ws://{}", l.local_addr().expect("addr"));
        let server = std::thread::spawn(move || {
            let (s, _) = l.accept().expect("accept");
            let mut c = accept_conn(s).expect("ws upgrade");
            let hdr = read_u32s(&mut c, 3).expect("hdr");
            let body = read_bytes(&mut c, hdr[2] as usize).expect("body");
            write_msg_bytes(&mut c, &[hdr[0], 0, body.len() as u32], &body).expect("echo");
        });
        let mut c = dial(&addr, std::time::Duration::from_secs(5)).expect("dial ws");
        let big: Vec<u8> = (0..20 * 1024 * 1024).map(|i| (i % 251) as u8).collect();
        write_msg_bytes(&mut c, &[0xBEEF, 0, big.len() as u32], &big).expect("send");
        let hdr = read_u32s(&mut c, 3).expect("hdr back");
        assert_eq!(hdr[0], 0xBEEF);
        let echoed = read_bytes(&mut c, hdr[2] as usize).expect("echo body");
        assert_eq!(echoed, big, "byte pipe must be transparent");
        server.join().expect("server thread");
    }
}