Aether/apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs

use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::management_tokens::{
    InMemoryManagementTokenRepository, ManagementTokenListQuery, ManagementTokenReadRepository,
};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::{
    InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation, StoredProxyNodeEvent,
};
use aether_data::repository::users::InMemoryUserReadRepository;
use axum::body::Body;
use axum::extract::ws::Message;
use axum::routing::any;
use axum::{extract::Request, Router};
use base64::Engine as _;
use http::StatusCode;
use serde_json::json;
use tokio::sync::watch;

use super::super::{
    authenticated_tunnel_control_plane_request, build_router_with_state, hash_management_token,
    sample_endpoint, sample_key, sample_management_token, sample_provider, sample_proxy_node,
    start_server, with_tunnel_control_plane_key, AppState, TUNNEL_CONTROL_PLANE_TEST_GENERATION,
    TUNNEL_CONTROL_PLANE_TEST_PSK,
};
use crate::constants::{
    GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
    TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::data::GatewayDataState;
use crate::maintenance::{
    record_proxy_upgrade_traffic_success, skip_proxy_upgrade_rollout_node,
    start_proxy_upgrade_rollout,
};
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};

async fn recv_tunnel_test_frame(
    proxy_rx: &mut aether_runtime::BoundedQueueReceiver<Message>,
    description: &str,
) -> Message {
    tokio::time::timeout(Duration::from_secs(5), proxy_rx.recv())
        .await
        .unwrap_or_else(|_| panic!("timed out waiting for {description}"))
        .unwrap_or_else(|| panic!("proxy channel closed before {description}"))
}

#[tokio::test]
async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal() {
    let upstream_hits = Arc::new(Mutex::new(0usize));
    let upstream_hits_clone = Arc::clone(&upstream_hits);
    let upstream = Router::new().route(
        "/api/admin/proxy-nodes",
        any(move |_request: Request| {
            let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
            async move {
                *upstream_hits_inner.lock().expect("mutex should lock") += 1;
                (StatusCode::OK, Body::from("unexpected upstream hit"))
            }
        }),
    );

    let mut manual_node = sample_proxy_node("proxy-node-manual");
    manual_node.name = "alpha-manual".to_string();
    manual_node.status = "online".to_string();
    manual_node.is_manual = true;
    manual_node.tunnel_mode = false;
    manual_node.tunnel_connected = false;
    manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
    manual_node.proxy_username = Some("alice".to_string());
    manual_node.proxy_password = Some("supersecret".to_string());
    manual_node.last_heartbeat_at_unix_secs = None;
    manual_node.tunnel_connected_at_unix_secs = None;

    let mut tunnel_node = sample_proxy_node("proxy-node-tunnel");
    tunnel_node.name = "zeta-tunnel".to_string();
    tunnel_node.status = "offline".to_string();

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
        tunnel_node,
        manual_node,
    ]));

    let (upstream_url, upstream_handle) = start_server(upstream).await;
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
                proxy_node_repository,
            )),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?status=online&skip=0&limit=10"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["total"], 1);
    assert_eq!(payload["skip"], 0);
    assert_eq!(payload["limit"], 10);
    assert!(payload["rollout"].is_null());

    let items = payload["items"].as_array().expect("items should be array");
    assert_eq!(items.len(), 1);
    assert_eq!(items[0]["id"], "proxy-node-manual");
    assert_eq!(items[0]["name"], "alpha-manual");
    assert_eq!(items[0]["status"], "online");
    assert_eq!(items[0]["is_manual"], true);
    assert_eq!(items[0]["proxy_url"], "http://proxy.example:8080");
    assert_eq!(items[0]["proxy_username"], "alice");
    assert_eq!(items[0]["has_proxy_password"], true);
    assert!(items[0].get("proxy_password").is_none());
    assert!(items[0]["created_at"].is_string());
    assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

    gateway_handle.abort();
    upstream_handle.abort();
}

