ilink-hub 0.3.0

//! Shared Hub state: metrics, long-poll tracking, and the composed [`HubState`]
//! with its `IlinkConnState` / `RoutingState` / `ClientState` sub-states.

use dashmap::DashMap;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex as StdMutex};
use tokio::sync::{broadcast, watch, Mutex, RwLock};

use crate::ilink::{QrLoginUiEvent, UpstreamSink};
use crate::store::Store;

// Hub-internal re-exports (Router, ClientRegistry, PairingRegistry,
// MessageQueue) and the `ilink_status` module come from the crate's `hub` module.
use super::*;

// ─── Concurrency limits ───────────────────────────────────────────────────────

/// Maximum number of concurrent `getupdates` long-polls allowed for a single vtoken.
///
/// A healthy backend has exactly one bridge process polling its vtoken at a time.
/// When two or more bridge processes share one credential/token, they race for
/// the destructive `drain` of the per-vtoken message queue and inbound messages
/// get stolen non-deterministically (split-brain). To stop a malicious or
/// misconfigured client from holding an unbounded number of long-polls (which
/// would saturate the Tokio worker pool), the Hub caps the concurrent poll
/// count per vtoken at this value and rejects additional polls with HTTP 429.
///
/// SEC-003: a single vtoken must not be able to exhaust Hub resources. The
/// cap is intentionally small — anything beyond ~3 is already a configuration
/// problem worth surfacing in the operator logs.
pub const MAX_CONCURRENT_POLLS_PER_VTOKEN: usize = 3;

/// Maximum number of concurrent `getupdates` long-polls allowed Hub-wide,
/// across all vtokens. Each long-poll holds an idle await, a DashMap shard
/// entry, and a `mpsc`/`Notify` channel; 8192 well-behaved clients are
/// already a lot, and any number beyond that strongly suggests either
/// runaway retry storms or a misbehaving bridge. The cap is enforced
/// *before* the per-vtoken gate so a single misconfigured vtoken cannot
/// starve the rest of the fleet.
pub const MAX_HUB_POLLS_DEFAULT: usize = 8192;

// ─── Metrics ──────────────────────────────────────────────────────────────────

/// Bucketed-latency histogram. Uses a fixed log-scale bucket layout to
/// answer "what is the P50 / P95 / P99 of X" without pulling in the full
/// `prometheus-client` crate (which would add a transitive dependency tree
/// of its own). The bucket layout is chosen for sub-second HTTP and CLI
/// latencies; for sub-millisecond paths it's coarse, for multi-second
/// paths it covers the long tail.
///
/// Bucket boundaries in milliseconds. The layout is `[1, 5, 25, 100, 500, 2_500, 10_000, +Inf]`
/// — 8 buckets plus overflow. Suitable for the metrics we currently
/// care about (HTTP round-trips, upstream long-poll cadence, CLI exec).
pub const HISTOGRAM_BUCKETS_MS: &[u64] = &[1, 5, 25, 100, 500, 2_500, 10_000];

/// Number of histogram buckets: one per boundary plus the overflow (+Inf) bucket.
pub const HISTOGRAM_BUCKET_COUNT: usize = HISTOGRAM_BUCKETS_MS.len() + 1;

/// Latency histogram observation. One per metric name; `observe` is hot
/// path (single fetch_add per bucket, no allocation), `snapshot` is the
/// Prometheus-scrape path (called every 15s, not hot).
///
/// N-02: `sum_us` stores the cumulative observation sum in **microseconds**
/// rather than milliseconds. With a millisecond-grained sum, every request
/// under 1 ms (which is the common case for the inbound dispatch path and
/// most in-process work) contributed zero to the sum — the Prometheus
/// `_sum` / `rate(_sum)` reading became useless for sub-millisecond paths.
/// Microsecond precision restores a non-zero sum for those observations,
/// at the cost of a single extra fetch_add per observation. The Prometheus
/// text output still emits `_sum` in milliseconds (`sum_us / 1000`) to
/// match existing dashboards and operator expectations.
#[derive(Debug)]
pub struct LatencyHistogram {
    /// Cumulative count of observations.
    pub count: AtomicU64,
    /// Sum of all observed latencies in **microseconds**. Rendered as
    /// milliseconds (`sum_us / 1000`) in the Prometheus `_sum` field; see
    /// the N-02 note above.
    pub sum_us: AtomicU64,
    /// Fixed-size bucket array: one slot per boundary in `HISTOGRAM_BUCKETS_MS` plus
    /// one overflow (+Inf) slot. Using a fixed-size array avoids heap allocation and
    /// eliminates bounds-check indirection on the hot observe() path.
    pub buckets: [AtomicU64; HISTOGRAM_BUCKET_COUNT],
}

