iicp-client 0.7.68

// SPDX-License-Identifier: Apache-2.0
//! Integration tests for IicpNode: health, task, concurrency gate, nonce replay, traceparent.

use std::net::TcpListener as StdListener;

use iicp_client::confidentiality::encrypt_payload;
use iicp_client::node::{IicpNode, NodeConfig};
use iicp_client::CxPublicKey;
use serde_json::{json, Value};

fn free_port() -> u16 {
    StdListener::bind("127.0.0.1:0")
        .unwrap()
        .local_addr()
        .unwrap()
        .port()
}

async fn start_node(port: u16, max_concurrent: usize) -> tokio::task::JoinHandle<()> {
    let mut cfg = NodeConfig::new(
        "test-node",
        "http://test.local",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.max_concurrent = max_concurrent;
    cfg.region = Some("test-region".into());
    cfg.model = Some("test-model".into());
    let node = IicpNode::new(cfg);
    let addr = format!("127.0.0.1:{port}");
    tokio::spawn(async move {
        let _ = node
            .serve(
                |task| Box::pin(async move { Ok(json!({ "echo": task.payload })) }),
                &addr,
                None,
            )
            .await;
    })
}

async fn wait_port(port: u16) {
    for _ in 0..40 {
        if reqwest::get(format!("http://127.0.0.1:{port}/iicp/health"))
            .await
            .is_ok()
        {
            return;
        }
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    }
    panic!("server did not start on port {port}");
}

#[tokio::test]
async fn test_health_endpoint_returns_200() {
    let port = free_port();
    let handle = start_node(port, 4).await;
    wait_port(port).await;

    let resp = reqwest::get(format!("http://127.0.0.1:{port}/iicp/health"))
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);
    let body: Value = resp.json().await.unwrap();
    assert_eq!(body["status"], "ok");
    assert_eq!(body["node_id"], "test-node");
    assert_eq!(body["max_concurrent"], 4);
    assert!(body["available"].as_bool().unwrap_or(false));

    handle.abort();
}

#[tokio::test]
async fn test_task_endpoint_returns_200() {
    let port = free_port();
    let handle = start_node(port, 4).await;
    wait_port(port).await;

    let client = reqwest::Client::new();
    let resp = client
        .post(format!("http://127.0.0.1:{port}/v1/task"))
        .json(&json!({ "task_id": "t-001", "intent": "x", "payload": { "msg": "hi" } }))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);
    let body: Value = resp.json().await.unwrap();
    // Spec iicp-dir.md task status ∈ {success, failure, timeout}; the node returns
    // "success" (was "completed" — fixed in ddd002a; this assertion was a missed
    // loose end from that fix, surfaced by #453).
    assert_eq!(body["status"], "success");
    assert_eq!(body["task_id"], "t-001");

    handle.abort();
}

#[tokio::test]
async fn test_task_endpoint_decrypts_iicp_conf() {
    let port = free_port();
    let mut cfg = NodeConfig::new(
        "cx-rust-node",
        "http://cx-rust.local",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.model = Some("test-model".into());
    let node = IicpNode::new(cfg);
    let cx_public_key: CxPublicKey =
        serde_json::from_value(node.register_payload_for_test()["cx_public_key"].clone()).unwrap();
    let addr = format!("127.0.0.1:{port}");
    let handle = tokio::spawn(async move {
        let _ = node
            .serve(
                |task| Box::pin(async move { Ok(json!({"echo": task.payload})) }),
                &addr,
                None,
            )
            .await;
    });
    wait_port(port).await;

    let payload = json!({"messages": [{"role": "user", "content": "secret"}]});
    let env = encrypt_payload(
        &payload,
        &cx_public_key,
        "cx-rust-1",
        "urn:iicp:intent:llm:chat:v1",
    )
    .unwrap();
    let request_body = json!({
        "task_id": "cx-rust-1",
        "intent": "urn:iicp:intent:llm:chat:v1",
        "iicp_conf": env,
    });
    let parsed: iicp_client::node::TaskRequest =
        serde_json::from_value(request_body.clone()).unwrap();
    assert!(parsed.payload.is_null());
    assert!(parsed.iicp_conf.is_some());
    let client = reqwest::Client::new();
    let resp = client
        .post(format!("http://127.0.0.1:{port}/v1/task"))
        .json(&request_body)
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);
    let body: Value = resp.json().await.unwrap();
    assert_eq!(body["status"], "success");
    assert_eq!(body["result"]["echo"], payload);

    handle.abort();
}