#[tokio::test]
async fn gateway_does_not_return_manual_proxy_password_in_node_detail() {
    let mut manual_node = sample_proxy_node("proxy-node-manual");
    manual_node.name = "alpha-manual".to_string();
    manual_node.status = "online".to_string();
    manual_node.is_manual = true;
    manual_node.tunnel_mode = false;
    manual_node.tunnel_connected = false;
    manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
    manual_node.proxy_username = Some("alice".to_string());
    manual_node.proxy_password = Some("supersecret".to_string());

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
                proxy_node_repository,
            )),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes/proxy-node-manual"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["node"]["id"], "proxy-node-manual");
    assert_eq!(payload["node"]["proxy_username"], "alice");
    assert_eq!(payload["node"]["has_proxy_password"], true);
    assert!(payload["node"].get("proxy_password").is_none());

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_reports_active_proxy_upgrade_rollout_in_proxy_node_list() {
    let mut alpha = sample_proxy_node("node-alpha");
    alpha.name = "alpha".to_string();
    alpha.status = "online".to_string();
    alpha.tunnel_connected = true;
    alpha.remote_config = None;
    alpha.proxy_metadata = Some(json!({ "version": "1.9.0" }));

    let mut beta = sample_proxy_node("node-beta");
    beta.name = "beta".to_string();
    beta.status = "online".to_string();
    beta.tunnel_connected = true;
    beta.remote_config = None;
    beta.proxy_metadata = Some(json!({ "version": "1.9.0" }));

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
    let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
        .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());

    let rollout = start_proxy_upgrade_rollout(
        &data_state,
        "2.0.0".to_string(),
        2,
        120,
        Some(crate::maintenance::ProxyUpgradeRolloutProbeConfig {
            url: "https://probe.example/health".to_string(),
            timeout_secs: 15,
        }),
    )
    .await
    .expect("rollout should start");
    assert_eq!(rollout.updated, 2);

    data_state
        .apply_proxy_node_heartbeat(&ProxyNodeHeartbeatMutation {
            node_id: "node-alpha".to_string(),
            expected_tunnel_generation: None,
            heartbeat_interval: None,
            active_connections: None,
            total_requests_delta: None,
            avg_latency_ms: None,
            failed_requests_delta: None,
            dns_failures_delta: None,
            stream_errors_delta: None,
            proxy_metadata: None,
            proxy_version: Some("2.0.0".to_string()),
        })
        .await
        .expect("heartbeat should apply");
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["rollout"]["version"], "2.0.0");
    assert_eq!(payload["rollout"]["batch_size"], 2);
    assert_eq!(payload["rollout"]["cooldown_secs"], 120);
    assert_eq!(
        payload["rollout"]["probe"]["url"],
        "https://probe.example/health"
    );
    assert_eq!(payload["rollout"]["probe"]["timeout_secs"], 15);
    assert_eq!(
        payload["rollout"]["pending_node_ids"],
        json!(["node-alpha", "node-beta"])
    );
    assert_eq!(payload["rollout"]["completed_node_ids"], json!([]));
    assert_eq!(payload["rollout"]["conflict_node_ids"], json!([]));
    assert_eq!(payload["rollout"]["blocked"], true);
    assert!(payload["rollout"]["started_at"].is_string());
    assert!(payload["rollout"]["last_dispatched_at"].is_string());
    assert!(payload["rollout"]["updated_at"].is_string());

    let tracked_nodes = payload["rollout"]["tracked_nodes"]
        .as_array()
        .expect("tracked_nodes should be array");
    assert_eq!(tracked_nodes.len(), 2);
    let alpha_status = tracked_nodes
        .iter()
        .find(|tracked| tracked["node_id"] == "node-alpha")
        .expect("alpha status should exist");
    assert_eq!(alpha_status["state"], "awaiting_traffic");
    assert!(alpha_status["version_confirmed_at"].is_string());
    assert!(alpha_status["traffic_confirmed_at"].is_null());

    let beta_status = tracked_nodes
        .iter()
        .find(|tracked| tracked["node_id"] == "node-beta")
        .expect("beta status should exist");
    assert_eq!(beta_status["state"], "awaiting_version");
    assert!(beta_status["version_confirmed_at"].is_null());
    assert!(beta_status["traffic_confirmed_at"].is_null());

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_cancels_active_proxy_upgrade_rollout_locally() {
    let mut alpha = sample_proxy_node("node-alpha");
    alpha.name = "alpha".to_string();
    alpha.status = "online".to_string();
    alpha.tunnel_connected = true;
    alpha.remote_config = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![alpha]));
    let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
        .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());

    let rollout = start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
        .await
        .expect("rollout should start");
    assert!(rollout.rollout_active);

    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state.clone()),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/upgrade/cancel"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["cancelled"], true);
    assert_eq!(payload["version"], "2.0.0");
    assert_eq!(payload["pending_node_ids"], json!(["node-alpha"]));
    assert_eq!(payload["conflict_node_ids"], json!([]));

    let list_response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    let list_payload: serde_json::Value =
        list_response.json().await.expect("json body should parse");
    assert!(list_payload["rollout"].is_null());

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_clears_proxy_upgrade_rollout_conflicts_locally() {
    let mut alpha = sample_proxy_node("node-alpha");
    alpha.name = "alpha".to_string();
    alpha.status = "online".to_string();
    alpha.tunnel_connected = true;
    alpha.remote_config = None;

    let mut beta = sample_proxy_node("node-beta");
    beta.name = "beta".to_string();
    beta.status = "online".to_string();
    beta.tunnel_connected = true;
    beta.remote_config = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
    let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
        .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());

    let rollout = start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
        .await
        .expect("rollout should start");
    assert_eq!(rollout.updated, 1);

    data_state
        .update_proxy_node_remote_config(
            &aether_data::repository::proxy_nodes::ProxyNodeRemoteConfigMutation {
                node_id: "node-beta".to_string(),
                expected_tunnel_generation: None,
                node_name: None,
                allowed_ports: None,
                log_level: None,
                heartbeat_interval: None,
                scheduling_state: None,
                upgrade_to: Some(Some("3.0.0".to_string())),
            },
        )
        .await
        .expect("conflict target should update");

    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state.clone()),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/upgrade/clear-conflicts"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["cleared"], 1);
    assert_eq!(payload["node_ids"], json!(["node-beta"]));
    assert_eq!(payload["blocked"], true);
    assert_eq!(payload["pending_node_ids"], json!(["node-alpha"]));

    let updated_beta = data_state
        .find_proxy_node("node-beta")
        .await
        .expect("node lookup should succeed")
        .expect("beta should exist");
    let beta_upgrade_to = updated_beta
        .remote_config
        .as_ref()
        .and_then(|value| value.get("upgrade_to"))
        .cloned()
        .unwrap_or(serde_json::Value::Null);
    assert!(beta_upgrade_to.is_null());

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_skips_proxy_upgrade_rollout_node_and_advances_next_wave_locally() {
    let mut alpha = sample_proxy_node("node-alpha");
    alpha.name = "alpha".to_string();
    alpha.status = "online".to_string();
    alpha.tunnel_connected = true;
    alpha.remote_config = None;

    let mut beta = sample_proxy_node("node-beta");
    beta.name = "beta".to_string();
    beta.status = "online".to_string();
    beta.tunnel_connected = true;
    beta.remote_config = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
    let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
        .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());

    let rollout = start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
        .await
        .expect("rollout should start");
    assert_eq!(rollout.node_ids, vec!["node-alpha"]);

    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state.clone()),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-alpha/upgrade/skip"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["node_id"], "node-alpha");
    assert_eq!(payload["skipped_node_ids"], json!(["node-alpha"]));
    assert_eq!(payload["updated"], 1);
    assert_eq!(payload["pending_node_ids"], json!(["node-beta"]));

    let alpha_after = data_state
        .find_proxy_node("node-alpha")
        .await
        .expect("node lookup should succeed")
        .expect("alpha should exist");
    let alpha_upgrade_to = alpha_after
        .remote_config
        .as_ref()
        .and_then(|value| value.get("upgrade_to"))
        .cloned()
        .unwrap_or(serde_json::Value::Null);
    assert!(alpha_upgrade_to.is_null());

    let list_response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    let list_payload: serde_json::Value =
        list_response.json().await.expect("json body should parse");
    assert_eq!(
        list_payload["rollout"]["skipped_node_ids"],
        json!(["node-alpha"])
    );

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_retries_proxy_upgrade_rollout_node_locally() {
    let mut alpha = sample_proxy_node("node-alpha");
    alpha.name = "alpha".to_string();
    alpha.status = "online".to_string();
    alpha.tunnel_connected = true;
    alpha.remote_config = None;

    let mut beta = sample_proxy_node("node-beta");
    beta.name = "beta".to_string();
    beta.status = "online".to_string();
    beta.tunnel_connected = true;
    beta.remote_config = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
    let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
        .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());

    start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
        .await
        .expect("rollout should start");
    let _ = skip_proxy_upgrade_rollout_node(&data_state, "node-alpha")
        .await
        .expect("skip should succeed");

    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state.clone()),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-alpha/upgrade/retry"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["node_id"], "node-alpha");
    assert_eq!(payload["skipped_node_ids"], json!([]));
    assert_eq!(payload["blocked"], true);

    let alpha_after = data_state
        .find_proxy_node("node-alpha")
        .await
        .expect("node lookup should succeed")
        .expect("alpha should exist");
    let alpha_upgrade_to = alpha_after
        .remote_config
        .as_ref()
        .and_then(|value| value.get("upgrade_to"))
        .and_then(serde_json::Value::as_str);
    assert_eq!(alpha_upgrade_to, Some("2.0.0"));

    let list_response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    let list_payload: serde_json::Value =
        list_response.json().await.expect("json body should parse");
    assert_eq!(list_payload["rollout"]["skipped_node_ids"], json!([]));
    let tracked_nodes = list_payload["rollout"]["tracked_nodes"]
        .as_array()
        .expect("tracked nodes should be array");
    let alpha_status = tracked_nodes
        .iter()
        .find(|tracked| tracked["node_id"] == "node-alpha")
        .expect("alpha should be tracked again");
    assert_eq!(alpha_status["state"], "awaiting_version");

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_restores_skipped_proxy_upgrade_rollout_nodes_locally() {
    let mut alpha = sample_proxy_node("node-alpha");
    alpha.name = "alpha".to_string();
    alpha.status = "online".to_string();
    alpha.tunnel_connected = true;
    alpha.remote_config = None;

    let mut beta = sample_proxy_node("node-beta");
    beta.name = "beta".to_string();
    beta.status = "online".to_string();
    beta.tunnel_connected = true;
    beta.remote_config = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
    let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
        .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());

    start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
        .await
        .expect("rollout should start");
    let _ = skip_proxy_upgrade_rollout_node(&data_state, "node-alpha")
        .await
        .expect("skip should succeed");

    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state.clone()),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/upgrade/restore-skipped"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["restored"], 1);
    assert_eq!(payload["node_ids"], json!(["node-alpha"]));
    assert_eq!(payload["skipped_node_ids"], json!([]));
    assert_eq!(payload["updated"], 0);
    assert_eq!(payload["blocked"], true);
    assert_eq!(payload["pending_node_ids"], json!(["node-beta"]));

    let list_response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    let list_payload: serde_json::Value =
        list_response.json().await.expect("json body should parse");
    assert_eq!(list_payload["rollout"]["skipped_node_ids"], json!([]));

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_registers_and_unregisters_proxy_nodes_locally_with_management_token_principal() {
    let upstream_hits = Arc::new(Mutex::new(0usize));
    let upstream_hits_clone = Arc::clone(&upstream_hits);
    let upstream = Router::new().route(
        "/api/admin/proxy-nodes/register",
        any(move |_request: Request| {
            let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
            async move {
                *upstream_hits_inner.lock().expect("mutex should lock") += 1;
                (StatusCode::OK, Body::from("unexpected upstream hit"))
            }
        }),
    );

    let raw_token = "ae-proxy-register-test";
    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::default());
    let state = AppState::new().expect("gateway should build");
    let admin_user = state
        .create_local_auth_user_with_settings(
            Some("[email protected]".to_string()),
            true,
            "admin".to_string(),
            "hash".to_string(),
            "admin".to_string(),
            None,
            None,
            None,
            None,
        )
        .await
        .expect("admin user should be created")
        .expect("admin user should exist");
    let mut management_token =
        sample_management_token("token-proxy-register", &admin_user.id, "proxy-admin", true);
    management_token.token.allowed_ips = None;
    let management_token_repository =
        Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
            vec![management_token],
            vec![(
                hash_management_token(raw_token),
                "token-proxy-register".to_string(),
            )],
        ));

    let (upstream_url, upstream_handle) = start_server(upstream).await;
    let state = state.with_data_state_for_tests(
        GatewayDataState::with_management_token_repository_for_tests(management_token_repository)
            .attach_proxy_node_repository_for_tests(proxy_node_repository),
    );
    let token_lookup = state
        .get_management_token_with_user_by_hash(&hash_management_token(raw_token))
        .await
        .expect("token lookup should succeed");
    assert!(token_lookup.is_some());
    let user_lookup = state
        .find_user_auth_by_id(&admin_user.id)
        .await
        .expect("user lookup should succeed");
    assert!(user_lookup.is_some());
    let gateway = build_router_with_state(state);
    let (gateway_url, gateway_handle) = start_server(gateway).await;
    let client = reqwest::Client::new();

    let register_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/register"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .bearer_auth(raw_token)
        .json(&json!({
            "name": "proxy-1",
            "ip": "1.1.1.1",
            "port": 0,
            "heartbeat_interval": 30,
            "tunnel_mode": true
        }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(register_response.status(), StatusCode::OK);
    let register_payload: serde_json::Value = register_response
        .json()
        .await
        .expect("json body should parse");
    let node_id = register_payload["node_id"]
        .as_str()
        .expect("node_id should be present")
        .to_string();
    assert_eq!(register_payload["node"]["name"], "proxy-1");
    assert_eq!(register_payload["node"]["status"], "offline");
    assert_eq!(
        register_payload["node"]["registered_by"],
        json!(admin_user.id)
    );

    let unregister_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/unregister"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .bearer_auth(raw_token)
        .json(&json!({ "node_id": node_id }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(unregister_response.status(), StatusCode::OK);
    let unregister_payload: serde_json::Value = unregister_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(unregister_payload["message"], "unregistered");

    assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

    gateway_handle.abort();
    upstream_handle.abort();
}

#[tokio::test]
async fn gateway_allows_management_token_with_proxy_nodes_write_permission() {
    let raw_token = "ae-proxy-register-write-permission";
    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::default());
    let state = AppState::new().expect("gateway should build");
    let admin_user = state
        .create_local_auth_user_with_settings(
            Some("[email protected]".to_string()),
            true,
            "admin".to_string(),
            "hash".to_string(),
            "admin".to_string(),
            None,
            None,
            None,
            None,
        )
        .await
        .expect("admin user should be created")
        .expect("admin user should exist");
    let mut management_token = sample_management_token(
        "token-proxy-register-write",
        &admin_user.id,
        "proxy-admin-write",
        true,
    );
    management_token.token.allowed_ips = None;
    management_token.token.permissions = Some(json!(["admin:proxy_nodes:write"]));
    let management_token_repository =
        Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
            vec![management_token],
            vec![(
                hash_management_token(raw_token),
                "token-proxy-register-write".to_string(),
            )],
        ));

    let state = state.with_data_state_for_tests(
        GatewayDataState::with_management_token_repository_for_tests(management_token_repository)
            .attach_proxy_node_repository_for_tests(proxy_node_repository),
    );
    let gateway = build_router_with_state(state);
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let register_response = reqwest::Client::new()
        .post(format!("{gateway_url}/api/admin/proxy-nodes/register"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .bearer_auth(raw_token)
        .json(&json!({
            "name": "proxy-write",
            "ip": "1.1.1.1",
            "port": 0,
            "heartbeat_interval": 30,
            "tunnel_mode": true
        }))
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(register_response.status(), StatusCode::OK);
    let register_payload: serde_json::Value = register_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(register_payload["node"]["name"], "proxy-write");
    assert_eq!(
        register_payload["node"]["registered_by"],
        json!(admin_user.id)
    );

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_rejects_management_token_without_required_admin_route_permission() {
    let upstream_hits = Arc::new(Mutex::new(0usize));
    let upstream_hits_clone = Arc::clone(&upstream_hits);
    let upstream = Router::new().route(
        "/api/admin/proxy-nodes/register",
        any(move |_request: Request| {
            let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
            async move {
                *upstream_hits_inner.lock().expect("mutex should lock") += 1;
                (StatusCode::OK, Body::from("unexpected upstream hit"))
            }
        }),
    );

    let raw_token = "ae-proxy-register-denied";
    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::default());
    let state = AppState::new().expect("gateway should build");
    let admin_user = state
        .create_local_auth_user_with_settings(
            Some("[email protected]".to_string()),
            true,
            "admin".to_string(),
            "hash".to_string(),
            "admin".to_string(),
            None,
            None,
            None,
            None,
        )
        .await
        .expect("admin user should be created")
        .expect("admin user should exist");
    let mut management_token = sample_management_token(
        "token-proxy-register-denied",
        &admin_user.id,
        "proxy-admin-denied",
        true,
    );
    management_token.token.allowed_ips = None;
    management_token.token.permissions = Some(json!(["admin:usage:read"]));
    let management_token_repository =
        Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
            vec![management_token],
            vec![(
                hash_management_token(raw_token),
                "token-proxy-register-denied".to_string(),
            )],
        ));

    let (upstream_url, upstream_handle) = start_server(upstream).await;
    let state = state.with_data_state_for_tests(
        GatewayDataState::with_management_token_repository_for_tests(management_token_repository)
            .attach_proxy_node_repository_for_tests(proxy_node_repository),
    );
    let gateway = build_router_with_state(state);
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let register_response = reqwest::Client::new()
        .post(format!("{gateway_url}/api/admin/proxy-nodes/register"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .bearer_auth(raw_token)
        .json(&json!({
            "name": "proxy-denied",
            "ip": "1.1.1.1",
            "port": 0,
            "heartbeat_interval": 30,
            "tunnel_mode": true
        }))
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(register_response.status(), StatusCode::FORBIDDEN);
    let denied_payload: serde_json::Value = register_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(
        denied_payload["detail"],
        json!("management token permission denied")
    );
    assert_eq!(
        denied_payload["required_permission"],
        json!("admin:proxy_nodes:write")
    );
    assert_eq!(denied_payload["route_family"], json!("proxy_nodes_manage"));
    assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

    gateway_handle.abort();
    upstream_handle.abort();
    drop(upstream_url);
}

#[tokio::test]
async fn proxy_node_install_session_requires_admin_and_preserves_parent_token_constraints() {
    let write_raw_token = "ae-proxy-install-write-only";
    let admin_raw_token = "ae-proxy-install-admin";
    let state = AppState::new().expect("gateway should build");
    let admin_user = state
        .create_local_auth_user_with_settings(
            Some("[email protected]".to_string()),
            true,
            "admin".to_string(),
            "hash".to_string(),
            "admin".to_string(),
            None,
            None,
            None,
            None,
        )
        .await
        .expect("admin user should be created")
        .expect("admin user should exist");

    let parent_allowed_ips = json!(["127.0.0.1"]);
    let parent_expires_at = 4_102_444_800;
    let mut write_parent = sample_management_token(
        "token-proxy-install-write",
        &admin_user.id,
        "proxy-install-write",
        true,
    );
    write_parent.token.allowed_ips = Some(parent_allowed_ips.clone());
    write_parent.token.permissions = Some(json!(["admin:proxy_nodes:write"]));
    write_parent.token.expires_at_unix_secs = Some(parent_expires_at);
    let mut admin_parent = sample_management_token(
        "token-proxy-install-admin",
        &admin_user.id,
        "proxy-install-admin",
        true,
    );
    admin_parent.token.allowed_ips = Some(parent_allowed_ips.clone());
    admin_parent.token.permissions = Some(json!(["admin:proxy_nodes:admin"]));
    admin_parent.token.expires_at_unix_secs = Some(parent_expires_at);

    let management_token_repository =
        Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
            vec![write_parent, admin_parent],
            vec![
                (
                    hash_management_token(write_raw_token),
                    "token-proxy-install-write".to_string(),
                ),
                (
                    hash_management_token(admin_raw_token),
                    "token-proxy-install-admin".to_string(),
                ),
            ],
        ));
    let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users([
        admin_user.clone()
    ]));
    let state = state.with_data_state_for_tests(
        GatewayDataState::with_management_token_repository_for_tests(Arc::clone(
            &management_token_repository,
        ))
        .with_user_reader(user_repository),
    );
    let gateway = build_router_with_state(state);
    let (gateway_url, gateway_handle) = start_server(gateway).await;
    let client = reqwest::Client::new();

    let denied = client
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/install-sessions"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .bearer_auth(write_raw_token)
        .json(&json!({ "node_name": "write-only-node" }))
        .send()
        .await
        .expect("write-only install request should complete");
    assert_eq!(denied.status(), StatusCode::FORBIDDEN);
    let denied_payload: serde_json::Value = denied.json().await.expect("denial should be json");
    assert_eq!(
        denied_payload["required_permission"],
        json!("admin:proxy_nodes:admin")
    );

    let accepted = client
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/install-sessions"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .bearer_auth(admin_raw_token)
        .json(&json!({ "node_name": "constrained-node" }))
        .send()
        .await
        .expect("admin install request should complete");
    let accepted_status = accepted.status();
    let accepted_body = accepted.text().await.expect("response body should read");
    assert_eq!(
        accepted_status,
        StatusCode::OK,
        "unexpected response body: {accepted_body}"
    );
    let accepted_payload: serde_json::Value =
        serde_json::from_str(&accepted_body).expect("install response should be JSON");
    let install_code = accepted_payload["install_code"]
        .as_str()
        .expect("install code should be returned")
        .to_string();

    let tokens = management_token_repository
        .list_management_tokens(&ManagementTokenListQuery {
            user_id: Some(admin_user.id.clone()),
            is_active: None,
            offset: 0,
            limit: 10,
        })
        .await
        .expect("management tokens should list");
    assert_eq!(tokens.total, 3);
    let child = tokens
        .items
        .iter()
        .find(|item| {
            !matches!(
                item.token.id.as_str(),
                "token-proxy-install-write" | "token-proxy-install-admin"
            )
        })
        .expect("install session should create one child management token");
    assert_eq!(child.token.allowed_ips, Some(parent_allowed_ips));
    assert_eq!(child.token.expires_at_unix_secs, Some(parent_expires_at));
    assert_eq!(
        child.token.permissions,
        Some(json!(["admin:proxy_nodes:write"]))
    );
    assert!(
        !child.token.is_active,
        "unused install-session token must remain disabled"
    );
    let child_id = child.token.id.clone();

    // The in-memory repository cannot atomically verify the administrator row and therefore
    // rejects one-time activation. SQL-backed repositories cover the successful atomic path.
    let install_script = client
        .get(format!("{gateway_url}/install-tunnel/{install_code}"))
        .send()
        .await
        .expect("tunnel install script should receive a response");
    assert_eq!(install_script.status(), StatusCode::NOT_FOUND);

    let consumed = management_token_repository
        .get_management_token_with_user(&child_id)
        .await
        .expect("consumed token lookup should succeed");
    assert!(
        consumed.is_none(),
        "failed one-time activation must discard the pending bearer"
    );

    let replay = client
        .get(format!("{gateway_url}/install-tunnel/{install_code}"))
        .send()
        .await
        .expect("tunnel install replay should receive a response");
    assert_eq!(replay.status(), StatusCode::NOT_FOUND);

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_rejects_proxy_node_registration_when_allowed_ips_is_json_null() {
    let raw_token = "ae_proxy_register_json_null";
    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::default());
    let state = AppState::new().expect("gateway should build");
    let admin_user = state
        .create_local_auth_user_with_settings(
            Some("[email protected]".to_string()),
            true,
            "admin".to_string(),
            "hash".to_string(),
            "admin".to_string(),
            None,
            None,
            None,
            None,
        )
        .await
        .expect("admin user should be created")
        .expect("admin user should exist");
    let mut management_token = sample_management_token(
        "token-proxy-register-json-null",
        &admin_user.id,
        "proxy-admin",
        true,
    );
    management_token.token.allowed_ips = Some(serde_json::Value::Null);
    let management_token_repository =
        Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
            vec![management_token],
            vec![(
                hash_management_token(raw_token),
                "token-proxy-register-json-null".to_string(),
            )],
        ));

    let state = state.with_data_state_for_tests(
        GatewayDataState::with_management_token_repository_for_tests(management_token_repository)
            .attach_proxy_node_repository_for_tests(proxy_node_repository),
    );
    let gateway = build_router_with_state(state);
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let register_response = reqwest::Client::new()
        .post(format!("{gateway_url}/api/admin/proxy-nodes/register"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .bearer_auth(raw_token)
        .json(&json!({
            "name": "proxy-json-null",
            "ip": "1.1.1.1",
            "port": 0,
            "heartbeat_interval": 30,
            "tunnel_mode": true
        }))
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(register_response.status(), StatusCode::UNAUTHORIZED);
    let _ = register_response
        .bytes()
        .await
        .expect("error body should be readable");

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
    let proxy_auths = Arc::new(Mutex::new(Vec::<Option<String>>::new()));
    let proxy_auths_clone = Arc::clone(&proxy_auths);
    let proxy = Router::new().fallback(any(move |request: Request| {
        let proxy_auths_inner = Arc::clone(&proxy_auths_clone);
        async move {
            proxy_auths_inner.lock().expect("mutex should lock").push(
                request
                    .headers()
                    .get("proxy-authorization")
                    .and_then(|value| value.to_str().ok())
                    .map(str::to_string),
            );
            (
                StatusCode::OK,
                Body::from("fl=1234\nip=203.0.113.10\nwarp=off\n"),
            )
        }
    }));
    let (proxy_url, proxy_handle) = start_server(proxy).await;
    let _probe_url_guard = crate::handlers::admin::override_proxy_connectivity_probe_url_for_tests(
        "http://probe.example/cdn-cgi/trace",
    );

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::default());
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
                proxy_node_repository,
            )),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;
    let client = reqwest::Client::new();

    let create_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/manual"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({
            "name": "manual-node",
            "proxy_url": proxy_url.clone(),
            "username": "alice",
            "password": "supersecret",
            "region": "US-West"
        }))
        .send()
        .await
        .expect("create request should succeed");
    let create_status = create_response.status();
    let create_body = create_response
        .text()
        .await
        .expect("body should read as text");
    assert_eq!(create_status, StatusCode::OK, "create body: {create_body}");
    let create_payload: serde_json::Value =
        serde_json::from_str(&create_body).expect("json body should parse");
    let node_id = create_payload["node_id"]
        .as_str()
        .expect("node id should exist")
        .to_string();
    assert_eq!(create_payload["node"]["is_manual"], true);
    assert_eq!(create_payload["node"]["status"], "online");
    assert_eq!(create_payload["node"]["proxy_url"], proxy_url);
    assert_eq!(create_payload["node"]["proxy_username"], "alice");
    assert_eq!(create_payload["node"]["has_proxy_password"], true);
    assert!(create_payload["node"].get("proxy_password").is_none());

    let test_url_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/test-url"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({
            "proxy_url": proxy_url.clone(),
            "username": "alice",
            "password": "supersecret"
        }))
        .send()
        .await
        .expect("test-url request should succeed");
    assert_eq!(test_url_response.status(), StatusCode::OK);
    let test_url_payload: serde_json::Value = test_url_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(test_url_payload["success"], true);
    assert!(test_url_payload["latency_ms"].is_u64());
    assert_eq!(test_url_payload["exit_ip"], "203.0.113.10");

    let test_node_response = client
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/{node_id}/test"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("test-node request should succeed");
    assert_eq!(test_node_response.status(), StatusCode::OK);
    let test_node_payload: serde_json::Value = test_node_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(test_node_payload["success"], true);
    assert_eq!(test_node_payload["exit_ip"], "203.0.113.10");

    let expected_proxy_auth = format!(
        "Basic {}",
        base64::engine::general_purpose::STANDARD.encode("alice:supersecret")
    );
    assert_eq!(
        proxy_auths.lock().expect("mutex should lock").as_slice(),
        [
            Some(expected_proxy_auth.clone()),
            Some(expected_proxy_auth.clone()),
        ]
    );

    let update_response = client
        .patch(format!("{gateway_url}/api/admin/proxy-nodes/{node_id}"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({
            "name": "manual-node-updated",
            "region": "US-East"
        }))
        .send()
        .await
        .expect("update request should succeed");
    assert_eq!(update_response.status(), StatusCode::OK);
    let update_payload: serde_json::Value = update_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(update_payload["node"]["name"], "manual-node-updated");
    assert_eq!(update_payload["node"]["region"], "US-East");
    assert_eq!(update_payload["node"]["proxy_url"], proxy_url);

    gateway_handle.abort();
    proxy_handle.abort();
}