impl Default for LatencyHistogram {
    fn default() -> Self {
        Self::new()
    }
}

impl LatencyHistogram {
    pub fn new() -> Self {
        Self {
            count: AtomicU64::new(0),
            sum_us: AtomicU64::new(0),
            // AtomicU64 is not Copy so we cannot use array-repeat syntax; initialize
            // each element explicitly. HISTOGRAM_BUCKET_COUNT is a compile-time constant.
            buckets: std::array::from_fn(|_| AtomicU64::new(0)),
        }
    }

    /// Record a single observation. `elapsed` is the wall-clock duration of
    /// the operation; we convert to microseconds internally so sub-millisecond
    /// observations remain visible in the cumulative `_sum`.
    ///
    /// Buckets are still bucketed by millisecond (Prometheus convention
    /// keeps the bucket boundaries in the unit operators care about), so
    /// we round down to the nearest millisecond when picking a bucket.
    /// The sum, however, is tracked at microsecond precision so a 500 μs
    /// dispatch contributes 500 to `sum_us`, not 0.
    pub fn observe(&self, elapsed: std::time::Duration) {
        let us = u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX);
        let ms = elapsed.as_millis() as u64;
        self.count.fetch_add(1, Ordering::Relaxed);
        self.sum_us.fetch_add(us, Ordering::Relaxed);
        // Linear scan is O(8) — fine for our small fixed layout.
        for (i, boundary) in HISTOGRAM_BUCKETS_MS.iter().enumerate() {
            if ms <= *boundary {
                self.buckets[i].fetch_add(1, Ordering::Relaxed);
                return;
            }
        }
        // Overflow bucket.
        self.buckets[HISTOGRAM_BUCKETS_MS.len()].fetch_add(1, Ordering::Relaxed);
    }
}

pub struct Metrics {
    pub messages_dispatched: AtomicU64,
    pub messages_dropped: AtomicU64,
    /// User-side (or command) messages taken from upstream and passed into routing
    /// (excludes bot-side echo copies with `message_type == 2`).
    pub upstream_user_messages: AtomicU64,
    /// Total sendmessage calls from backend clients.
    pub sendmessage_total: AtomicU64,
    /// sendmessage calls that were rejected (unknown token, missing context, etc.).
    pub sendmessage_errors: AtomicU64,
    /// Number of QR re-login attempts triggered (manual or automatic).
    pub relogin_attempts: AtomicU64,
    /// Number of message history persist tasks dropped because the semaphore was exhausted.
    /// A non-zero value indicates the DB write pool is too slow for the current message rate.
    pub messages_persist_dropped: AtomicU64,
    /// Latency of `getupdates` long-polls, measured from handler entry to
    /// handler exit (covers both the wakeup wait and the drain).
    pub getupdates_latency_ms: LatencyHistogram,
    /// Latency of upstream `sendmessage` HTTP round-trips. Excludes Hub
    /// internal bookkeeping (context translation, footer append, etc.).
    pub sendmessage_upstream_latency_ms: LatencyHistogram,
    /// Latency of the inbound dispatch pipeline (in-memory only — no DB
    /// I/O). Excludes `tokio::spawn` wall-clock time; this is the
    /// synchronous time from `dispatch_message` entry to its first
    /// `push_to_queue` call.
    pub dispatch_latency_ms: LatencyHistogram,
    /// Unix timestamp (seconds) when this `Metrics` instance was created — emitted as
    /// `*_created` in Prometheus text format so scrape tools can compute per-second rates
    /// correctly even after a process restart.
    pub process_start_unix_secs: f64,
    /// Number of times the broadcast-channel receiver on the dispatcher lagged
    /// (dropped messages due to a full buffer). Incremented inside the
    /// `broadcast::error::RecvError::Lagged` arm.
    pub dispatcher_lagged: AtomicU64,
    /// Failed fire-and-forget persist operations on the unicast forward path.
    pub persist_fire_and_forget_failures_forward: AtomicU64,
    /// Failed fire-and-forget persist operations on the broadcast fan-out path.
    pub persist_fire_and_forget_failures_broadcast: AtomicU64,
    /// Number of quote-reply messages where all three DB fallback layers (timestamp,
    /// content-prefix, footer) returned `None`. Non-zero values indicate quote-reply
    /// routing silently fell back to base routing — useful for detecting DB write gaps.
    pub quote_resolve_miss_total: AtomicU64,
}

