iicp_client/

node.rs

// SPDX-License-Identifier: Apache-2.0
//! IICP provider node — registration, heartbeats, and task serving.
//!
//! Implements:
//! - `GET  /iicp/health`   — liveness / capacity (always 200)
//! - `GET  /metrics`       — Prometheus text (503 if `metrics` feature absent)
//! - `POST /v1/task`       — task handler with concurrency gate (IICP-E021),
//!   nonce replay protection (IICP-E011), and W3C traceparent propagation.

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use axum::{
    extract::State,
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response},
    routing::{get, post},
    Json, Router,
};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::net::TcpListener;
use tokio::sync::Mutex;

use crate::errors::{IicpError, Result};

const DEFAULT_DIRECTORY: &str = "https://iicp.network/api";
const HEARTBEAT_INTERVAL_SECS: u64 = 30;
const NONCE_TTL_SECS: u64 = 300;

/// Configuration for an IICP provider node.
#[derive(Debug, Clone)]
pub struct NodeConfig {
    pub node_id: String,
    pub endpoint: String,
    pub intent: String,
    pub model: Option<String>,
    pub region: Option<String>,
    pub capabilities: Vec<String>,
    pub directory_url: String,
    pub timeout_ms: u64,
    /// Maximum concurrent tasks; excess requests receive 429 IICP-E021.
    pub max_concurrent: usize,
    /// Tokens-per-minute capacity declared to directory (`limits.tokens_per_min`).
    pub tokens_per_min: u32,
    /// Per-request token cap declared on the capability object (`capabilities[].max_tokens`).
    pub max_tokens: u32,
    /// Optional native IICP binary endpoint (spec/iicp-dir.md v0.7.0).
    /// Scheme MUST be `iicp://` (plaintext) or `iicpsec://` (TLS).
    /// Default IICP port is 9484 (ADR-040). When set, the directory persists it
    /// and clients SHOULD prefer it over `endpoint` for task CALLs.
    pub transport_endpoint: Option<String>,
    /// #331 Phase A.1 / ADR-041 — NAT-traversal observability fields surfaced
    /// to the directory in the register payload. Populated by
    /// [`IicpNode::apply_nat_profile`] when an operator runs detect_nat at
    /// startup, OR set manually if the operator already knows their topology.
    ///
    /// `transport_method` is one of `direct` / `upnp_mapped` / `stun_hole_punch`
    /// / `turn_relay` / `external_tunnel` / `unknown`.
    pub transport_method: Option<String>,
    /// One of `full_cone` / `restricted_cone` / `port_restricted` / `symmetric`
    /// / `unknown` (observability only).
    pub nat_type: Option<String>,
    /// Forward-compat slot for ADR-041 transport_candidates[] + relay_endpoint.
    pub transport_metadata: Option<serde_json::Value>,
    /// ADR-043 §9 — 8-category exposure_mode, computed by `qualify_service` and set
    /// in `apply_nat_profile`. Surfaced to the directory `nodes.exposure_mode` column (#344).
    pub exposure_mode: Option<String>,
    /// S.12 §2.1 CIP policy block surfaced to the directory register payload.
    /// When `None`, register() falls back to the module-level
    /// [`crate::cip_policy::get_cip_policy`] — operators can configure once
    /// and have it apply to all nodes that don't override.
    pub cip_policy: Option<std::sync::Arc<crate::cip_policy::CooperativeInferencePolicy>>,
    /// ADR-019 declarative pricing block. When `None`, the SDK does not
    /// advertise pricing and the directory defaults to a 1.0 multiplier.
    pub pricing: Option<crate::pricing::PricingConfig>,
    /// Operator-provisioned HMAC key for ADR-019 pricing signatures. When
    /// empty, the SDK captures the directory-issued key from the register
    /// response and uses it for subsequent signing.
    pub node_hmac_key: String,
    /// Phase 3+ availability windows (ADR-006). Local-time "HH:MM" windows that
    /// shape the effective capacity advertised to the directory and gated at
    /// serve time. Empty → always full capacity. See [`crate::availability`].
    pub availability_windows: Vec<crate::availability::Window>,
    /// ADR-010 task_id idempotency. `false` by default to preserve the pre-0.6
    /// contract (a task_id may be resubmitted). When `true`, a duplicate task_id
    /// within the 5-minute window is rejected with IICP-E010.
    pub enable_idempotency: bool,
    /// Phase 2 mesh (ADR-009/022). When `true`, serve() gossips peers and exposes
    /// POST /v1/peers. Default false.
    pub enable_mesh: bool,
    /// When `true`, serve() exposes POST /v1/relay to forward tasks to peers learned
    /// via gossip (ADR-022). Requires `enable_mesh`. Default false.
    pub relay_capable: bool,
    /// Port for the RelayAcceptServer (R1 relay-as-last-resort, #341).
    /// Workers behind CGNAT connect here outbound and send RELAY_BIND. Default 9485.
    pub relay_accept_port: u16,
    /// R2: when set, this node acts as a relay WORKER — connects outbound to the
    /// specified relay endpoint. Format: "host:port" (e.g. "relay.example.com:9485").
    pub relay_worker_endpoint: Option<String>,
}