#[tokio::test]
async fn gateway_tests_disconnected_tunnel_proxy_nodes_locally() {
    let _probe_url_guard = crate::handlers::admin::override_proxy_connectivity_probe_url_for_tests(
        "https://www.cloudflare.com/cdn-cgi/trace",
    );

    let proxy_node_repository =
        Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
            "node-offline",
        )]));
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
                proxy_node_repository,
            )),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .post(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-offline/test"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["success"], false);
    assert_eq!(payload["error"], "tunnel 未连接");
    assert_eq!(
        payload["probe_url"],
        "https://www.cloudflare.com/cdn-cgi/trace"
    );
    assert_eq!(payload["timeout_secs"], 10);

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_tests_connected_tunnel_proxy_nodes_with_active_probe() {
    let _probe_url_guard = crate::handlers::admin::override_proxy_connectivity_probe_url_for_tests(
        "https://probe.example/cdn-cgi/trace",
    );

    let mut node = with_tunnel_control_plane_key(
        sample_proxy_node("node-online"),
        TUNNEL_CONTROL_PLANE_TEST_PSK,
    );
    node.status = "online".to_string();
    node.tunnel_connected = true;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![node]));
    let state = AppState::new()
        .expect("gateway should build")
        .with_data_state_for_tests(
            GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
                .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
        );
    let tunnel_state = state.tunnel.app_state();
    let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
    let (proxy_close_tx, _) = watch::channel(false);
    tunnel_state.hub.register_proxy(Arc::new(
        TunnelProxyConn::new(
            500,
            "node-online".to_string(),
            "Node Online".to_string(),
            proxy_tx,
            proxy_close_tx,
            16,
            2,
        )
        .with_tunnel_generation(TUNNEL_CONTROL_PLANE_TEST_GENERATION.to_string())
        .with_authenticated_key(TUNNEL_CONTROL_PLANE_TEST_PSK.to_string()),
    ));

    let gateway = build_router_with_state(state);
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let request_task = tokio::spawn({
        let gateway_url = gateway_url.clone();
        async move {
            reqwest::Client::new()
                .post(format!(
                    "{gateway_url}/api/admin/proxy-nodes/node-online/test"
                ))
                .header(GATEWAY_HEADER, "rust-phase3b")
                .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
                .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
                .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
                .send()
                .await
        }
    });

    let request_headers = match recv_tunnel_test_frame(&mut proxy_rx, "probe headers frame").await {
        Message::Binary(data) => data,
        other => panic!("unexpected message: {other:?}"),
    };
    let request_header =
        tunnel_protocol::FrameHeader::parse(&request_headers).expect("request header should parse");
    assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
    let meta_payload = tunnel_protocol::decode_payload(&request_headers, &request_header)
        .expect("request header payload should decode");
    let meta: tunnel_protocol::RequestMeta =
        serde_json::from_slice(&meta_payload).expect("request meta should parse");
    assert_eq!(meta.method, "GET");
    assert_eq!(meta.url, "https://probe.example/cdn-cgi/trace");
    assert_eq!(meta.follow_redirects, Some(false));

    let request_body = match recv_tunnel_test_frame(&mut proxy_rx, "probe body frame").await {
        Message::Binary(data) => data,
        other => panic!("unexpected message: {other:?}"),
    };
    let request_body_header =
        tunnel_protocol::FrameHeader::parse(&request_body).expect("request body should parse");
    assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY);
    assert_ne!(
        request_body_header.flags & tunnel_protocol::FLAG_END_STREAM,
        0,
        "probe body frame should close the stream"
    );

    let response_meta = tunnel_protocol::ResponseMeta {
        status: 200,
        headers: vec![("content-type".to_string(), "text/plain".to_string())],
    };
    let response_meta_bytes =
        serde_json::to_vec(&response_meta).expect("response meta should serialize");
    let mut response_headers_frame = tunnel_protocol::encode_frame(
        request_header.stream_id,
        tunnel_protocol::RESPONSE_HEADERS,
        0,
        &response_meta_bytes,
    );
    tunnel_state
        .hub
        .handle_proxy_frame(500, &mut response_headers_frame)
        .await;

    let mut response_body_frame = tunnel_protocol::encode_frame(
        request_header.stream_id,
        tunnel_protocol::RESPONSE_BODY,
        0,
        b"fl=1234\nip=203.0.113.10\nwarp=off\n",
    );
    tunnel_state
        .hub
        .handle_proxy_frame(500, &mut response_body_frame)
        .await;

    let mut response_end_frame = tunnel_protocol::encode_frame(
        request_header.stream_id,
        tunnel_protocol::STREAM_END,
        0,
        &[],
    );
    tunnel_state
        .hub
        .handle_proxy_frame(500, &mut response_end_frame)
        .await;

    let response = request_task
        .await
        .expect("request task should complete")
        .expect("test-node request should succeed");
    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["success"], true);
    assert!(payload["latency_ms"].is_u64());
    assert_eq!(payload["exit_ip"], "203.0.113.10");
    assert_eq!(payload["probe_url"], "https://probe.example/cdn-cgi/trace");
    assert_eq!(payload["timeout_secs"], 10);

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_deletes_proxy_nodes_and_clears_proxy_refs_locally() {
    let mut manual_node = sample_proxy_node("manual-node-1");
    manual_node.name = "manual-node-1".to_string();
    manual_node.status = "online".to_string();
    manual_node.is_manual = true;
    manual_node.tunnel_mode = false;
    manual_node.tunnel_connected = false;
    manual_node.proxy_url = Some("http://127.0.0.1:8899".to_string());
    manual_node.last_heartbeat_at_unix_secs = None;
    manual_node.tunnel_connected_at_unix_secs = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
    let mut provider = sample_provider("provider-1", "OpenAI", 10);
    provider.proxy = Some(json!({ "node_id": "manual-node-1", "enabled": true }));
    let mut endpoint = sample_endpoint(
        "endpoint-1",
        "provider-1",
        "openai:chat",
        "https://example.com/v1",
    );
    endpoint.proxy = Some(json!({ "node_id": "manual-node-1", "enabled": true }));
    let mut key = sample_key("key-1", "provider-1", "openai:chat", "sk-test");
    key.proxy = Some(json!({ "node_id": "manual-node-1", "enabled": true }));
    let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
        vec![provider],
        vec![endpoint],
        vec![key],
    ));

    let data_state =
        GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&proxy_node_repository))
            .attach_provider_catalog_repository_for_tests(Arc::clone(&provider_catalog_repository))
            .with_system_config_values_for_tests(vec![
                ("system_proxy_node_id".to_string(), json!("manual-node-1")),
                (
                    "external_models_proxy_node_id".to_string(),
                    json!("manual-node-1"),
                ),
            ]);
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state.clone()),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .delete(format!("{gateway_url}/api/admin/proxy-nodes/manual-node-1"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    assert_eq!(payload["cleared_system_proxy"], true);
    assert_eq!(payload["cleared_external_models_proxy"], true);
    assert_eq!(payload["cleared_providers"], 1);
    assert_eq!(payload["cleared_endpoints"], 1);
    assert_eq!(payload["cleared_keys"], 1);

    assert!(data_state
        .find_proxy_node("manual-node-1")
        .await
        .expect("node lookup should succeed")
        .is_none());
    assert_eq!(
        data_state
            .find_system_config_value("system_proxy_node_id")
            .await
            .expect("system config lookup should succeed"),
        Some(serde_json::Value::Null)
    );
    assert_eq!(
        data_state
            .find_system_config_value("external_models_proxy_node_id")
            .await
            .expect("external models proxy config lookup should succeed"),
        Some(serde_json::Value::Null)
    );

    let provider_ids = vec!["provider-1".to_string()];
    let providers = data_state
        .list_provider_catalog_providers(false)
        .await
        .expect("provider list should succeed");
    assert!(providers.iter().all(|provider| provider.proxy.is_none()));
    let endpoints = data_state
        .list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
        .await
        .expect("endpoint list should succeed");
    assert!(endpoints.iter().all(|endpoint| endpoint.proxy.is_none()));
    let keys = data_state
        .list_provider_catalog_keys_by_provider_ids(&provider_ids)
        .await
        .expect("key list should succeed");
    assert!(keys.iter().all(|key| key.proxy.is_none()));

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_continues_proxy_reference_cleanup_when_external_models_cache_delete_fails() {
    let mut manual_node = sample_proxy_node("manual-node-cache-failure");
    manual_node.status = "online".to_string();
    manual_node.is_manual = true;
    manual_node.tunnel_mode = false;
    manual_node.tunnel_connected = false;
    manual_node.proxy_url = Some("http://127.0.0.1:8899".to_string());
    manual_node.last_heartbeat_at_unix_secs = None;
    manual_node.tunnel_connected_at_unix_secs = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
    let mut provider = sample_provider("provider-cache-failure", "OpenAI", 10);
    provider.proxy = Some(json!({
        "node_id": "manual-node-cache-failure",
        "enabled": true
    }));
    let mut endpoint = sample_endpoint(
        "endpoint-cache-failure",
        "provider-cache-failure",
        "openai:chat",
        "https://example.com/v1",
    );
    endpoint.proxy = Some(json!({
        "node_id": "manual-node-cache-failure",
        "enabled": true
    }));
    let mut key = sample_key(
        "key-cache-failure",
        "provider-cache-failure",
        "openai:chat",
        "sk-test",
    );
    key.proxy = Some(json!({
        "node_id": "manual-node-cache-failure",
        "enabled": true
    }));
    let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
        vec![provider],
        vec![endpoint],
        vec![key],
    ));
    let data_state =
        GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&proxy_node_repository))
            .attach_provider_catalog_repository_for_tests(Arc::clone(&provider_catalog_repository))
            .with_system_config_values_for_tests([(
                "external_models_proxy_node_id".to_string(),
                json!("manual-node-cache-failure"),
            )]);
    let state = AppState::new()
        .expect("gateway should build")
        .with_data_state_for_tests(data_state.clone());

    let cleanup = crate::handlers::admin::clear_proxy_node_references_with_cache_failure_for_tests(
        &state,
        "manual-node-cache-failure",
    )
    .await
    .expect("cache failure should not abort persistent reference cleanup");
    assert_eq!(cleanup["cleared_external_models_proxy"], true);
    assert_eq!(cleanup["external_models_cache_clear_succeeded"], false);
    assert_eq!(cleanup["cleared_providers"], 1);
    assert_eq!(cleanup["cleared_endpoints"], 1);
    assert_eq!(cleanup["cleared_keys"], 1);
    assert_eq!(
        data_state
            .find_system_config_value("external_models_proxy_node_id")
            .await
            .expect("external models proxy config lookup should succeed"),
        Some(serde_json::Value::Null)
    );
    assert!(data_state
        .find_proxy_node("manual-node-cache-failure")
        .await
        .expect("node lookup should succeed")
        .is_some());

    let provider_ids = vec!["provider-cache-failure".to_string()];
    let providers = data_state
        .list_provider_catalog_providers(false)
        .await
        .expect("provider list should succeed");
    assert!(providers.iter().all(|provider| provider.proxy.is_none()));
    let endpoints = data_state
        .list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
        .await
        .expect("endpoint list should succeed");
    assert!(endpoints.iter().all(|endpoint| endpoint.proxy.is_none()));
    let keys = data_state
        .list_provider_catalog_keys_by_provider_ids(&provider_ids)
        .await
        .expect("key list should succeed");
    assert!(keys.iter().all(|key| key.proxy.is_none()));
}