impl Metrics {
    pub fn new() -> Self {
        let process_start_unix_secs = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs_f64())
            .unwrap_or(0.0);
        Self {
            messages_dispatched: AtomicU64::new(0),
            messages_dropped: AtomicU64::new(0),
            upstream_user_messages: AtomicU64::new(0),
            sendmessage_total: AtomicU64::new(0),
            sendmessage_errors: AtomicU64::new(0),
            relogin_attempts: AtomicU64::new(0),
            messages_persist_dropped: AtomicU64::new(0),
            getupdates_latency_ms: LatencyHistogram::new(),
            sendmessage_upstream_latency_ms: LatencyHistogram::new(),
            dispatch_latency_ms: LatencyHistogram::new(),
            process_start_unix_secs,
            dispatcher_lagged: AtomicU64::new(0),
            persist_fire_and_forget_failures_forward: AtomicU64::new(0),
            persist_fire_and_forget_failures_broadcast: AtomicU64::new(0),
            quote_resolve_miss_total: AtomicU64::new(0),
        }
    }
}

impl Default for Metrics {
    fn default() -> Self {
        Self::new()
    }
}

/// RAII guard that records elapsed time into a `LatencyHistogram` on drop.
/// Shared between the `getupdates` handler (`HistoGuard` alias in routes.rs)
/// and the inbound dispatch pipeline (`dispatch.rs`).
pub struct LatencyGuard<'a> {
    start: std::time::Instant,
    histogram: &'a LatencyHistogram,
}

impl<'a> LatencyGuard<'a> {
    pub fn new(histogram: &'a LatencyHistogram) -> Self {
        Self {
            start: std::time::Instant::now(),
            histogram,
        }
    }
}

impl Drop for LatencyGuard<'_> {
    fn drop(&mut self) {
        self.histogram.observe(self.start.elapsed());
    }
}

// ─── Concurrent long-poll tracker ─────────────────────────────────────────────

/// Tracks how many `getupdates` long-polls are concurrently active per vtoken.
///
/// A healthy backend has at most one process polling its vtoken at a time. Two or more
/// concurrent polls for the same vtoken mean multiple bridge processes share one
/// credential/token and are competing for the same per-vtoken message queue (`drain` is a
/// destructive read), so inbound messages get stolen non-deterministically. This tracker
/// lets the Hub surface that misconfiguration instead of failing silently.
#[derive(Debug, Default)]
pub struct PollTracker {
    /// Per-vtoken concurrent poll counter. Public for test-only access so
    /// integration tests can poison the mutex to verify the let-Ok
    /// panic-safety path (F-M2-2); production code should only call
    /// `enter` / rely on `Drop`.
    pub counts: StdMutex<HashMap<String, usize>>,
    /// Hub-wide total of in-flight long-polls. Guarded by an AtomicUsize
    /// (lock-free fast path) so the per-request cost of the global gate is
    /// a single fetch_add — orders of magnitude cheaper than the
    /// per-vtoken `StdMutex` we already pay.
    total: AtomicUsize,
    /// Operator-configured Hub-wide cap. Defaults to [`MAX_HUB_POLLS_DEFAULT`].
    hub_cap: AtomicUsize,
}

impl PollTracker {
    /// Set the Hub-wide cap. Must be called once at startup before the Hub
    /// serves any traffic; subsequent changes are not thread-safe w.r.t.
    /// in-flight `enter` calls (they observe either the old or the new cap).
    pub fn set_hub_cap(&self, cap: usize) {
        self.hub_cap.store(cap, Ordering::Relaxed);
    }

    /// Current Hub-wide total of in-flight polls. For metrics / tests.
    pub fn total_polls(&self) -> usize {
        self.total.load(Ordering::Relaxed)
    }