impl NodeConfig {
    pub fn new(
        node_id: impl Into<String>,
        endpoint: impl Into<String>,
        intent: impl Into<String>,
    ) -> Self {
        Self {
            node_id: node_id.into(),
            endpoint: endpoint.into(),
            intent: intent.into(),
            model: None,
            region: None,
            capabilities: vec![],
            directory_url: DEFAULT_DIRECTORY.into(),
            timeout_ms: 5_000,
            max_concurrent: 4,
            tokens_per_min: 10_000,
            max_tokens: 8_192,
            transport_endpoint: None,
            transport_method: None,
            nat_type: None,
            transport_metadata: None,
            exposure_mode: None,
            cip_policy: None,
            pricing: None,
            node_hmac_key: String::new(),
            availability_windows: Vec::new(),
            enable_idempotency: false,
            enable_mesh: false,
            relay_capable: false,
            relay_accept_port: 9485,
            relay_worker_endpoint: None,
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct TaskRequest {
    pub task_id: String,
    pub intent: String,
    pub payload: Value,
    pub constraints: Option<Value>,
    pub auth: Option<Value>,
    pub nonce: Option<String>,
    /// Injected server-side from the W3C `traceparent` header — not from the JSON body.
    #[serde(skip_deserializing)]
    pub _trace: Option<Value>,
}

#[derive(Debug, Serialize)]
pub struct TaskResponse {
    pub task_id: String,
    pub status: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<Value>,
}

pub type TaskHandlerFn = Arc<
    dyn Fn(
            TaskRequest,
        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Value>> + Send>>
        + Send
        + Sync,
>;

struct AppState {
    handler: TaskHandlerFn,
    node_id: String,
    region: String,
    intent: String,
    model: String,
    active_jobs: Arc<AtomicUsize>,
    max_concurrent: usize,
    availability: Arc<crate::availability::AvailabilityEvaluator>,
    idempotency: Arc<crate::idempotency::IdempotencyGuard>,
    enable_idempotency: bool,
    peer_manager: Arc<crate::peer_manager::PeerManager>,
    http: reqwest::Client,
    nonce_cache: Arc<Mutex<HashMap<String, Instant>>>,
    /// #343 — shared pinhole state for /iicp/health surface.
    pinhole_uid: Arc<std::sync::RwLock<Option<u32>>>,
    pinhole_lease_seconds: Arc<std::sync::RwLock<u32>>,
    /// R1 relay-as-last-resort (#341): sessions from workers binding outbound.
    #[cfg(feature = "iicp-tcp")]
    relay_sessions: Arc<crate::relay_session::RelaySessionRegistry>,
}

// ── GET /iicp/health ─────────────────────────────────────────────────────────

async fn health_endpoint(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    let active = state.active_jobs.load(Ordering::Relaxed);
    let uid = state.pinhole_uid.read().ok().and_then(|g| *g);
    let lease = state
        .pinhole_lease_seconds
        .read()
        .map(|g| *g)
        .unwrap_or(3600);
    let pinhole_state = if let Some(uid) = uid {
        json!({ "active": true, "unique_id": uid, "lease_seconds": lease })
    } else {
        json!({ "active": false })
    };
    let eff_max = state
        .availability
        .effective_max_concurrent(state.max_concurrent);
    Json(json!({
        "status": "ok",
        "node_id": state.node_id,
        "region": state.region,
        "load": (active as f64 / state.max_concurrent.max(1) as f64),
        "active_jobs": active,
        "max_concurrent": state.max_concurrent,
        "effective_max_concurrent": eff_max,
        "available": active < eff_max,
        "model": state.model,
        "intent": state.intent,
        "pinhole_state": pinhole_state,
    }))
}

// ── GET /metrics ─────────────────────────────────────────────────────────────

async fn metrics_endpoint() -> Response {
    #[cfg(feature = "metrics")]
    {
        use prometheus::{Encoder, TextEncoder};
        let encoder = TextEncoder::new();
        let mf = prometheus::gather();
        let mut buf = Vec::new();
        if encoder.encode(&mf, &mut buf).is_ok() {
            return (
                StatusCode::OK,
                [(
                    axum::http::header::CONTENT_TYPE,
                    "text/plain; version=0.0.4",
                )],
                buf,
            )
                .into_response();
        }
    }
    (
        StatusCode::SERVICE_UNAVAILABLE,
        "metrics feature not enabled",
    )
        .into_response()
}

// ── POST /v1/peers (ADR-009 gossip exchange) ──────────────────────────────────

async fn peers_endpoint(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    body: axum::body::Bytes,
) -> Response {
    let sig = headers
        .get("x-iicp-signature")
        .and_then(|v| v.to_str().ok());
    if !state.peer_manager.verify_exchange(&body, sig) {
        return (
            StatusCode::UNAUTHORIZED,
            Json(json!({"error":{"code":"IICP-E012","message":"invalid_signature"}})),
        )
            .into_response();
    }
    if let Ok(parsed) = serde_json::from_slice::<Value>(&body) {
        if let Some(arr) = parsed.get("known_peers").and_then(Value::as_array) {
            let dicts: Vec<Value> = arr.iter().filter(|p| p.is_object()).cloned().collect();
            state.peer_manager.merge_peers(&dicts);
        }
    }
    let peers: Vec<Value> = state
        .peer_manager
        .get_peers()
        .iter()
        .map(|p| {
            json!({
                "node_id": p.node_id,
                "endpoint": p.endpoint,
                "region": p.region,
                "last_seen": p.last_seen,
            })
        })
        .collect();
    Json(json!({ "peers": peers })).into_response()
}

// ── POST /v1/relay (ADR-022 mesh relay) ───────────────────────────────────────

async fn relay_endpoint(
    State(state): State<Arc<AppState>>,
    Json(payload): Json<Value>,
) -> Response {
    let target_id = payload
        .get("target_node_id")
        .and_then(Value::as_str)
        .unwrap_or("");
    let task = payload.get("task");
    if target_id.is_empty() || task.is_none() {
        return (
            StatusCode::UNPROCESSABLE_ENTITY,
            Json(
                json!({"error":{"code":"IICP-E000","message":"target_node_id and task required"}}),
            ),
        )
            .into_response();
    }
    let task_val = task.expect("checked above").clone();

    // R1: check relay session registry first (CGNAT workers with no inbound endpoint)
    #[cfg(feature = "iicp-tcp")]
    if let Some(session) = state.relay_sessions.get(target_id) {
        match session.forward_task(&task_val, 120).await {
            Ok(result) => {
                let task_id = task_val
                    .get("task_id")
                    .and_then(Value::as_str)
                    .unwrap_or("");
                return Json(json!({
                    "task_id": task_id,
                    "status": "completed",
                    "result": result
                }))
                .into_response();
            }
            Err(e) => {
                return (
                    StatusCode::BAD_GATEWAY,
                    Json(json!({"error":{"code":"IICP-E031","message":format!("relay session forward failed: {e}")}})),
                )
                    .into_response();
            }
        }
    }

    // Fall back to HTTP forwarding for routable peers (ADR-022)
    let target = match state.peer_manager.relay_target(target_id) {
        Some(t) => t,
        None => {
            return (
                StatusCode::NOT_FOUND,
                Json(json!({"error":{"code":"IICP-E030","message":"target not in peer list and not a bound relay worker"}})),
            )
                .into_response();
        }
    };
    let url = format!("{}/v1/task", target.endpoint.trim_end_matches('/'));
    match state
        .http
        .post(&url)
        .timeout(Duration::from_secs(120))
        .json(&task_val)
        .send()
        .await
    {
        Ok(resp) => {
            let status = StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::OK);
            let bytes = resp.bytes().await.unwrap_or_default();
            (status, bytes).into_response()
        }
        Err(e) => (
            StatusCode::BAD_GATEWAY,
            Json(json!({"error":{"code":"IICP-E031","message":format!("relay failed: {e}")}})),
        )
            .into_response(),
    }
}

// ── POST /v1/task ─────────────────────────────────────────────────────────────

/// Try to claim a concurrency slot. On `true` the caller owns one increment of
/// `active_jobs` and MUST `fetch_sub` it on every exit path. realtime/interactive
/// wait briefly for a slot; other tiers fail fast so the proxy sees back-pressure
/// immediately (ADR-006; see [`crate::scheduler`]).
async fn admit(state: &AppState, qos: &str) -> bool {
    // Effective cap folds in availability windows (ADR-006): a reduced/closed
    // window lowers capacity below max_concurrent.
    let cap = state
        .availability
        .effective_max_concurrent(state.max_concurrent);
    let prev = state.active_jobs.fetch_add(1, Ordering::Relaxed);
    if prev < cap {
        return true;
    }
    state.active_jobs.fetch_sub(1, Ordering::Relaxed);
    if !crate::scheduler::is_queue_eligible(qos) {
        return false;
    }
    let deadline = Instant::now() + crate::scheduler::QUEUE_WAIT;
    while Instant::now() < deadline {
        tokio::time::sleep(Duration::from_millis(50)).await;
        let cap = state
            .availability
            .effective_max_concurrent(state.max_concurrent);
        let prev = state.active_jobs.fetch_add(1, Ordering::Relaxed);
        if prev < cap {
            return true;
        }
        state.active_jobs.fetch_sub(1, Ordering::Relaxed);
    }
    false
}

async fn task_endpoint(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    Json(mut req): Json<TaskRequest>,
) -> Response {
    // QoS-aware admission — IICP-E021
    let qos = req
        .constraints
        .as_ref()
        .and_then(|c| c.get("qos_class"))
        .and_then(|v| v.as_str())
        .unwrap_or("best_effort")
        .to_string();
    if !admit(&state, &qos).await {
        return (
            StatusCode::TOO_MANY_REQUESTS,
            [("Retry-After", "2"), ("Content-Type", "application/json")],
            Json(json!({
                "error": {
                    "code": "IICP-E021",
                    "message": "capacity_exceeded",
                    "qos_class": qos,
                    "retry_after_ms": 2000,
                }
            })),
        )
            .into_response();
    }

    // Nonce replay protection — IICP-E011
    if let Some(ref nonce) = req.nonce {
        let mut cache = state.nonce_cache.lock().await;
        cache.retain(|_, inserted_at| inserted_at.elapsed().as_secs() < NONCE_TTL_SECS);
        if cache.contains_key(nonce) {
            state.active_jobs.fetch_sub(1, Ordering::Relaxed);
            return (
                StatusCode::CONFLICT,
                Json(json!({
                    "error": { "code": "IICP-E011", "message": "replay_detected" }
                })),
            )
                .into_response();
        }
        cache.insert(nonce.clone(), Instant::now());
    }

    // Idempotency — duplicate task_id within the retry window (ADR-010). Opt-in
    // (NodeConfig.enable_idempotency) to preserve the pre-0.6 contract.
    if state.enable_idempotency && !state.idempotency.check_and_register(&req.task_id) {
        state.active_jobs.fetch_sub(1, Ordering::Relaxed);
        return (
            StatusCode::CONFLICT,
            Json(json!({
                "error": { "code": "IICP-E010", "message": "duplicate_task" }
            })),
        )
            .into_response();
    }

    // W3C traceparent propagation
    if let Some(tp) = headers.get("traceparent").and_then(|v| v.to_str().ok()) {
        req._trace = Some(json!({ "traceparent": tp }));
    }

    let task_id = req.task_id.clone();
    // ADR-014 TRACE-02 — iicp.task.execute span via `tracing` crate.
    // `tracing-opentelemetry` bridge propagates this to an OTLP collector when
    // OTEL_EXPORTER_OTLP_ENDPOINT is set and the operator configures the bridge
    // at startup (e.g. via opentelemetry-otlp + tracing-opentelemetry).
    let result = {
        let span = tracing::info_span!(
            "iicp.task.execute",
            "iicp.task_id" = %task_id,
            "iicp.intent" = %req.intent,
        );
        let _guard = span.enter();
        (state.handler)(req).await
    };
    state.active_jobs.fetch_sub(1, Ordering::Relaxed);

    match result {
        Ok(value) => Json(TaskResponse {
            task_id,
            status: "completed".into(),
            result: Some(value),
            error: None,
        })
        .into_response(),
        Err(e) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(TaskResponse {
                task_id,
                status: "error".into(),
                result: None,
                error: Some(json!({ "message": e.to_string() })),
            }),
        )
            .into_response(),
    }
}

// ── IicpNode ──────────────────────────────────────────────────────────────────

/// IICP provider node — handles registration, heartbeats, and task serving.
pub struct IicpNode {
    cfg: NodeConfig,
    http: Client,
    /// ADR-019 HMAC key used for signing pricing declarations. Initialized
    /// from `cfg.node_hmac_key`; populated from the directory's response on
    /// first register() so subsequent re-registrations sign with the
    /// directory-issued key.
    runtime_hmac_key: std::sync::RwLock<String>,
    /// BUG-5: token stashed by register() so deregister()/heartbeat don't need it re-passed.
    runtime_token: std::sync::RwLock<String>,
    /// #343 — UPnP IPv6 pinhole UID captured by `apply_nat_profile`, revoked
    /// on shutdown via [`Self::revoke_pinhole`]. Only read under the `nat`
    /// feature; allowed dead_code so non-nat builds compile cleanly.
    #[allow(dead_code)]
    pinhole_uid: std::sync::RwLock<Option<u32>>,
    #[allow(dead_code)]
    pinhole_lease_seconds: std::sync::RwLock<u32>,
}

impl IicpNode {
    pub fn new(cfg: NodeConfig) -> Self {
        let http = Client::builder()
            .timeout(Duration::from_millis(cfg.timeout_ms + 2_000))
            .use_rustls_tls()
            .build()
            .expect("failed to build HTTP client");
        let runtime_hmac_key = std::sync::RwLock::new(cfg.node_hmac_key.clone());
        Self {
            cfg,
            http,
            runtime_hmac_key,
            runtime_token: std::sync::RwLock::new(String::new()),
            pinhole_uid: std::sync::RwLock::new(None),
            pinhole_lease_seconds: std::sync::RwLock::new(3600),
        }
    }

    /// Current HMAC key in use for ADR-019 pricing signatures (empty if
    /// unregistered AND no operator-provisioned key).
    pub fn node_hmac_key(&self) -> String {
        self.runtime_hmac_key.read().expect("poisoned").clone()
    }

    /// Borrow this node's configuration. Useful for callers (e.g.
    /// [`crate::conformance::run_conformance_checks`]) that need to inspect
    /// `directory_url`, `endpoint`, or `node_id` without owning the config.
    pub fn cfg(&self) -> &NodeConfig {
        &self.cfg
    }

    /// Set the relay-worker endpoint after construction. Used by the CLI when a
    /// relay is auto-elected post-NAT-detection (tier ≥ 3): `serve()` reads
    /// `self.cfg.relay_worker_endpoint` to start the outbound relay session.
    pub fn set_relay_worker_endpoint(&mut self, endpoint: String) {
        self.cfg.relay_worker_endpoint = Some(endpoint);
    }

    /// Populate `endpoint`, `transport_endpoint`, and the NAT observability
    /// fields from a `NatProfile` produced by [`crate::nat_detection::detect_nat`].
    ///
    /// Operators typically call this right after `detect_nat()` and before
    /// `register()` so the directory receives the discovered public endpoint
    /// + transport_method/nat_type/transport_metadata in the same payload.
    ///
    /// Defensive: tier-4 (unreachable) profiles do NOT overwrite a manually-
    /// set endpoint, and `transport_method == "unreachable"` is filtered out
    /// before register.
    #[cfg(feature = "nat")]
    pub fn apply_nat_profile(&mut self, profile: &crate::nat_detection::NatProfile) {
        if profile.is_reachable() {
            if let Some(pub_ep) = &profile.public_endpoint {
                self.cfg.endpoint = pub_ep.clone();
            }
        }
        if let Some(tep) = &profile.transport_endpoint {
            self.cfg.transport_endpoint = Some(tep.clone());
        }
        let tm = match profile.transport_method {
            crate::nat_detection::TransportMethod::Direct => Some("direct"),
            crate::nat_detection::TransportMethod::UpnpMapped => Some("upnp_mapped"),
            crate::nat_detection::TransportMethod::StunHolePunch => Some("stun_hole_punch"),
            crate::nat_detection::TransportMethod::TurnRelay => Some("turn_relay"),
            crate::nat_detection::TransportMethod::ExternalTunnel => Some("external_tunnel"),
            crate::nat_detection::TransportMethod::Unreachable => None,
        };
        if let Some(name) = tm {
            self.cfg.transport_method = Some(name.into());
        }
        if self.cfg.nat_type.is_none() {
            self.cfg.nat_type = Some("unknown".into());
        }
        let tail: Vec<&str> = profile
            .detection_log
            .iter()
            .rev()
            .take(1)
            .map(|s| s.as_str())
            .collect();
        self.cfg.transport_metadata = Some(serde_json::json!({
            "tier": profile.tier,
            "detection_log_tail": tail,
        }));
        // ADR-043 §9 (#344) — derive the canonical 8-category exposure_mode and
        // advertise it so the directory can store nodes.exposure_mode for routing.
        self.cfg.exposure_mode = Some(
            crate::qualify::qualify_service(profile)
                .exposure_mode
                .to_string(),
        );
        // #343 — capture the IPv6 firewall pinhole UID and lease so we can renew and revoke.
        if let Some(v6) = &profile.ipv6 {
            if v6.pinhole_active {
                if let Some(uid) = v6.pinhole_unique_id {
                    if let Ok(mut slot) = self.pinhole_uid.write() {
                        *slot = Some(uid);
                    }
                }
                if let Some(lease) = v6.pinhole_lease_seconds {
                    if let Ok(mut slot) = self.pinhole_lease_seconds.write() {
                        *slot = lease;
                    }
                }
            }
        }
    }

    /// #343 — close the UPnP IPv6 firewall pinhole if one is tracked. Best-effort.
    #[cfg(feature = "nat")]
    pub async fn revoke_pinhole(&self) -> bool {
        let uid = match self.pinhole_uid.write() {
            Ok(mut slot) => slot.take(),
            Err(_) => None,
        };
        match uid {
            Some(uid) => crate::nat_detection::delete_ipv6_pinhole(uid).await,
            None => false,
        }
    }

    /// Tell the directory this node is going away.
    ///
    /// Mirrors `iicp_client.IicpNode.deregister` (Python iter-1471) and
    /// `IicpNode.deregister` (TS iter-1474). Best-effort: shutdown paths
    /// swallow failures so a flaky directory connection doesn't block exit.
    /// Deregister from the directory. `node_token` defaults to the token stashed by
    /// `register()` (BUG-5) when `None` — pass `Some(token)` to override.
    pub async fn deregister(&self, node_token: Option<&str>) -> Result<()> {
        let stashed = self.runtime_token.read().expect("poisoned").clone();
        let token = node_token.map(str::to_string).unwrap_or(stashed);
        if token.is_empty() {
            return Err(crate::errors::IicpError::Node(
                "deregister() requires a node_token (none stashed — call register() first)".into(),
            ));
        }
        let url = format!(
            "{}/v1/register",
            self.cfg.directory_url.trim_end_matches('/')
        );
        let resp = self
            .http
            .delete(&url)
            .bearer_auth(&token)
            .json(&serde_json::json!({"node_id": self.cfg.node_id}))
            .send()
            .await?;
        let status = resp.status();
        if !status.is_success() && status.as_u16() != 404 {
            return Err(crate::errors::IicpError::Node(format!(
                "Deregister failed: {status}"
            )));
        }
        Ok(())
    }

    /// Register with the directory and return the assigned `node_token`.
    ///
    /// Payload conforms to spec/iicp-dir.md §3.1 REGISTER plus the v0.7.0
    /// dual-endpoint extension (`transport_endpoint`). Pre-iter-1413
    /// builds sent a non-spec flat-`intent` shape that the production
    /// directory rejects with 422; fixed here.
    pub async fn register(&self) -> Result<String> {
        // Build the spec-compliant capability object. Legacy
        // `capabilities: Vec<String>` is folded into the models array.
        let mut models: Vec<String> = match &self.cfg.model {
            Some(m) => vec![m.clone()],
            None => Vec::new(),
        };
        for cap in &self.cfg.capabilities {
            if !models.contains(cap) {
                models.push(cap.clone());
            }
        }
        let region = self
            .cfg
            .region
            .clone()
            .unwrap_or_else(|| "eu-central".to_string());

        let mut payload = json!({
            "endpoint": self.cfg.endpoint,
            "region": region,
            "capabilities": [{
                "intent": self.cfg.intent,
                "models": models,
                "max_tokens": self.cfg.max_tokens,
            }],
            "limits": {
                "max_concurrent": self.cfg.max_concurrent,
                "tokens_per_min": self.cfg.tokens_per_min,
            },
        });
        if !self.cfg.node_id.is_empty() {
            payload["node_id"] = json!(self.cfg.node_id);
        }
        // spec v0.7.0 — native IICP binary endpoint
        if let Some(t) = &self.cfg.transport_endpoint {
            payload["transport_endpoint"] = json!(t);
        }
        // #331 / ADR-041 — NAT-traversal observability (set manually or via
        // apply_nat_profile after detect_nat)
        if let Some(m) = &self.cfg.transport_method {
            payload["transport_method"] = json!(m);
        }
        if let Some(n) = &self.cfg.nat_type {
            payload["nat_type"] = json!(n);
        }
        if let Some(md) = &self.cfg.transport_metadata {
            payload["transport_metadata"] = md.clone();
        }
        // ADR-043 §9 (#344) — 8-category network exposure classification
        if let Some(e) = &self.cfg.exposure_mode {
            payload["exposure_mode"] = json!(e);
        }

        // SDK self-identification — directory surfaces these on /v1/discover
        // so dashboards can render a language badge. Free-form so future
        // SDKs in other languages can self-tag without a directory change.
        payload["sdk_language"] = json!("rust");
        payload["sdk_version"] = json!(env!("CARGO_PKG_VERSION"));

        // S.12 §2.1 CIP-D1 policy block. Use the per-config policy if set,
        // otherwise fall back to the module-level cip_policy::get_cip_policy().
        let policy_arc = self
            .cfg
            .cip_policy
            .clone()
            .unwrap_or_else(crate::cip_policy::get_cip_policy);
        if let Some(block) = policy_arc.as_register_policy_block() {
            payload["policy"] = block;
        }

        // ADR-019 — declarative pricing block. Operator opt-in.
        if let Some(pricing) = &self.cfg.pricing {
            let hmac_key = self.runtime_hmac_key.read().expect("poisoned").clone();
            payload["pricing"] = crate::pricing::build_pricing_block(pricing, &hmac_key);
        }
        if !self.cfg.node_hmac_key.is_empty() {
            payload["node_hmac_key"] = json!(self.cfg.node_hmac_key);
        }

        let resp = self
            .http
            .post(format!(
                "{}/v1/register",
                self.cfg.directory_url.trim_end_matches('/')
            ))
            .json(&payload)
            .send()
            .await
            .map_err(|e| IicpError::Node(e.to_string()))?;

        if !resp.status().is_success() {
            return Err(IicpError::Node(format!(
                "register failed: {}",
                resp.status()
            )));
        }
        let data: Value = resp
            .json()
            .await
            .map_err(|e| IicpError::Node(e.to_string()))?;
        let token = data["node_token"]
            .as_str()
            .or_else(|| data["token"].as_str())
            .ok_or_else(|| IicpError::Node(format!("no node_token in response: {data}")))?;
        // BUG-5: stash the token so deregister()/heartbeat don't need it re-passed.
        *self.runtime_token.write().expect("poisoned") = token.to_string();
        // ADR-019: capture directory-issued HMAC key for subsequent signing.
        // Operator-provisioned key (cfg.node_hmac_key) wins — we only set the
        // runtime key from the response when the operator hasn't set one.
        if self.cfg.node_hmac_key.is_empty() {
            if let Some(dir_key) = data["node_hmac_key"].as_str() {
                if !dir_key.is_empty() {
                    let mut guard = self.runtime_hmac_key.write().expect("poisoned");
                    *guard = dir_key.to_string();
                }
            }
        }
        Ok(token.to_string())
    }

    /// Send a single heartbeat to the directory.
    pub async fn heartbeat(&self, node_token: &str) -> Result<()> {
        let resp = self
            .http
            // /v1/heartbeat — default directory_url already ends in /api;
            // the prior /api/v1/heartbeat path doubled the prefix and 404'd,
            // so last_seen never updated and nodes vanished from /v1/stats.
            .post(format!(
                "{}/v1/heartbeat",
                self.cfg.directory_url.trim_end_matches('/')
            ))
            // NodeTokenAuth middleware requires Bearer auth; the body
            // token is retained for back-compat with older directory builds.
            .bearer_auth(node_token)
            .json(&json!({
                "node_id": self.cfg.node_id,
                "node_token": node_token,
                "status": "available",
                // Live capacity after availability shaping (ADR-006).
                "max_concurrent": crate::availability::AvailabilityEvaluator::new(
                    self.cfg.availability_windows.clone(),
                )
                .effective_max_concurrent(self.cfg.max_concurrent),
            }))
            .send()
            .await
            .map_err(|e| IicpError::Node(e.to_string()))?;

        if !resp.status().is_success() {
            return Err(IicpError::Node(format!(
                "heartbeat failed: {}",
                resp.status()
            )));
        }
        Ok(())
    }

    /// Start the task server (blocks until cancelled).
    ///
    /// Serves `POST /v1/task`, `GET /iicp/health`, `GET /metrics`.
    /// Starts a background heartbeat loop when `node_token` is provided.
    pub async fn serve<F, Fut>(
        &self,
        handler: F,
        addr: &str,
        node_token: Option<String>,
    ) -> Result<()>
    where
        F: Fn(TaskRequest) -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = Result<Value>> + Send + 'static,
    {
        let handler: TaskHandlerFn = Arc::new(move |req| Box::pin(handler(req)));
        // Clone before handler is potentially moved into the relay worker closure (iicp-tcp only).
        #[cfg(feature = "iicp-tcp")]
        let handler_for_relay = Arc::clone(&handler);
        // Extract bind host before `addr` is shadowed by SocketAddr (iicp-tcp only).
        #[cfg(feature = "iicp-tcp")]
        let bind_host: String = addr.split(':').next().unwrap_or("0.0.0.0").to_string();
        let active_jobs = Arc::new(AtomicUsize::new(0));
        let nonce_cache = Arc::new(Mutex::new(HashMap::new()));
        // #343 — shared pinhole state: pass to AppState (health endpoint) and renewal task.
        let shared_pinhole_uid: Arc<std::sync::RwLock<Option<u32>>> = Arc::new(
            std::sync::RwLock::new(self.pinhole_uid.read().ok().and_then(|g| *g)),
        );
        let shared_pinhole_lease: Arc<std::sync::RwLock<u32>> = Arc::new(std::sync::RwLock::new(
            self.pinhole_lease_seconds
                .read()
                .map(|g| *g)
                .unwrap_or(3600),
        ));

        let state = Arc::new(AppState {
            handler,
            node_id: self.cfg.node_id.clone(),
            region: self.cfg.region.clone().unwrap_or_else(|| "unknown".into()),
            intent: self.cfg.intent.clone(),
            model: self.cfg.model.clone().unwrap_or_default(),
            active_jobs,
            max_concurrent: self.cfg.max_concurrent,
            availability: Arc::new(crate::availability::AvailabilityEvaluator::new(
                self.cfg.availability_windows.clone(),
            )),
            idempotency: Arc::new(crate::idempotency::IdempotencyGuard::default()),
            enable_idempotency: self.cfg.enable_idempotency,
            peer_manager: Arc::new(crate::peer_manager::PeerManager::with_opts(
                self.cfg.directory_url.clone(),
                self.cfg.node_hmac_key.clone(),
                crate::peer_manager::PeerManagerOpts {
                    relay_capable: self.cfg.relay_capable,
                    relay_accept_port: self.cfg.relay_accept_port,
                },
            )),
            http: self.http.clone(),
            nonce_cache,
            pinhole_uid: Arc::clone(&shared_pinhole_uid),
            pinhole_lease_seconds: Arc::clone(&shared_pinhole_lease),
            #[cfg(feature = "iicp-tcp")]
            relay_sessions: Arc::new(crate::relay_session::RelaySessionRegistry::new()),
        });

        // Capture the availability handle before `state` is moved into the router,
        // so the heartbeat loop below can report effective capacity.
        let hb_availability = Arc::clone(&state.availability);
        // Phase 2 mesh: bootstrap + gossip when enabled (before `state` is moved).
        if self.cfg.enable_mesh {
            let pm = Arc::clone(&state.peer_manager);
            let node_id = self.cfg.node_id.clone();
            let own_endpoint = self.cfg.endpoint.clone();
            tokio::spawn(async move {
                pm.start(&node_id, &own_endpoint).await;
                let interval = pm.gossip_interval();
                loop {
                    tokio::time::sleep(interval).await;
                    pm.gossip_round().await;
                }
            });
        }

        let mut app = Router::new()
            .route("/v1/task", post(task_endpoint))
            .route("/iicp/health", get(health_endpoint))
            .route("/metrics", get(metrics_endpoint));
        if self.cfg.enable_mesh {
            app = app.route("/v1/peers", post(peers_endpoint));
        }
        if self.cfg.relay_capable {
            app = app.route("/v1/relay", post(relay_endpoint));
        }
        // R1: capture relay_sessions Arc before state is moved into the router.
        #[cfg(feature = "iicp-tcp")]
        let relay_sessions_arc = Arc::clone(&state.relay_sessions);
        let app = app.with_state(state);

        let addr: SocketAddr = addr
            .parse()
            .map_err(|e| IicpError::Node(format!("invalid addr: {e}")))?;
        let listener = TcpListener::bind(addr)
            .await
            .map_err(|e| IicpError::Node(e.to_string()))?;

        tracing::info!("IICP node {} listening on {}", self.cfg.node_id, addr);

        if let Some(token) = node_token {
            let node_id = self.cfg.node_id.clone();
            let dir = self.cfg.directory_url.clone();
            let http = self.http.clone();
            let avail = Arc::clone(&hb_availability);
            let max_c = self.cfg.max_concurrent;
            tokio::spawn(async move {
                loop {
                    tokio::time::sleep(Duration::from_secs(HEARTBEAT_INTERVAL_SECS)).await;
                    if let Err(e) = http
                        // /v1/heartbeat — see heartbeat() above for the doubled-prefix
                        // history. Same fix applied here in the background loop.
                        .post(format!("{}/v1/heartbeat", dir.trim_end_matches('/')))
                        .bearer_auth(&token)
                        .json(&json!({
                            "node_id": &node_id,
                            "node_token": &token,
                            "status": "available",
                            // Live capacity after availability shaping (ADR-006).
                            "max_concurrent": avail.effective_max_concurrent(max_c),
                        }))
                        .send()
                        .await
                    {
                        tracing::warn!("heartbeat failed: {e}");
                    }
                }
            });
        }

        // #343 — pinhole renewal task: extends the UPnP IPv6 firewall pinhole at lease/2.
        #[cfg(feature = "nat")]
        {
            let uid_arc = Arc::clone(&shared_pinhole_uid);
            let lease_arc = Arc::clone(&shared_pinhole_lease);
            tokio::spawn(async move {
                loop {
                    let (_uid, lease) = {
                        let u = uid_arc.read().ok().and_then(|g| *g);
                        let l = lease_arc.read().map(|g| *g).unwrap_or(3600);
                        (u, l)
                    };
                    let delay = Duration::from_secs(u64::from((lease / 2).max(60)));
                    tokio::time::sleep(delay).await;
                    let uid = match uid_arc.read().ok().and_then(|g| *g) {
                        Some(u) => u,
                        None => return,
                    };
                    let ok = crate::nat_detection::renew_ipv6_pinhole(uid, lease).await;
                    if ok {
                        tracing::debug!("UPnP IPv6 pinhole uid={uid} renewed (lease={lease}s)");
                    } else {
                        tracing::warn!("UPnP IPv6 pinhole uid={uid} renewal failed — will retry");
                    }
                }
            });
        }

        // R1: start RelayAcceptServer when relay-capable (#341)
        #[cfg(feature = "iicp-tcp")]
        if self.cfg.relay_capable {
            let relay_reg = relay_sessions_arc;
            let relay_host_str = bind_host.clone();
            let relay_port = self.cfg.relay_accept_port;
            tokio::spawn(async move {
                let srv = Arc::new(crate::relay_session::RelayAcceptServer::new(
                    (*relay_reg).clone(),
                    relay_host_str,
                    relay_port,
                ));
                if let Err(e) = srv.serve().await {
                    tracing::warn!("Relay accept server error: {e}");
                }
            });
        }

        // R2: start relay worker client if relay_worker_endpoint is configured (#341)
        #[cfg(feature = "iicp-tcp")]
        if let Some(ref ep) = self.cfg.relay_worker_endpoint {
            let ep = ep.clone();
            let node_id = self.cfg.node_id.clone();
            let intent = self.cfg.intent.clone();
            let models = self.cfg.model.clone().map(|m| vec![m]).unwrap_or_default();
            let handler_fn: crate::relay_worker_client::RelayHandlerFn =
                Arc::new(move |task: Value| {
                    let h = Arc::clone(&handler_for_relay);
                    Box::pin(async move {
                        let req = crate::node::TaskRequest {
                            task_id: task
                                .get("task_id")
                                .and_then(|v| v.as_str())
                                .unwrap_or("")
                                .to_string(),
                            intent: task
                                .get("intent")
                                .and_then(|v| v.as_str())
                                .unwrap_or("")
                                .to_string(),
                            payload: task.get("payload").cloned().unwrap_or(Value::Null),
                            constraints: task.get("constraints").cloned(),
                            auth: task.get("auth").cloned(),
                            nonce: None,
                            _trace: None,
                        };
                        h(req)
                            .await
                            .unwrap_or_else(|e| json!({"error": e.to_string()}))
                    })
                });
            let (rhost, rport) = {
                if let Some(pos) = ep.rfind(':') {
                    let port = ep[pos + 1..].parse::<u16>().unwrap_or(9485);
                    (ep[..pos].to_string(), port)
                } else {
                    (ep.clone(), 9485u16)
                }
            };
            // on_bind: re-register with the relay's public endpoint so the node
            // appears ACTIVE in directory + stats (#358).
            let http_client = self.http.clone();
            let dir_url = self.cfg.directory_url.clone();
            let on_bind_cb: crate::relay_worker_client::OnBindFn = Arc::new(
                move |rh: String, rp: u16, _wid: String| {
                    let http = http_client.clone();
                    let dir = dir_url.clone();
                    Box::pin(async move {
                        // A full re-register would require the IicpNode reference here,
                        // which isn't available. For v0.7.0 we log the bind event.
                        // The node operator should use the cli bin which has the full
                        // context to re-register. Full wiring tracked in #341 R2.
                        tracing::info!(
                            "Relay worker bound to relay {}:{} — update directory registration to use relay endpoint",
                            rh, rp,
                        );
                        let _ = (http, dir); // suppress unused warnings
                    })
                },
            );
            tokio::spawn(async move {
                let rwc = Arc::new(
                    crate::relay_worker_client::RelayWorkerClient::new(
                        node_id, intent, rhost, rport, handler_fn, models,
                    )
                    .with_on_bind(on_bind_cb),
                );
                rwc.run().await;
            });
        }

        axum::serve(listener, app)
            .await
            .map_err(|e| IicpError::Node(e.to_string()))
    }
}