#[tokio::test]
async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_principal() {
    let upstream_hits = Arc::new(Mutex::new(0usize));
    let upstream_hits_clone = Arc::clone(&upstream_hits);
    let upstream = Router::new().route(
        "/api/admin/proxy-nodes/node-1/events",
        any(move |_request: Request| {
            let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
            async move {
                *upstream_hits_inner.lock().expect("mutex should lock") += 1;
                (StatusCode::OK, Body::from("unexpected upstream hit"))
            }
        }),
    );

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed_with_events(
        vec![sample_proxy_node("node-1")],
        vec![
            StoredProxyNodeEvent {
                id: 1,
                node_id: "node-1".to_string(),
                event_type: "connected".to_string(),
                detail: Some("older".to_string()),
                event_metadata: None,
                created_at_unix_ms: Some(1_710_000_000),
            },
            StoredProxyNodeEvent {
                id: 2,
                node_id: "node-1".to_string(),
                event_type: "disconnected".to_string(),
                detail: Some("newer".to_string()),
                event_metadata: None,
                created_at_unix_ms: Some(1_710_000_100),
            },
        ],
    ));

    let (upstream_url, upstream_handle) = start_server(upstream).await;
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
                proxy_node_repository,
            )),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;

    let response = reqwest::Client::new()
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-1/events?limit=1"
        ))
        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");

    assert_eq!(response.status(), StatusCode::OK);
    let payload: serde_json::Value = response.json().await.expect("json body should parse");
    let items = payload["items"].as_array().expect("items should be array");
    assert_eq!(items.len(), 1);
    assert_eq!(items[0]["id"], 2);
    assert_eq!(items[0]["event_type"], "disconnected");
    assert_eq!(items[0]["detail"], "newer");
    assert_eq!(items[0]["created_at"], "2024-03-09T16:01:40Z");
    assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

    gateway_handle.abort();
    upstream_handle.abort();
}