    /// Register a new active poll for `vtoken`. Returns the per-vtoken count
    /// (always >= 1 on the success path, 0 when the per-vtoken mutex is poisoned)
    /// and a guard that decrements *both* the per-vtoken count and the Hub-wide
    /// total on drop.
    ///
    /// The Hub-wide gate runs *before* the per-vtoken gate so a single
    /// misbehaving vtoken cannot starve the rest of the fleet: even if all
    /// per-vtoken slots are full, the Hub still serves polls from other vtokens
    /// up to the global cap.
    ///
    /// F-M2-2: never panic on mutex poisoning. If the per-vtoken `counts` mutex
    /// is poisoned, the guard is still produced but the count is reported as
    /// 0 (so the per-vtoken 429 gate won't trip) and the drop handler becomes
    /// a best-effort no-op. A poisoned `counts` map is a process-wide bug, but
    /// it must not take the Tokio worker down on every subsequent long-poll.
    /// The Hub-wide counter is lock-free AtomicUsize, so it cannot poison.
    pub fn enter(self: &Arc<Self>, vtoken: &str) -> EnterOutcome {
        // Hub-wide gate first. fetch_add returns the previous value; we then
        // check the new total against the cap. If we're over, decrement back
        // and refuse. The decrement is safe because we just incremented — the
        // counter cannot have wrapped in between on any platform usize can
        // represent.
        let cap = self.hub_cap.load(Ordering::Relaxed);
        let prev_total = self.total.fetch_add(1, Ordering::AcqRel);
        if prev_total >= cap {
            self.total.fetch_sub(1, Ordering::AcqRel);
            return EnterOutcome::HubLimitReached {
                total: prev_total,
                cap,
            };
        }

        let count = {
            let Ok(mut counts) = self.counts.lock() else {
                // Do NOT roll back the Hub-wide increment here. The guard
                // we return is responsible for decrementing `total` on drop,
                // keeping the counter accurate for the duration of the
                // (poisoned) request. Rolling back here and then letting the
                // guard decrement again would cause an underflow.
                return EnterOutcome::Poisoned {
                    guard: PollGuard {
                        tracker: Arc::clone(self),
                        vtoken: vtoken.to_string(),
                    },
                };
            };
            let c = counts.entry(vtoken.to_string()).or_insert(0);
            *c += 1;
            *c
        };
        EnterOutcome::Ok {
            per_vtoken: count,
            guard: PollGuard {
                tracker: Arc::clone(self),
                vtoken: vtoken.to_string(),
            },
        }
    }
}

/// Result of [`PollTracker::enter`]. The caller inspects the variant and
/// either serves the long-poll (Ok), rejects it as 503 (HubLimitReached),
/// or treats the per-vtoken count as advisory and serves anyway (Poisoned).
#[derive(Debug)]
pub enum EnterOutcome {
    Ok { per_vtoken: usize, guard: PollGuard },
    HubLimitReached { total: usize, cap: usize },
    Poisoned { guard: PollGuard },
}

impl EnterOutcome {
    /// Convenience: extract the guard, regardless of variant. `HubLimitReached`
    /// has no guard (the Hub-wide increment was rolled back); callers must
    /// surface the rejection *before* calling this.
    #[allow(dead_code)]
    pub fn guard(self) -> Option<PollGuard> {
        match self {
            EnterOutcome::Ok { guard, .. } | EnterOutcome::Poisoned { guard } => Some(guard),
            EnterOutcome::HubLimitReached { .. } => None,
        }
    }
}

/// RAII guard returned by [`PollTracker::enter`]; decrements the per-vtoken poll count
/// and the Hub-wide total when the long-poll handler returns (success, timeout, shutdown,
/// or client disconnect).
#[derive(Debug)]
pub struct PollGuard {
    tracker: Arc<PollTracker>,
    vtoken: String,
}

impl Drop for PollGuard {
    fn drop(&mut self) {
        // F-M2-2: best-effort decrement; a poisoned mutex here would otherwise
        // propagate a panic into the Tokio worker that called the handler.
        // The Hub-wide counter cannot be poisoned (it's an AtomicUsize), so
        // we always decrement it. saturating_sub guards against the
        // (theoretically impossible) underflow.
        self.tracker.total.fetch_sub(1, Ordering::AcqRel);
        let Ok(mut counts) = self.tracker.counts.lock() else {
            return;
        };
        if let Some(c) = counts.get_mut(&self.vtoken) {
            *c = c.saturating_sub(1);
            if *c == 0 {
                counts.remove(&self.vtoken);
            }
        }
    }
}