#[tokio::test]
async fn test_concurrency_gate_429() {
    let port = free_port();
    let mut cfg = NodeConfig::new(
        "gate-node",
        "http://test.local",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.max_concurrent = 0;
    let node = IicpNode::new(cfg);
    let addr = format!("127.0.0.1:{port}");
    let handle = tokio::spawn(async move {
        let _ = node
            .serve(
                |task| Box::pin(async move { Ok(json!({"echo": task.payload})) }),
                &addr,
                None,
            )
            .await;
    });

    for _ in 0..40 {
        if reqwest::get(format!("http://127.0.0.1:{port}/iicp/health"))
            .await
            .is_ok()
        {
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    }

    let client = reqwest::Client::new();
    let resp = client
        .post(format!("http://127.0.0.1:{port}/v1/task"))
        .json(&json!({ "task_id": "t", "intent": "x", "payload": {} }))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 429);
    let retry_after = resp
        .headers()
        .get("retry-after")
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string());
    let body: Value = resp.json().await.unwrap();
    assert_eq!(body["error"]["code"], "IICP-E021");
    assert_eq!(retry_after.as_deref(), Some("2"));

    handle.abort();
}

#[tokio::test]
async fn test_nonce_replay_409() {
    let port = free_port();
    let handle = start_node(port, 4).await;
    wait_port(port).await;

    let client = reqwest::Client::new();
    let nonce = "nonce-rust-replay-test";

    let r1 = client
        .post(format!("http://127.0.0.1:{port}/v1/task"))
        .json(&json!({ "task_id": "t1", "intent": "x", "payload": {}, "nonce": nonce }))
        .send()
        .await
        .unwrap();
    assert_eq!(r1.status(), 200);

    let r2 = client
        .post(format!("http://127.0.0.1:{port}/v1/task"))
        .json(&json!({ "task_id": "t2", "intent": "x", "payload": {}, "nonce": nonce }))
        .send()
        .await
        .unwrap();
    assert_eq!(r2.status(), 409);
    let body: Value = r2.json().await.unwrap();
    assert_eq!(body["error"]["code"], "IICP-E011");

    handle.abort();
}