#[tokio::test]
async fn gateway_reports_proxy_node_metrics_and_filters_events_locally() {
    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
        with_tunnel_control_plane_key(sample_proxy_node("node-1"), TUNNEL_CONTROL_PLANE_TEST_PSK),
    ]));
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
                proxy_node_repository,
            )),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;
    let client = reqwest::Client::new();
    let now_unix_secs = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("system time should be after epoch")
        .as_secs();

    let baseline_heartbeat = json!({
        "node_id": "node-1",
        "heartbeat_session_id": "node-1-session",
        "heartbeat_id": 90,
        "heartbeat_interval": 30,
        "active_connections": 0,
        "proxy_metadata": {
            "tunnel_metrics": {
                "connect_errors": 0,
                "disconnects": 0,
                "error_events_total": 0,
                "ws_in_bytes": 0,
                "ws_out_bytes": 0,
                "ws_in_frames": 0,
                "ws_out_frames": 0,
                "heartbeat_rtt_last_ms": 0
            }
        },
        "proxy_version": "2.0.0"
    });
    let baseline_heartbeat_response = authenticated_tunnel_control_plane_request(
        &client,
        format!("{gateway_url}/api/internal/tunnel/heartbeat"),
        "/api/internal/tunnel/heartbeat",
        "node-1",
        &baseline_heartbeat,
    )
    .send()
    .await
    .expect("baseline heartbeat request should succeed");
    assert_eq!(baseline_heartbeat_response.status(), StatusCode::OK);

    let heartbeat = json!({
        "node_id": "node-1",
        "heartbeat_session_id": "node-1-session",
        "heartbeat_id": 91,
        "heartbeat_interval": 30,
        "active_connections": 7,
        "proxy_metadata": {
            "tunnel_metrics": {
                "connect_errors": 3,
                "disconnects": 1,
                "error_events_total": 1,
                "ws_in_bytes": 1000,
                "ws_out_bytes": 2000,
                "ws_in_frames": 10,
                "ws_out_frames": 20,
                "heartbeat_rtt_last_ms": 42
            },
            "recent_tunnel_errors": [{
                "timestamp_unix_secs": now_unix_secs,
                "category": "tcp_connect_timeout",
                "message": "tunnel TCP connect timeout"
            }]
        },
        "proxy_version": "2.0.0"
    });
    let heartbeat_response = authenticated_tunnel_control_plane_request(
        &client,
        format!("{gateway_url}/api/internal/tunnel/heartbeat"),
        "/api/internal/tunnel/heartbeat",
        "node-1",
        &heartbeat,
    )
    .send()
    .await
    .expect("heartbeat request should succeed");
    assert_eq!(heartbeat_response.status(), StatusCode::OK);

    let from = now_unix_secs.saturating_sub(120);
    let to = now_unix_secs.saturating_add(120);
    let metrics_response = client
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-1/metrics?from={from}&to={to}&step=1m"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("metrics request should succeed");
    assert_eq!(metrics_response.status(), StatusCode::OK);
    let metrics_payload: serde_json::Value = metrics_response
        .json()
        .await
        .expect("metrics json should parse");
    assert_eq!(metrics_payload["step"], "1m");
    assert_eq!(metrics_payload["summary"]["samples"], 2);
    assert_eq!(metrics_payload["summary"]["uptime_samples"], 2);
    assert_eq!(metrics_payload["summary"]["active_connections_max"], 7);
    assert_eq!(metrics_payload["summary"]["heartbeat_rtt_ms_sum"], 42);
    assert_eq!(metrics_payload["summary"]["heartbeat_rtt_ms_avg"], 21.0);
    assert_eq!(metrics_payload["summary"]["connect_errors_delta"], 3);
    assert_eq!(metrics_payload["summary"]["ws_out_frames_delta"], 20);
    let metric_items = metrics_payload["items"]
        .as_array()
        .expect("metrics items should be array");
    assert!(!metric_items.is_empty());
    assert!(metric_items.iter().any(|item| item["node_id"] == "node-1"));
    assert!(metric_items
        .iter()
        .all(|item| item["bucket_start"].is_string()));

    let fleet_response = client
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes/metrics/fleet?from={from}&to={to}&step=1m"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("fleet metrics request should succeed");
    assert_eq!(fleet_response.status(), StatusCode::OK);
    let fleet_payload: serde_json::Value = fleet_response
        .json()
        .await
        .expect("fleet json should parse");
    assert_eq!(fleet_payload["summary"]["samples"], 2);
    assert_eq!(fleet_payload["summary"]["error_events_delta"], 1);

    let events_response = client
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-1/events?from={from}&to={to}&event_type=tunnel_err"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("events request should succeed");
    assert_eq!(events_response.status(), StatusCode::OK);
    let events_payload: serde_json::Value = events_response
        .json()
        .await
        .expect("events json should parse");
    let event_items = events_payload["items"]
        .as_array()
        .expect("event items should be array");
    assert_eq!(event_items.len(), 1);
    assert_eq!(event_items[0]["event_type"], "tunnel_err");
    assert_eq!(
        event_items[0]["event_metadata"]["category"],
        "tcp_connect_timeout"
    );

    let invalid_metrics_response = client
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-1/metrics?from={from}&to={to}&step=5m"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("invalid metrics request should succeed");
    assert_eq!(invalid_metrics_response.status(), StatusCode::BAD_REQUEST);

    gateway_handle.abort();
}