// ─── Shared Hub State ─────────────────────────────────────────────────────────

/// State tied to the iLink upstream WebSocket connection.
///
/// Anything that mutates only when iLink connects, logs in, or sends a QR-ready
/// event lives here. Callers that need to send a message upstream, observe a QR
/// login, or trigger a re-login take a reference to this sub-state rather than
/// touching the whole `HubState`.
///
/// `upstream` is held as a trait object so end-to-end tests can inject a
/// recording mock in place of [`UpstreamClient`]. The polling loop owns the
/// concrete `UpstreamClient` separately and does not go through this field.
/// The observability counters on the polling loop are exposed through the
/// `UpstreamSink::polls_ok` / `polls_err` / `relogin_attempts` accessors.
pub struct IlinkConnState {
    pub upstream: Arc<dyn UpstreamSink>,
    /// Shared with Axum graceful shutdown; long-poll handlers exit early when this becomes `true`.
    pub shutdown: watch::Receiver<bool>,
    /// Current iLink upstream status (see [`ilink_status`] constants).
    pub ilink_status: Arc<AtomicU8>,
    /// Broadcasts QR login UI events to SSE subscribers.
    pub qr_tx: broadcast::Sender<QrLoginUiEvent>,
    /// Last QR Ready event — replayed to new SSE subscribers that connect after it was sent.
    ///
    /// Uses a synchronous `std::sync::Mutex` because the value is `Option<Event>`
    /// (cheap, no async work needed under the lock) and the SSE handler clones
    /// the value into the response stream before returning. A `tokio` mutex
    /// here would only add scheduling overhead.
    pub qr_last_ready: Arc<std::sync::Mutex<Option<QrLoginUiEvent>>>,
    /// Signals the polling loop to initiate a fresh QR re-login.
    pub relogin_tx: broadcast::Sender<()>,
    /// Single-use, short-lived tickets that authenticate the QR SSE stream
    /// without putting the admin token in the URL. See [`SseTicketStore`].
    pub qr_ticket: crate::server::sse_ticket::SseTicketStore,
}

impl IlinkConnState {
    pub(crate) fn new(upstream: Arc<dyn UpstreamSink>, shutdown: watch::Receiver<bool>) -> Self {
        let (qr_tx, _) = broadcast::channel(16);
        let (relogin_tx, _) = broadcast::channel(4);
        Self {
            upstream,
            shutdown,
            ilink_status: Arc::new(AtomicU8::new(ilink_status::UNKNOWN)),
            qr_tx,
            qr_last_ready: Arc::new(std::sync::Mutex::new(None)),
            relogin_tx,
            qr_ticket: crate::server::sse_ticket::SseTicketStore::new(),
        }
    }
}

/// Routing-layer state: per-message dispatch decisions and conversation vctx
/// mapping. Pure in-memory; no I/O.
pub struct RoutingState {
    pub router: Mutex<Router>,
}

impl RoutingState {
    pub(crate) fn new() -> Self {
        Self {
            router: Mutex::new(Router::new(None)),
        }
    }
}

/// Registered backend clients, paired devices, the per-vtoken message queue,
/// and long-poll concurrency tracking.
pub struct ClientState {
    pub registry: RwLock<ClientRegistry>,
    pub pairing: RwLock<PairingRegistry>,
    /// Notified whenever a pairing session transitions state (scanned/confirmed).
    /// `get_qrcode_status` waits on this instead of sleep-polling every 1s.
    pub pairing_notify: Arc<tokio::sync::Notify>,
    pub queue: Arc<dyn MessageQueue>,
    /// Tracks concurrent `getupdates` long-polls per vtoken to detect bridges that share one
    /// credential/token (queue split-brain).
    pub poll_tracker: Arc<PollTracker>,
    /// Lock-free last-seen timestamps (Unix seconds) per vtoken.
    /// `getupdates` updates this atomically without acquiring the registry write lock.
    /// `spawn_health_checker` reads it to mark clients offline.
    pub last_seen: Arc<DashMap<String, AtomicU64>>,
}