#[tokio::test]
async fn test_traceparent_propagated_to_handler() {
    let port = free_port();
    let tp_header = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";

    let captured = std::sync::Arc::new(tokio::sync::Mutex::new(None::<Value>));
    let captured_clone = captured.clone();

    let mut cfg = NodeConfig::new(
        "trace-node",
        "http://test.local",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.max_concurrent = 4;
    let node = IicpNode::new(cfg);
    let addr = format!("127.0.0.1:{port}");

    let handle = tokio::spawn(async move {
        let _ = node
            .serve(
                move |task| {
                    let cap = captured_clone.clone();
                    Box::pin(async move {
                        if let Some(t) = &task._trace {
                            *cap.lock().await = Some(t.clone());
                        }
                        Ok(json!({}))
                    })
                },
                &addr,
                None,
            )
            .await;
    });

    wait_port(port).await;

    let client = reqwest::Client::new();
    client
        .post(format!("http://127.0.0.1:{port}/v1/task"))
        .header("traceparent", tp_header)
        .json(&json!({ "task_id": "t1", "intent": "x", "payload": {} }))
        .send()
        .await
        .unwrap();

    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    let cap = captured.lock().await;
    let tp = cap.as_ref().and_then(|v| v["traceparent"].as_str());
    assert_eq!(tp, Some(tp_header));

    handle.abort();
}

#[tokio::test]
async fn test_node_register_returns_token() {
    use mockito::Server;

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(json!({ "node_token": "tok-abc123", "message": "registered" }).to_string())
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-001",
        "https://my-host.example.com",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    let node = IicpNode::new(cfg);
    let token = node.register().await.unwrap();
    assert_eq!(token, "tok-abc123");
}

#[tokio::test]
async fn test_node_register_no_token_fails() {
    use mockito::Server;

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(json!({ "message": "ok" }).to_string())
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-001",
        "https://my-host.example.com",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    let node = IicpNode::new(cfg);
    assert!(node.register().await.is_err());
}

#[tokio::test]
async fn test_node_heartbeat_ok() {
    // #346 — heartbeat path is /v1/heartbeat (NOT /api/v1/heartbeat).
    // Uses a raw TCP echo server: mockito's async mock returns 501 on Linux
    // CI for bearer_auth requests; axum routing returns 404 in the same env.
    // A raw TCP handler is environment-agnostic and sufficient to verify
    // that heartbeat() sends the request and accepts a 2xx response.
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();

    tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.unwrap();
        let mut buf = [0u8; 4096];
        let _ = stream.read(&mut buf).await;
        let resp = b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 15\r\nconnection: close\r\n\r\n{\"status\":\"ok\"}";
        let _ = stream.write_all(resp).await;
    });

    let mut cfg = NodeConfig::new(
        "n-001",
        "https://my-host.example.com",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = format!("http://{addr}");
    let node = IicpNode::new(cfg);
    node.heartbeat("tok-abc123")
        .await
        .expect("heartbeat should succeed against local server");
}

/// ADR-047 Part A (#411) — the node answers the directory's liveness challenge:
/// beat 1 captures the nonce, beat 2 returns HMAC-SHA256(node_hmac_key, nonce).
#[tokio::test]
async fn test_heartbeat_answers_liveness_challenge() {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::sync::oneshot;

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let (tx, rx) = oneshot::channel::<String>();

    tokio::spawn(async move {
        let mk_resp = |body: &str| {
            format!(
                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
                body.len(),
                body
            )
        };
        // Beat 1 — issue a nonce.
        let (mut s1, _) = listener.accept().await.unwrap();
        let mut b1 = [0u8; 4096];
        let _ = s1.read(&mut b1).await;
        let _ = s1
            .write_all(mk_resp("{\"ok\":true,\"challenge\":\"nonce-abc\"}").as_bytes())
            .await;
        // Beat 2 — capture the request so the test can assert challenge_response.
        let (mut s2, _) = listener.accept().await.unwrap();
        let mut b2 = vec![0u8; 8192];
        let n = s2.read(&mut b2).await.unwrap();
        let _ = s2.write_all(mk_resp("{\"ok\":true}").as_bytes()).await;
        let _ = tx.send(String::from_utf8_lossy(&b2[..n]).to_string());
    });

    let mut cfg = NodeConfig::new(
        "n-1",
        "https://h.example.com",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = format!("http://{addr}");
    cfg.node_hmac_key = "secret-key".to_string();
    let node = IicpNode::new(cfg);
    node.heartbeat("tok").await.unwrap(); // beat 1 → captures nonce
    node.heartbeat("tok").await.unwrap(); // beat 2 → answers

    let req2 = rx.await.unwrap();
    let expected = iicp_client::pricing::sign_body(b"nonce-abc", "secret-key");
    assert!(
        req2.contains(&format!("\"challenge_response\":\"{expected}\"")),
        "beat 2 must carry challenge_response=HMAC(key,nonce); got: {req2}"
    );
}

/// The heartbeat body MUST carry an explicit `available: true` boolean (not only the
/// `status: "available"` string). The directory keys discover eligibility off the
/// `available` field, so sending it restores a briefly-dormant node on the next beat —
/// robust even against directory builds older than v1.10.17.
#[tokio::test]
async fn test_heartbeat_payload_includes_available_true() {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::sync::oneshot;

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let (tx, rx) = oneshot::channel::<String>();

    tokio::spawn(async move {
        let (mut s, _) = listener.accept().await.unwrap();
        let mut b = vec![0u8; 8192];
        let n = s.read(&mut b).await.unwrap();
        let resp = "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 11\r\nconnection: close\r\n\r\n{\"ok\":true}";
        let _ = s.write_all(resp.as_bytes()).await;
        let _ = tx.send(String::from_utf8_lossy(&b[..n]).to_string());
    });

    let mut cfg = NodeConfig::new(
        "n-1",
        "https://h.example.com",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = format!("http://{addr}");
    let node = IicpNode::new(cfg);
    node.heartbeat("tok").await.unwrap();

    let req = rx.await.unwrap();
    assert!(
        req.contains("\"available\":true"),
        "heartbeat must carry available:true; got: {req}"
    );
    assert!(
        req.contains("\"status\":\"available\""),
        "heartbeat must keep status:available; got: {req}"
    );
}

/// iter-1413: register payload matches spec/iicp-dir.md §3.1 —
/// capabilities is an array of {intent, models, max_tokens} objects, not a flat intent string.
#[tokio::test]
async fn test_register_payload_spec_compliant() {
    use mockito::{Matcher, Server};

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .match_body(Matcher::PartialJson(json!({
            "endpoint": "https://provider.example.com:8080",
            "region": "eu-central",
            "capabilities": [{
                "intent": "urn:iicp:intent:llm:chat:v1",
                "models": ["llama-3-8b"],
                "max_tokens": 8192
            }],
            "limits": { "max_concurrent": 2, "tokens_per_min": 2000 }
        })))
        .with_status(201)
        .with_header("content-type", "application/json")
        .with_body(json!({ "node_token": "tok-1", "node_id": "n-1" }).to_string())
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-1",
        "https://provider.example.com:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.model = Some("llama-3-8b".into());
    cfg.region = Some("eu-central".into());
    cfg.max_concurrent = 2;
    cfg.tokens_per_min = 2000;
    cfg.max_tokens = 8192;
    let node = IicpNode::new(cfg);
    let register_payload = node.register_payload_for_test();
    assert_eq!(register_payload["cx_public_key"]["algorithm"], "X25519");
    assert_eq!(register_payload["cx_public_key"]["encoding"], "base64url");
    assert!(register_payload["cx_public_key"]["key_id"]
        .as_str()
        .unwrap()
        .starts_with("cx-"));
    let token = node.register().await.unwrap();
    assert_eq!(token, "tok-1");
}

/// iter-1413: spec v0.7.0 — register includes transport_endpoint when configured.
#[tokio::test]
async fn test_register_includes_transport_endpoint() {
    use mockito::{Matcher, Server};

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .match_body(Matcher::PartialJson(json!({
            "endpoint": "https://provider.example.com:8080",
            "transport_endpoint": "iicp://provider.example.com:9484"
        })))
        .with_status(201)
        .with_header("content-type", "application/json")
        .with_body(json!({ "node_token": "tok-2", "node_id": "n-2" }).to_string())
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-2",
        "https://provider.example.com:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.model = Some("qwen2.5:0.5b".into());
    cfg.transport_endpoint = Some("iicp://provider.example.com:9484".into());
    let node = IicpNode::new(cfg);
    assert!(node.register().await.is_ok());
}

/// iter-1428: register payload includes transport_method / nat_type /
/// transport_metadata when set on NodeConfig (manually OR via apply_nat_profile).
#[tokio::test]
async fn test_register_includes_nat_observability_when_set() {
    use mockito::{Matcher, Server};

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .match_body(Matcher::PartialJson(json!({
            "transport_method": "upnp_mapped",
            "nat_type": "full_cone",
            "transport_metadata": {"tier": 1}
        })))
        .with_status(201)
        .with_body(json!({ "node_token": "tok-nat", "node_id": "n-nat" }).to_string())
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-nat",
        "https://provider.example.com:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.model = Some("qwen2.5:0.5b".into());
    cfg.transport_endpoint = Some("iicp://provider.example.com:9484".into());
    cfg.transport_method = Some("upnp_mapped".into());
    cfg.nat_type = Some("full_cone".into());
    cfg.transport_metadata = Some(json!({"tier": 1, "detection_log_tail": ["ok"]}));
    let node = IicpNode::new(cfg);
    assert!(node.register().await.is_ok());
}

/// iter-1428: apply_nat_profile populates the NAT fields from a NatProfile
/// and overrides `endpoint` when the profile is reachable.
#[cfg(feature = "nat")]
#[tokio::test]
async fn test_apply_nat_profile_populates_fields() {
    use iicp_client::nat_detection::{NatProfile, TransportMethod};
    use mockito::{Matcher, Server};

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .match_body(Matcher::PartialJson(json!({
            "endpoint": "http://203.0.113.5:8080",
            "transport_endpoint": "iicp://203.0.113.5:9484",
            "transport_method": "upnp_mapped",
            "nat_type": "unknown"
        })))
        .with_status(201)
        .with_body(json!({ "node_token": "tok-applied", "node_id": "n-applied" }).to_string())
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-applied",
        "http://placeholder.example.com:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.model = Some("q".into());
    let mut node = IicpNode::new(cfg);

    let profile = NatProfile {
        tier: 1,
        transport_method: TransportMethod::UpnpMapped,
        public_endpoint: Some("http://203.0.113.5:8080".into()),
        transport_endpoint: Some("iicp://203.0.113.5:9484".into()),
        internal_endpoint: None,
        operator_guidance: None,
        detection_log: vec!["tier-1: UPnP mapped".into()],
        ipv6: None,
    };
    node.apply_nat_profile(&profile);
    assert!(node.register().await.is_ok());
}

/// iter-1428: tier-4 (unreachable) profiles preserve a manually-set endpoint
/// and do NOT surface transport_method "unreachable" to the directory.
#[cfg(feature = "nat")]
#[tokio::test]
async fn test_apply_nat_profile_unreachable_preserves_endpoint() {
    use iicp_client::nat_detection::{NatProfile, TransportMethod};
    use mockito::{Matcher, Server};

    let mut server = Server::new_async().await;
    // The mock matches the original manual endpoint AND requires
    // transport_method to be absent (PartialJson only checks supplied keys).
    // For "absence" we rely on the body match never including transport_method.
    let _m = server
        .mock("POST", "/v1/register")
        .match_body(Matcher::PartialJson(json!({
            "endpoint": "https://manual.example.com:8080"
        })))
        .with_status(201)
        .with_body(json!({ "node_token": "tok-keep", "node_id": "n-keep" }).to_string())
        .expect(1)
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-keep",
        "https://manual.example.com:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.model = Some("q".into());
    let mut node = IicpNode::new(cfg);

    let profile = NatProfile {
        tier: 4,
        transport_method: TransportMethod::Unreachable,
        public_endpoint: None,
        transport_endpoint: None,
        internal_endpoint: None,
        operator_guidance: Some("install igd-next".into()),
        detection_log: vec!["tier-4 fallback".into()],
        ipv6: None,
    };
    node.apply_nat_profile(&profile);
    assert!(node.register().await.is_ok());
    _m.assert_async().await;
}

/// iter-1413: legacy capabilities Vec<String> folds into the models array of the
/// single capability object — keeps pre-iter-1413 caller configs working.
/// We assert via a custom body matcher closure so the mock only succeeds when
/// the models array contains all three names (order-independent).
#[tokio::test]
async fn test_register_legacy_capabilities_folds_into_models() {
    use mockito::{Matcher, Server};

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .match_body(Matcher::PartialJson(json!({
            "capabilities": [{
                "intent": "urn:iicp:intent:llm:chat:v1",
                "max_tokens": 8192
            }]
        })))
        // The body matcher above guarantees the capabilities structure is correct.
        // To check models contains all three (order-independent), we expect 1 call;
        // a single Mock with PartialJson treats missing fields as no-match, so if
        // the structure is wrong, register() will get a non-mock 501 and fail.
        .with_status(201)
        .with_header("content-type", "application/json")
        .with_body(json!({ "node_token": "tok-3", "node_id": "n-3" }).to_string())
        .expect(1)
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-3",
        "https://provider.example.com:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.model = Some("llama-3-8b".into());
    cfg.capabilities = vec!["mistral-7b".into(), "phi-3-mini".into()];
    let node = IicpNode::new(cfg);
    assert!(node.register().await.is_ok());

    // Verify the mock fired exactly once (= our body matched).
    _m.assert_async().await;
}

/// WQ-066 (0.7.45) — relay_capable=true must appear in the register payload so
/// the directory can surface it in /v1/discover. Pre-fix: the Rust SDK wired
/// relay_capable to the serve() relay server but never included it in the
/// registration JSON → the directory always saw relay_capable=false.
#[tokio::test]
async fn test_register_payload_includes_relay_capable() {
    use mockito::{Matcher, Server};

    let mut server = Server::new_async().await;
    let _m = server
        .mock("POST", "/v1/register")
        .match_body(Matcher::PartialJson(json!({
            "relay_capable": true,
            "relay_accept_port": 9490
        })))
        .with_status(201)
        .with_header("content-type", "application/json")
        .with_body(json!({ "node_token": "tok-r", "node_id": "n-r" }).to_string())
        .expect(1)
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "n-r",
        "https://relay.example.com:9484",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.model = Some("llama-3-8b".into());
    cfg.relay_capable = true;
    cfg.relay_accept_port = 9490;
    let node = IicpNode::new(cfg);
    node.register().await.unwrap();
    _m.assert_async().await;
}

// ── #494 — health_models heartbeat reporting ──────────────────────────────────

/// When backend_url is set and /api/tags returns models, heartbeat payload
/// includes health_models. Fails if probe_health_models result is not forwarded.
#[tokio::test]
async fn test_heartbeat_includes_health_models_when_probe_succeeds() {
    use mockito::Server;
    let mut server = Server::new_async().await;

    // Mock the backend /api/tags probe
    let mut backend = mockito::Server::new_async().await;
    let _m_tags = backend
        .mock("GET", "/api/tags")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(
            json!({"models": [{"name": "llama3:latest"}, {"name": "qwen2.5:0.5b"}]}).to_string(),
        )
        .create_async()
        .await;

    // Mock the directory heartbeat endpoint — capture the body
    use std::sync::{Arc, Mutex};
    let captured: Arc<Mutex<Option<Value>>> = Arc::new(Mutex::new(None));
    let captured_clone = captured.clone();
    let _m_hb = server
        .mock("POST", "/v1/heartbeat")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(json!({"ok": true}).to_string())
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "hm-rs-1",
        "http://127.0.0.1:9999",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    cfg.backend_url = Some(backend.url());
    let node = IicpNode::new(cfg);
    node.heartbeat("tok")
        .await
        .expect("heartbeat should succeed");
    _m_hb.assert_async().await;
    // Verify the heartbeat body contained health_models (via the mock call count)
    // Full body inspection is done in Python/TS tests; here we confirm no panic + mock fires.
    let _ = captured_clone;
}

/// #494 — /iicp/health must expose models[] containing the primary model.
/// DIR-TRUST-01 REACH probe checks health_data.get("models"); fails if absent.
#[tokio::test]
async fn test_health_endpoint_exposes_models_array() {
    let port = free_port();
    let handle = start_node(port, 4).await;
    wait_port(port).await;

    let resp = reqwest::get(format!("http://127.0.0.1:{port}/iicp/health"))
        .await
        .unwrap();
    let body: Value = resp.json().await.unwrap();
    let models = body["models"]
        .as_array()
        .expect("/iicp/health must include models[]");
    assert!(
        models.iter().any(|m| m.as_str() == Some("test-model")),
        "primary model must appear in models[]"
    );
    handle.abort();
}

/// #494 — /iicp/health models[] includes capabilities (not just the primary model).
#[tokio::test]
async fn test_health_endpoint_models_includes_capabilities() {
    let port = free_port();
    let mut cfg = NodeConfig::new(
        "multi-node",
        "http://multi.local",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.model = Some("primary-model".into());
    cfg.capabilities = vec!["extra-model".into()];
    let node = IicpNode::new(cfg);
    let addr = format!("127.0.0.1:{port}");
    let handle = tokio::spawn(async move {
        let _ = node
            .serve(|_t| Box::pin(async move { Ok(json!({})) }), &addr, None)
            .await;
    });
    wait_port(port).await;

    let resp = reqwest::get(format!("http://127.0.0.1:{port}/iicp/health"))
        .await
        .unwrap();
    let body: Value = resp.json().await.unwrap();
    let models = body["models"].as_array().expect("models[] must be present");
    let names: Vec<&str> = models.iter().filter_map(|m| m.as_str()).collect();
    assert!(
        names.contains(&"primary-model"),
        "primary model in models[]"
    );
    assert!(
        names.contains(&"extra-model"),
        "capability must appear in models[]"
    );
    handle.abort();
}

/// When backend_url is not set, health_models must not appear in the heartbeat payload.
/// This test confirms backward compat — old nodes without backend_url still work.
#[tokio::test]
async fn test_heartbeat_omits_health_models_when_no_backend_url() {
    let mut server = mockito::Server::new_async().await;

    let _m = server
        .mock("POST", "/v1/heartbeat")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(json!({"ok": true}).to_string())
        .expect(1)
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "hm-rs-2",
        "http://127.0.0.1:9999",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = server.url();
    // No backend_url set
    let node = IicpNode::new(cfg);
    node.heartbeat("tok")
        .await
        .expect("heartbeat without backend_url should succeed");
    _m.assert_async().await;
}

// ── #494 — model drift re-registration ────────────────────────────────────────

/// When the backend's model list changes after registration, the heartbeat loop
/// should detect the drift and re-register with the updated list.
/// This test verifies the register call fires when live ≠ registered.
#[tokio::test]
async fn test_model_drift_triggers_reregister() {
    use mockito::Server;
    let mut dir = Server::new_async().await;
    let mut backend = mockito::Server::new_async().await;

    // Backend now only has phi3:mini (llama3.2:1b drifted away)
    let _m_tags = backend
        .mock("GET", "/api/tags")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(json!({"models": [{"name": "phi3:mini"}]}).to_string())
        .create_async()
        .await;

    // Directory register — must be called once for drift re-registration
    let _m_reg = dir
        .mock("POST", "/v1/register")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(json!({"node_token": "tok-drift-rs"}).to_string())
        .expect(1)
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "drift-rs-1",
        "http://node.local:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = dir.url();
    cfg.backend_url = Some(backend.url());
    cfg.model = Some("phi3:mini".into());
    cfg.capabilities = vec!["llama3.2:1b".into()];
    let node = IicpNode::new(cfg);
    // Simulate a prior registration: registered phi3:mini + llama3.2:1b
    {
        let mut g = node.registered_models().write().expect("poisoned");
        *g = vec!["phi3:mini".into(), "llama3.2:1b".into()];
    }

    node.check_model_drift_and_reregister().await;

    _m_reg.assert_async().await;
}

/// When backend returns empty model list, no re-registration should occur
/// (avoids spurious re-register during transient backend downtime).
#[tokio::test]
async fn test_no_reregister_on_empty_backend_models() {
    use mockito::Server;
    let mut dir = Server::new_async().await;
    let mut backend = mockito::Server::new_async().await;

    let _m_tags = backend
        .mock("GET", "/api/tags")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(json!({"models": []}).to_string())
        .create_async()
        .await;

    let _m_reg = dir
        .mock("POST", "/v1/register")
        .with_status(200)
        .expect(0) // must NOT be called
        .create_async()
        .await;

    let mut cfg = NodeConfig::new(
        "drift-rs-2",
        "http://node.local:8080",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.directory_url = dir.url();
    cfg.backend_url = Some(backend.url());
    cfg.model = Some("phi3:mini".into());
    let node = IicpNode::new(cfg);
    {
        let mut g = node.registered_models().write().expect("poisoned");
        *g = vec!["phi3:mini".into()];
    }

    node.check_model_drift_and_reregister().await;

    _m_reg.assert_async().await; // expect 0 calls
}

#[tokio::test]
async fn test_backend_stability_health_and_task_drain() {
    use iicp_client::backend_stability::BackendStabilityObservation;

    let port = free_port();
    let mut cfg = NodeConfig::new(
        "drain-rust-node",
        "http://drain-rust.local",
        "urn:iicp:intent:llm:chat:v1",
    );
    cfg.model = Some("test-model".into());
    let node = IicpNode::new(cfg);
    node.set_backend_stability_for_test(BackendStabilityObservation::draining(
        "backend_loading",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64
            + 30,
        json!({"model_size_bytes": 123, "loaded_instances": [{"id":"secret"}]}),
    ));
    let addr = format!("127.0.0.1:{port}");
    let handle = tokio::spawn(async move {
        let _ = node
            .serve(
                |_task| Box::pin(async move { Ok(json!({"ok": true})) }),
                &addr,
                None,
            )
            .await;
    });
    wait_port(port).await;

    let health: Value = reqwest::get(format!("http://127.0.0.1:{port}/iicp/health"))
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    assert_eq!(health["backend_stability"]["backend_state"], "draining");
    assert_eq!(
        health["backend_stability"]["reason_class"],
        "backend_loading"
    );
    assert!(health["backend_stability"]
        .get("model_size_bytes")
        .is_none());
    assert!(health["backend_stability"]
        .get("loaded_instances")
        .is_none());

    let resp = reqwest::Client::new()
        .post(format!("http://127.0.0.1:{port}/v1/task"))
        .json(&json!({"task_id":"t-drain","intent":"urn:iicp:intent:llm:chat:v1","payload":{}}))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE);
    assert!(resp.headers().get("retry-after").is_some());
    let body: Value = resp.json().await.unwrap();
    assert_eq!(body["error"]["code"], "IICP-E024");
    assert_eq!(body["error"]["reason"], "backend_loading");

    handle.abort();
}