#[tokio::test]
async fn gateway_updates_proxy_node_config_and_dispatches_upgrade_targets_locally() {
    let upstream_hits = Arc::new(Mutex::new(0usize));
    let upstream_hits_clone = Arc::clone(&upstream_hits);
    let upstream = Router::new()
        .route(
            "/api/admin/proxy-nodes/node-online/config",
            any(move |_request: Request| {
                let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
                async move {
                    *upstream_hits_inner.lock().expect("mutex should lock") += 1;
                    (StatusCode::OK, Body::from("unexpected upstream hit"))
                }
            }),
        )
        .route(
            "/api/admin/proxy-nodes/upgrade",
            any(|| async move { (StatusCode::OK, Body::from("unexpected upstream hit")) }),
        );

    let mut online_node = sample_proxy_node("node-online");
    online_node = with_tunnel_control_plane_key(online_node, TUNNEL_CONTROL_PLANE_TEST_PSK);
    online_node.status = "online".to_string();
    online_node.tunnel_connected = true;
    let mut online_node_2 = sample_proxy_node("node-zeta");
    online_node_2.name = "zeta-online".to_string();
    online_node_2.status = "online".to_string();
    online_node_2.tunnel_connected = true;
    online_node_2.remote_config = None;
    let mut offline_node = sample_proxy_node("node-offline");
    offline_node.status = "offline".to_string();
    offline_node.tunnel_connected = false;
    offline_node.remote_config = None;
    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
        online_node,
        online_node_2,
        offline_node,
    ]));

    let (upstream_url, upstream_handle) = start_server(upstream).await;
    let data_state =
        GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&proxy_node_repository))
            .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state.clone()),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;
    let client = reqwest::Client::new();

    let config_response = client
        .put(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-online/config"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({
            "node_name": "edge-online",
            "allowed_ports": [443, 8443],
            "log_level": "info",
            "heartbeat_interval": 45,
            "upgrade_to": null
        }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(config_response.status(), StatusCode::OK);
    let config_payload: serde_json::Value = config_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(config_payload["node_id"], "node-online");
    assert_eq!(config_payload["config_version"], 8);
    assert_eq!(config_payload["node"]["name"], "edge-online");
    assert_eq!(config_payload["remote_config"]["log_level"], "info");
    assert!(config_payload["remote_config"].get("upgrade_to").is_none());

    let upgrade_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({ "version": "2.0.0", "cooldown_secs": 0 }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(upgrade_response.status(), StatusCode::OK);
    let upgrade_payload: serde_json::Value = upgrade_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(upgrade_payload["version"], "2.0.0");
    assert_eq!(upgrade_payload["eligible_total"], 3);
    assert_eq!(upgrade_payload["updated"], 3);
    assert_eq!(upgrade_payload["skipped"], 0);
    assert_eq!(
        upgrade_payload["node_ids"],
        json!(["node-online", "node-offline", "node-zeta"])
    );
    assert_eq!(upgrade_payload["rollout_cancelled"], false);

    let blocked_upgrade_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({ "version": "2.0.0", "cooldown_secs": 0 }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(blocked_upgrade_response.status(), StatusCode::OK);
    let blocked_upgrade_payload: serde_json::Value = blocked_upgrade_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(blocked_upgrade_payload["updated"], 0);
    assert_eq!(blocked_upgrade_payload["skipped"], 3);

    let heartbeat = json!({
        "node_id": "node-online",
        "heartbeat_session_id": "node-online-session",
        "heartbeat_id": 77,
        "heartbeat_interval": 45,
        "active_connections": 3,
        "total_requests": 5,
        "avg_latency_ms": 10.0,
        "proxy_metadata": { "arch": "arm64" },
        "proxy_version": "2.0.0"
    });
    let heartbeat_response = authenticated_tunnel_control_plane_request(
        &client,
        format!("{gateway_url}/api/internal/tunnel/heartbeat"),
        "/api/internal/tunnel/heartbeat",
        "node-online",
        &heartbeat,
    )
    .send()
    .await
    .expect("request should succeed");
    assert_eq!(heartbeat_response.status(), StatusCode::OK);
    let heartbeat_payload: serde_json::Value = heartbeat_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(heartbeat_payload["heartbeat_id"], 77);
    assert_eq!(heartbeat_payload["config_version"], 10);
    assert!(heartbeat_payload.get("upgrade_to").is_none());
    assert_eq!(heartbeat_payload["remote_config"]["allowed_ports"][1], 8443);
    assert_eq!(heartbeat_payload["remote_config"]["log_level"], "info");

    let post_heartbeat_upgrade_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({ "version": "2.0.0", "cooldown_secs": 0 }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(post_heartbeat_upgrade_response.status(), StatusCode::OK);
    let post_heartbeat_upgrade_payload: serde_json::Value = post_heartbeat_upgrade_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(post_heartbeat_upgrade_payload["node_ids"], json!([]));
    assert_eq!(post_heartbeat_upgrade_payload["updated"], 0);
    assert_eq!(post_heartbeat_upgrade_payload["skipped"], 3);

    assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

    gateway_handle.abort();
    upstream_handle.abort();
}

#[tokio::test]
async fn gateway_still_dispatches_upgrade_targets_to_draining_proxy_nodes() {
    let mut alpha = sample_proxy_node("node-alpha");
    alpha.name = "alpha".to_string();
    alpha.status = "online".to_string();
    alpha.tunnel_connected = true;
    alpha.remote_config = None;

    let mut zeta = sample_proxy_node("node-zeta");
    zeta.name = "zeta".to_string();
    zeta.status = "online".to_string();
    zeta.tunnel_connected = true;
    zeta.remote_config = None;

    let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![zeta, alpha]));
    let data_state =
        GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&proxy_node_repository))
            .with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
    let gateway = build_router_with_state(
        AppState::new()
            .expect("gateway should build")
            .with_data_state_for_tests(data_state),
    );
    let (gateway_url, gateway_handle) = start_server(gateway).await;
    let client = reqwest::Client::new();

    let config_response = client
        .put(format!(
            "{gateway_url}/api/admin/proxy-nodes/node-alpha/config"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({
            "scheduling_state": "draining"
        }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(config_response.status(), StatusCode::OK);
    let config_payload: serde_json::Value = config_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(
        config_payload["remote_config"]["scheduling_state"],
        "draining"
    );

    let list_response = client
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(list_response.status(), StatusCode::OK);
    let list_payload: serde_json::Value =
        list_response.json().await.expect("json body should parse");
    let alpha_payload = list_payload["items"]
        .as_array()
        .expect("items should be array")
        .iter()
        .find(|item| item["id"] == "node-alpha")
        .expect("alpha should exist");
    assert_eq!(
        alpha_payload["remote_config"]["scheduling_state"],
        "draining"
    );

    let upgrade_response = client
        .post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .json(&json!({ "version": "2.0.0", "batch_size": 2, "cooldown_secs": 0 }))
        .send()
        .await
        .expect("request should succeed");
    assert_eq!(upgrade_response.status(), StatusCode::OK);
    let upgrade_payload: serde_json::Value = upgrade_response
        .json()
        .await
        .expect("json body should parse");
    assert_eq!(upgrade_payload["eligible_total"], 2);
    assert_eq!(upgrade_payload["updated"], 2);
    assert_eq!(
        upgrade_payload["node_ids"],
        json!(["node-alpha", "node-zeta"])
    );

    let list_response_after_upgrade = client
        .get(format!(
            "{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
        ))
        .header(GATEWAY_HEADER, "rust-phase3b")
        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
        .send()
        .await
        .expect("request should succeed");
    let list_payload_after_upgrade: serde_json::Value = list_response_after_upgrade
        .json()
        .await
        .expect("json body should parse");
    let alpha_after_upgrade = list_payload_after_upgrade["items"]
        .as_array()
        .expect("items should be array")
        .iter()
        .find(|item| item["id"] == "node-alpha")
        .expect("alpha should exist after upgrade");
    assert_eq!(alpha_after_upgrade["remote_config"]["upgrade_to"], "2.0.0");

    gateway_handle.abort();
}