impl ClientState {
    pub(crate) fn new(queue: Arc<dyn MessageQueue>) -> Self {
        let poll_tracker = Arc::new(PollTracker::default());
        // Initialize the Hub-wide cap to the default. Operators can override
        // via `ILINK_MAX_HUB_POLLS`; see [`crate::runtime::serve::RuntimeConfig`].
        poll_tracker.set_hub_cap(MAX_HUB_POLLS_DEFAULT);
        Self {
            registry: RwLock::new(ClientRegistry::new()),
            pairing: RwLock::new(PairingRegistry::new()),
            pairing_notify: Arc::new(tokio::sync::Notify::new()),
            queue,
            poll_tracker,
            last_seen: Arc::new(DashMap::new()),
        }
    }
}

/// Maximum number of concurrent fire-and-forget persist tasks. Applying this limit
/// bounds the number of SQLite pool-acquire waiters during message bursts and
/// prevents them from growing without bound. Tasks that cannot acquire a permit
/// drop their work and increment the relevant failure counter — the same observable
/// behaviour as before, but now with natural backpressure.
const MAX_CONCURRENT_PERSIST_TASKS: usize = 32;

/// Top-level hub state. Groups related state into cohesive sub-states so that
/// internal helpers (dispatcher, hub-command handler, etc.) take the smallest
/// slice they need instead of the entire blob.
///
/// External callers (server routes, pairing, etc.) continue to access fields
/// through the same `state.field` paths they always have — the sub-state
/// fields are re-exported as direct `pub` fields on `HubState` for backward
/// compatibility. New code is encouraged to take `&RoutingState` /
/// `&IlinkConnState` / `&ClientState` parameters to make the dependency
/// explicit.
/// Admin authentication configuration, parsed once at startup from env vars and
/// stored in `HubState`. Routes call `state.admin.check(headers)` instead of
/// reading env vars at request time via `OnceLock`.
#[derive(Debug, Clone)]
pub struct AdminConfig {
    /// The required Bearer token value, or `None` if `ILINK_ADMIN_TOKEN` is unset.
    pub token: Option<String>,
    /// True when `ILINK_ADMIN_INSECURE_NO_AUTH` is set and no token is configured.
    pub insecure_no_auth: bool,
    /// Parsed once at startup from `ILINKHUB_OUTBOUND_ORIGIN_LABEL`.
    /// `None` = env var absent (default behaviour); `Some(v)` = explicit override.
    pub outbound_origin_label: Option<String>,
}

impl AdminConfig {
    pub fn from_env() -> Self {
        let token = std::env::var("ILINK_ADMIN_TOKEN")
            .ok()
            .filter(|s| !s.is_empty());
        let insecure_no_auth = token.is_none()
            && std::env::var("ILINK_ADMIN_INSECURE_NO_AUTH")
                .ok()
                .map(|v| matches!(v.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes"))
                .unwrap_or(false);
        let outbound_origin_label = std::env::var("ILINKHUB_OUTBOUND_ORIGIN_LABEL").ok();
        Self {
            token,
            insecure_no_auth,
            outbound_origin_label,
        }
    }
}

pub struct HubState {
    /// iLink upstream connection and shutdown signal.
    pub ilink: IlinkConnState,
    /// Per-message routing, vctx mapping, and quote-reply tracking.
    pub routing: RoutingState,
    /// Registered clients, paired devices, message queue, long-poll tracking.
    pub clients: ClientState,
    /// Persistent store (SQLx pool-backed). Cross-cutting; not part of any sub-state.
    pub store: Arc<Store>,
    /// Observability counters. Cross-cutting; not part of any sub-state.
    pub metrics: Arc<Metrics>,
    /// Backpressure semaphore for fire-and-forget context-token persist tasks.
    pub persist_sem: Arc<tokio::sync::Semaphore>,
    /// Per-process random secret shared with the in-process relay client so the Hub
    /// can distinguish trusted relay-forwarded XFF headers from local-process spoofing.
    /// The relay client injects `X-Ilink-Relay-Secret: <secret>` on every forwarded
    /// request; Hub's pair_confirm trusts X-Forwarded-For only when this matches.
    pub relay_secret: String,
    /// Admin authentication config, parsed once at startup.
    pub admin: AdminConfig,
    /// Pending Agent-to-Agent reply waiters.  MCP `call_agent` registers a
    /// oneshot here; the target Agent's `sendmessage` resolves it.
    pub a2a_waiter: Arc<crate::mcp::A2aWaiter>,
}

impl HubState {
    /// Build a new [`HubState`]. The `relay_secret` must be supplied by the
    /// caller; use [`crate::paths::load_or_create_relay_secret`] for the
    /// standard "load from disk, else generate and persist" path. We pass
    /// it in rather than computing it here so the constructor stays sync
    /// (callers from async contexts can `await` the I/O themselves) and so
    /// tests can pin a deterministic value.
    pub fn new(
        upstream: Arc<dyn UpstreamSink>,
        store: Arc<Store>,
        queue: Arc<dyn MessageQueue>,
        shutdown: watch::Receiver<bool>,
        relay_secret: String,
        admin: AdminConfig,
    ) -> Arc<Self> {
        Arc::new(Self {
            ilink: IlinkConnState::new(upstream, shutdown),
            routing: RoutingState::new(),
            clients: ClientState::new(queue),
            store,
            metrics: Arc::new(Metrics::new()),
            persist_sem: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_PERSIST_TASKS)),
            relay_secret,
            admin,
            a2a_waiter: Arc::new(crate::mcp::A2aWaiter::new()),
        })
    }

    /// Acquire `router` and `registry` in the canonical lock order (router
    /// before registry) and run `f` with both guards. The guards are dropped
    /// before `f` returns, so any `.await` inside `f` happens with no HubState
    /// lock held. Use this whenever a code path needs both locks; the
    /// pre-facade code had several places that re-acquired them in a
    /// different order or held one across an await, which is a deadlock
    /// hazard.
    ///
    /// Returns `None` if either lock is poisoned. Callers may treat that as
    /// a transient error and surface a 5xx — there is no recovery path that
    /// keeps the request meaningful once the lock is poisoned.
    pub async fn with_router_and_registry<R>(
        self: &Arc<Self>,
        f: impl FnOnce(&mut Router, &ClientRegistry) -> R,
    ) -> Option<R> {
        let mut router_guard = self.routing.router.lock().await;
        let registry_guard = self.clients.registry.read().await;
        let result = f(&mut router_guard, &registry_guard);
        drop(registry_guard);
        drop(router_guard);
        Some(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::Ordering;
    use std::time::Duration;

    // ─── LatencyHistogram::observe ───────────────────────────────────────────

    #[test]
    fn observe_increments_count_and_sum_us() {
        let h = LatencyHistogram::new();
        h.observe(Duration::from_millis(10));
        assert_eq!(h.count.load(Ordering::Relaxed), 1);
        // 10ms = 10_000 µs
        assert_eq!(h.sum_us.load(Ordering::Relaxed), 10_000);
    }

    /// N-02 regression: sub-millisecond observations must contribute a
    /// non-zero value to `sum_us` so rate calculations are meaningful.
    #[test]
    fn observe_sub_ms_contributes_non_zero_to_sum_us() {
        let h = LatencyHistogram::new();
        h.observe(Duration::from_micros(500)); // 0.5 ms = 500 µs
        assert_eq!(
            h.sum_us.load(Ordering::Relaxed),
            500,
            "500µs observation must add 500 to sum_us, not 0"
        );
    }

    /// 1 ms sits exactly on the first boundary — must land in bucket[0] (≤ 1ms).
    #[test]
    fn observe_exactly_at_boundary_lands_in_lower_bucket() {
        let h = LatencyHistogram::new();
        h.observe(Duration::from_millis(1));
        assert_eq!(
            h.buckets[0].load(Ordering::Relaxed),
            1,
            "1ms (= boundary) must land in bucket[0] (≤ 1ms)"
        );
        assert_eq!(
            h.buckets[1].load(Ordering::Relaxed),
            0,
            "bucket[1] must be empty"
        );
    }

    /// 2 ms is above the 1 ms boundary — must NOT land in bucket[0].
    #[test]
    fn observe_just_above_boundary_lands_in_next_bucket() {
        // HISTOGRAM_BUCKETS_MS = [1, 5, 25, 100, 500, 2500, 10000]
        // 2ms > 1ms → not bucket[0]; 2ms ≤ 5ms → bucket[1]
        let h = LatencyHistogram::new();
        h.observe(Duration::from_millis(2));
        assert_eq!(
            h.buckets[0].load(Ordering::Relaxed),
            0,
            "2ms must NOT land in bucket[0] (> 1ms boundary)"
        );
        assert_eq!(
            h.buckets[1].load(Ordering::Relaxed),
            1,
            "2ms must land in bucket[1] (≤ 5ms)"
        );
    }

    /// Values above all boundaries must land in the overflow (+Inf) bucket.
    #[test]
    fn observe_above_all_boundaries_lands_in_overflow_bucket() {
        let h = LatencyHistogram::new();
        h.observe(Duration::from_millis(11_000)); // > 10_000ms — overflow
        let overflow_idx = HISTOGRAM_BUCKETS_MS.len();
        assert_eq!(
            h.buckets[overflow_idx].load(Ordering::Relaxed),
            1,
            "11s must land in overflow bucket"
        );
        for i in 0..HISTOGRAM_BUCKETS_MS.len() {
            assert_eq!(
                h.buckets[i].load(Ordering::Relaxed),
                0,
                "bucket[{i}] must be empty"
            );
        }
    }

    // ─── PollTracker::enter ──────────────────────────────────────────────────

    fn make_tracker(cap: usize) -> Arc<PollTracker> {
        let t = Arc::new(PollTracker::default());
        t.set_hub_cap(cap);
        t
    }

    #[test]
    fn enter_first_poll_returns_ok_with_per_vtoken_count_one() {
        let t = make_tracker(10);
        let outcome = t.enter("vtok-A");
        assert!(
            matches!(outcome, EnterOutcome::Ok { per_vtoken: 1, .. }),
            "first enter must be Ok with per_vtoken=1"
        );
        assert_eq!(t.total_polls(), 1);
    }

    #[test]
    fn enter_second_poll_same_vtoken_returns_count_two() {
        let t = make_tracker(10);
        let _g1 = t.enter("vtok-A");
        let outcome = t.enter("vtok-A");
        assert!(
            matches!(outcome, EnterOutcome::Ok { per_vtoken: 2, .. }),
            "second enter on same vtoken must be Ok with per_vtoken=2"
        );
    }

    #[test]
    fn enter_hub_cap_exceeded_returns_hub_limit_reached() {
        let t = make_tracker(1);
        let _g1 = t.enter("vtok-A"); // fills the cap
        let outcome = t.enter("vtok-B");
        assert!(
            matches!(outcome, EnterOutcome::HubLimitReached { .. }),
            "cap=1 with 1 active poll must reject next enter"
        );
    }

    #[test]
    fn enter_hub_limit_reached_rolls_back_total_counter() {
        let t = make_tracker(1);
        let _g1 = t.enter("vtok-A"); // fills cap
        let _rejected = t.enter("vtok-B"); // triggers rollback
                                           // total must be 1 (not 2) because the increment was rolled back
        assert_eq!(
            t.total_polls(),
            1,
            "rejected enter must roll back Hub-wide total"
        );
    }

    #[test]
    fn poll_guard_drop_decrements_total_and_removes_vtoken_entry() {
        let t = make_tracker(10);
        {
            let outcome = t.enter("vtok-A");
            assert_eq!(t.total_polls(), 1);
            drop(outcome); // drops EnterOutcome → drops PollGuard → runs Drop impl
        }
        assert_eq!(
            t.total_polls(),
            0,
            "PollGuard drop must decrement Hub-wide total to zero"
        );
        let counts = t.counts.lock().unwrap();
        assert!(
            !counts.contains_key("vtok-A"),
            "zero-count vtoken entry must be removed from map"
        );
    }

    #[test]
    fn enter_outcome_guard_some_for_ok_none_for_hub_limit_reached() {
        let t = make_tracker(10);
        let ok = t.enter("vtok-A");
        assert!(ok.guard().is_some(), "Ok variant must return Some(guard)");

        let t2 = make_tracker(0);
        let rejected = t2.enter("vtok-X");
        assert!(
            rejected.guard().is_none(),
            "HubLimitReached must return None from guard()"
        );
    }
}