datum-agent 0.11.2

use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::{Duration, Instant},
};

use datum::{Channel, Keep, Sink, Source};
use datum_agent::{
    Agent, AgentConfig, JobMat, JobSpec,
    dcp::{
        ClientKind, ClusterJobList, ClusterJobStart, ClusterNodeList, ClusterViewProvider,
        DcpClient, DcpJobFactories, DcpQuicServerConfig, DcpServer, DcpServerConfig,
        DcpTcpServerConfig, Hello, RememberClusterAssignment, ResponseStatus, SubmitClusterJob,
        SubscribeMetrics,
        proto::{
            DcpFrame, JobStatus, MetricSample, NodeMetric, Request, StreamMetric, dcp_frame,
            request,
        },
        server::ClusterViewFuture,
    },
};
use datum_net::quic::{
    crypto::rustls::{QuicClientConfig, QuicServerConfig},
    quinn,
    rustls::{
        ClientConfig as RustlsClientConfig, RootCertStore, ServerConfig as RustlsServerConfig,
        pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer},
        server::WebPkiClientVerifier,
    },
};
use prost::Message as ProstMessage;
use rcgen::{CertifiedKey, generate_simple_self_signed};
use tokio::{
    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    sync::mpsc,
};

fn agent() -> datum_agent::AgentHandle {
    Agent::start_with_config(AgentConfig {
        poll_interval: Duration::from_millis(5),
        event_buffer: 128,
        ..AgentConfig::default()
    })
    .expect("agent starts")
}

fn never_job(instance_name: String) -> JobSpec {
    JobSpec::new(instance_name, |context| {
        let control = context.control();
        Ok(
            Source::tick(Duration::ZERO, Duration::from_millis(10), 1_u64)
                .instrumented(
                    format!("{}:{}", context.name(), context.generation()),
                    context.instrumentation_registry(),
                )
                .via_mat(context.drain_flow(), Keep::right)
                .to_mat(Sink::ignore(), move |_switch, completion| {
                    JobMat::new(completion, control.clone())
                }),
        )
    })
}

fn never_factories() -> DcpJobFactories {
    let factories = DcpJobFactories::new();
    factories
        .register("never", |instance_name, _params| {
            Ok(never_job(instance_name))
        })
        .expect("factory registers");
    factories
}

async fn start_tcp_server(
    factories: DcpJobFactories,
    metrics_interval: Duration,
) -> (
    datum_agent::AgentHandle,
    datum_agent::dcp::DcpServerHandle,
    SocketAddr,
) {
    let agent = agent();
    let server = DcpServer::from_agent(
        &agent,
        factories,
        DcpServerConfig {
            node_id: "test-agent".to_owned(),
            tcp: Some(DcpTcpServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
            }),
            quic: None,
            auth_token: None,
            metrics_interval,
            frame_buffer: 128,
        },
    );
    let handle = server.start().await.expect("DCP server starts");
    let addr = handle.tcp_addr().expect("TCP address");
    (agent, handle, addr)
}

async fn connect_tcp(addr: SocketAddr) -> DcpClient {
    DcpClient::connect_tcp(addr, Hello::new("test-client", ClientKind::Cli))
        .await
        .expect("DCP client connects")
}

async fn wait_status(
    client: &DcpClient,
    name: &str,
    predicate: impl Fn(&datum_agent::dcp::proto::JobStatus) -> bool,
) -> datum_agent::dcp::proto::JobStatus {
    wait_until(Duration::from_secs(3), || async {
        let status = client.job_status(name).await.expect("job status");
        predicate(&status).then_some(status)
    })
    .await
}

async fn wait_until<T, Fut>(timeout: Duration, mut condition: impl FnMut() -> Fut) -> T
where
    Fut: std::future::Future<Output = Option<T>>,
{
    let deadline = Instant::now() + timeout;
    loop {
        if let Some(value) = condition().await {
            return value;
        }
        assert!(Instant::now() < deadline, "condition timed out");
        tokio::time::sleep(Duration::from_millis(5)).await;
    }
}

#[derive(Clone)]
struct StaticClusterMetrics {
    sample: MetricSample,
}

impl StaticClusterMetrics {
    fn new(sample: MetricSample) -> Self {
        Self { sample }
    }
}

impl ClusterViewProvider for StaticClusterMetrics {
    fn subscribe_cluster_metrics(
        &self,
        _interval: Duration,
    ) -> Option<mpsc::Receiver<MetricSample>> {
        let sample = self.sample.clone();
        let (sender, receiver) = mpsc::channel(4);
        tokio::spawn(async move {
            loop {
                if sender.send(sample.clone()).await.is_err() {
                    break;
                }
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
        });
        Some(receiver)
    }

    fn submit_cluster_job(
        &self,
        _request: SubmitClusterJob,
        _timeout: Duration,
    ) -> ClusterViewFuture<'_, JobStatus> {
        unavailable_cluster_view()
    }

    fn list_cluster_jobs(&self, _timeout: Duration) -> ClusterViewFuture<'_, ClusterJobList> {
        unavailable_cluster_view()
    }

    fn cluster_node_info(&self, _timeout: Duration) -> ClusterViewFuture<'_, ClusterNodeList> {
        unavailable_cluster_view()
    }

    fn cluster_job_status(
        &self,
        _name: String,
        _timeout: Duration,
    ) -> ClusterViewFuture<'_, JobStatus> {
        unavailable_cluster_view()
    }

    fn drain_cluster_job(
        &self,
        _name: String,
        _timeout: Duration,
    ) -> ClusterViewFuture<'_, JobStatus> {
        unavailable_cluster_view()
    }

    fn stop_cluster_job(
        &self,
        _name: String,
        _timeout: Duration,
    ) -> ClusterViewFuture<'_, JobStatus> {
        unavailable_cluster_view()
    }

    fn remember_cluster_assignment(
        &self,
        _request: RememberClusterAssignment,
    ) -> ClusterViewFuture<'_, ()> {
        unavailable_cluster_view()
    }

    fn register_restarted_cluster_assignment(
        &self,
        _instance_name: String,
        _assignment: ClusterJobStart,
        _timeout: Duration,
    ) -> ClusterViewFuture<'_, ()> {
        unavailable_cluster_view()
    }
}

fn unavailable_cluster_view<T: Send + 'static>() -> ClusterViewFuture<'static, T> {
    Box::pin(async {
        Err(datum_agent::dcp::DcpError::Response {
            status: ResponseStatus::Failed,
            message: "cluster view unavailable in test".to_owned(),
        })
    })
}

#[tokio::test]
async fn hello_version_negotiation_accepts_and_rejects_unknown_major() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;

    let mut bad_hello = Hello::new("bad-client", ClientKind::Cli);
    bad_hello.protocol_version = "1.0.0".to_owned();
    let rejected = match DcpClient::connect_tcp(addr, bad_hello).await {
        Ok(_) => panic!("unknown major is rejected"),
        Err(error) => error,
    };
    assert!(
        matches!(
            rejected,
            datum_agent::dcp::DcpError::Response {
                status: ResponseStatus::ProtocolMismatch,
                ..
            }
        ),
        "unexpected rejection: {rejected:?}"
    );

    let client = connect_tcp(addr).await;
    assert!(client.list_jobs().await.expect("list jobs").is_empty());

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn commands_cover_start_list_status_drain_stop_and_restart() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let client = connect_tcp(addr).await;

    let started = client
        .start_job("never", "cmd-job", HashMap::new())
        .await
        .expect("job starts");
    assert_eq!(started.name, "cmd-job");
    assert_eq!(started.state, "Running");

    let listed = client.list_jobs().await.expect("jobs list");
    assert_eq!(listed.len(), 1);
    assert_eq!(listed[0].name, "cmd-job");

    let status = client.job_status("cmd-job").await.expect("status");
    assert_eq!(status.state, "Running");

    let restarted = client.restart_job("cmd-job").await.expect("restart");
    assert_eq!(restarted.state, "Running");
    assert_eq!(restarted.generation, 2);

    client.drain_job("cmd-job").await.expect("drain starts");
    let drained = wait_status(&client, "cmd-job", |status| status.state == "Drained").await;
    assert_eq!(drained.state, "Drained");

    let stopped_job = client
        .start_job("never", "stop-job", HashMap::new())
        .await
        .expect("second job starts");
    assert_eq!(stopped_job.state, "Running");
    let stopped = client.stop_job("stop-job").await.expect("stop");
    assert_eq!(stopped.state, "Stopped");

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn subscribe_events_preserves_lifecycle_ordering() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let client = connect_tcp(addr).await;
    let mut events = client.subscribe_events().await.expect("events subscribe");

    client
        .start_job("never", "event-job", HashMap::new())
        .await
        .expect("job starts");

    let submitted = tokio::time::timeout(Duration::from_secs(2), events.recv())
        .await
        .expect("submitted timeout")
        .expect("submitted event");
    let started = tokio::time::timeout(Duration::from_secs(2), events.recv())
        .await
        .expect("started timeout")
        .expect("started event");

    assert_eq!(submitted.kind, "Submitted");
    assert_eq!(started.kind, "Started");
    assert_eq!(submitted.name, "event-job");
    assert_eq!(started.name, "event-job");
    assert!(submitted.sequence < started.sequence);

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn subscribe_metrics_delivers_samples_while_job_runs() {
    let channel = Channel::bounded(16);
    let producer = channel.clone();
    let factories = DcpJobFactories::new();
    factories
        .register("metered", move |instance_name, _params| {
            let channel = channel.clone();
            Ok(JobSpec::new(instance_name, move |context| {
                let control = context.control();
                Ok(channel
                    .source()
                    .instrumented(
                        format!("{}:{}", context.name(), context.generation()),
                        context.instrumentation_registry(),
                    )
                    .via_mat(context.drain_flow(), Keep::right)
                    .to_mat(Sink::ignore(), move |_switch, completion| {
                        JobMat::new(completion, control.clone())
                    }))
            }))
        })
        .expect("factory registers");

    let (agent, handle, addr) = start_tcp_server(factories, Duration::from_millis(10)).await;
    let client = connect_tcp(addr).await;
    let mut metrics = client
        .subscribe_metrics(10)
        .await
        .expect("metrics subscribe");

    client
        .start_job("metered", "metrics-job", HashMap::new())
        .await
        .expect("job starts");
    for value in 0_u64..4 {
        producer.try_send(value).expect("metric item sends");
    }

    let deadline = Instant::now() + Duration::from_secs(3);
    let sample =
        loop {
            let sample = tokio::time::timeout(Duration::from_millis(200), metrics.recv())
                .await
                .ok()
                .flatten()
                .expect("metrics stream remains open");
            if sample.streams.iter().any(|metric| {
                metric.name.starts_with("metrics-job:") && metric.elements_through >= 4
            }) {
                break sample;
            }
            assert!(Instant::now() < deadline, "timed out waiting for metrics");
        };
    assert!(!sample.streams.is_empty());
    let node = sample
        .nodes
        .iter()
        .find(|metric| metric.node_id == "test-agent")
        .expect("local node telemetry");
    assert!(node.sample_ts_ms > 0);
    #[cfg(target_os = "linux")]
    {
        assert!(node.rss_bytes > 0);
        assert!(node.mem_total_bytes.is_some_and(|total| total > 0));
        assert!(node.threads > 0);
    }

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn subscribe_metrics_filters_remote_streams_at_dcp_boundary() {
    let agent = agent();
    let server = DcpServer::from_agent(
        &agent,
        never_factories(),
        DcpServerConfig {
            node_id: "filter-node".to_owned(),
            tcp: Some(DcpTcpServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
            }),
            quic: None,
            auth_token: None,
            metrics_interval: Duration::from_millis(5),
            frame_buffer: 128,
        },
    );
    server.set_cluster_view(Arc::new(StaticClusterMetrics::new(MetricSample {
        timestamp_ms: 10,
        streams: vec![
            stream_metric(1, "target-job:source", 11),
            stream_metric(2, "target-job:sink", 11),
            stream_metric(3, "unrelated-remote-local-only:source", 99),
        ],
        nodes: vec![node_metric("remote-node")],
    })));
    let handle = server.start().await.expect("DCP server starts");
    let addr = handle.tcp_addr().expect("TCP address");

    let named = raw_metrics_sample(addr, vec!["target-job".to_owned()], |sample| {
        !sample.streams.is_empty()
    })
    .await;
    assert_eq!(
        stream_names(&named),
        vec!["target-job:sink".to_owned(), "target-job:source".to_owned()]
    );

    let unfiltered = raw_metrics_sample(addr, Vec::new(), |sample| sample.streams.len() >= 3).await;
    assert_eq!(
        stream_names(&unfiltered),
        vec![
            "target-job:sink".to_owned(),
            "target-job:source".to_owned(),
            "unrelated-remote-local-only:source".to_owned(),
        ]
    );

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn metrics_keep_recent_drain_then_prune_stale_stream_after_restart() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(10)).await;
    let client = connect_tcp(addr).await;
    let mut metrics = client
        .subscribe_metrics(10)
        .await
        .expect("metrics subscribe");

    client
        .start_job("never", "restart-metrics", HashMap::new())
        .await
        .expect("job starts");
    let first_deadline = Instant::now() + Duration::from_secs(3);
    loop {
        let sample = tokio::time::timeout(Duration::from_millis(200), metrics.recv())
            .await
            .ok()
            .flatten()
            .expect("metrics stream remains open");
        if sample
            .streams
            .iter()
            .any(|metric| metric.name == "restart-metrics:1")
        {
            break;
        }
        assert!(Instant::now() < first_deadline, "first metric timed out");
    }

    client
        .restart_job("restart-metrics")
        .await
        .expect("job restarts");

    let recent_drain_deadline = Instant::now() + Duration::from_secs(3);
    let recent_drain = loop {
        let sample = tokio::time::timeout(Duration::from_millis(200), metrics.recv())
            .await
            .ok()
            .flatten()
            .expect("metrics stream remains open");
        let new_running = sample
            .streams
            .iter()
            .any(|metric| metric.name == "restart-metrics:2" && metric.state == "Running");
        let old_draining = sample
            .streams
            .iter()
            .any(|metric| metric.name == "restart-metrics:1" && metric.state == "Draining");
        if new_running && old_draining {
            break sample;
        }
        assert!(
            Instant::now() < recent_drain_deadline,
            "recent draining metric timed out"
        );
    };
    assert_eq!(
        recent_drain
            .streams
            .iter()
            .filter(|metric| metric.name.starts_with("restart-metrics:"))
            .map(|metric| metric.name.as_str())
            .collect::<Vec<_>>(),
        vec!["restart-metrics:1", "restart-metrics:2"]
    );

    let stale_deadline = Instant::now() + Duration::from_secs(6);
    let sample = loop {
        let sample = tokio::time::timeout(Duration::from_millis(200), metrics.recv())
            .await
            .ok()
            .flatten()
            .expect("metrics stream remains open");
        let new_running = sample
            .streams
            .iter()
            .any(|metric| metric.name == "restart-metrics:2");
        let old_retained = sample
            .streams
            .iter()
            .any(|metric| metric.name == "restart-metrics:1");
        if new_running && !old_retained {
            break sample;
        }
        assert!(
            Instant::now() < stale_deadline,
            "stale terminal metric was not pruned after its grace period"
        );
    };
    assert_eq!(
        sample
            .streams
            .iter()
            .filter(|metric| metric.name.starts_with("restart-metrics:"))
            .map(|metric| metric.name.as_str())
            .collect::<Vec<_>>(),
        vec!["restart-metrics:2"]
    );

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn loopback_tcp_mode_works_and_non_loopback_plaintext_is_refused() {
    let (agent_handle, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let client = connect_tcp(addr).await;
    assert!(client.list_jobs().await.expect("list").is_empty());
    handle.shutdown().await;
    agent_handle
        .registry()
        .shutdown()
        .expect("registry shuts down");

    let agent = agent();
    let server = DcpServer::from_agent(
        &agent,
        never_factories(),
        DcpServerConfig {
            tcp: Some(DcpTcpServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
            }),
            quic: None,
            ..DcpServerConfig::default()
        },
    );
    assert!(server.start().await.is_err());
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn malformed_frame_closes_connection_without_panicking() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let mut stream = TcpStream::connect(addr).await.expect("raw client connects");
    stream
        .write_all(&0_u32.to_be_bytes())
        .await
        .expect("malformed frame writes");
    let mut header = [0_u8; 4];
    let _ = tokio::time::timeout(Duration::from_secs(1), stream.read_exact(&mut header)).await;

    let client = connect_tcp(addr).await;
    assert!(
        client
            .list_jobs()
            .await
            .expect("server survived")
            .is_empty()
    );

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn concurrent_clients_can_issue_requests() {
    let (agent, handle, addr) =
        start_tcp_server(never_factories(), Duration::from_millis(20)).await;
    let mut tasks = Vec::new();
    for index in 0..8 {
        tasks.push(tokio::spawn(async move {
            let client = DcpClient::connect_tcp(
                addr,
                Hello::new(format!("client-{index}"), ClientKind::Cli),
            )
            .await
            .expect("client connects");
            client.list_jobs().await.expect("list jobs")
        }));
    }

    for task in tasks {
        assert!(task.await.expect("client task joins").is_empty());
    }

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

#[tokio::test]
async fn quic_mtls_handshake_failure_surfaces_cleanly() {
    let (server_config, client_config, no_cert_client_config) = mtls_quic_configs();
    let agent = agent();
    let server = DcpServer::from_agent(
        &agent,
        never_factories(),
        DcpServerConfig {
            node_id: "quic-agent".to_owned(),
            tcp: None,
            quic: Some(DcpQuicServerConfig {
                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
                server_config,
            }),
            auth_token: None,
            metrics_interval: Duration::from_millis(20),
            frame_buffer: 128,
        },
    );
    let handle = server.start().await.expect("QUIC server starts");
    let addr = handle.quic_addr().expect("QUIC addr");

    let failed = DcpClient::connect_quic(
        addr,
        "localhost",
        no_cert_client_config,
        Hello::new("bad-quic", ClientKind::Cli),
    )
    .await;
    assert!(failed.is_err(), "client without cert should fail");

    let client = DcpClient::connect_quic(
        addr,
        "localhost",
        client_config,
        Hello::new("good-quic", ClientKind::Cli),
    )
    .await
    .expect("mTLS client connects");
    assert!(client.list_jobs().await.expect("list over QUIC").is_empty());

    handle.shutdown().await;
    agent.registry().shutdown().expect("registry shuts down");
}

fn mtls_quic_configs() -> (
    quinn::ServerConfig,
    quinn::ClientConfig,
    quinn::ClientConfig,
) {
    let CertifiedKey {
        cert: server_cert,
        key_pair: server_key,
    } = generate_simple_self_signed(["localhost".to_owned()]).expect("server cert");
    let CertifiedKey {
        cert: client_cert,
        key_pair: client_key,
    } = generate_simple_self_signed(["datum-client".to_owned()]).expect("client cert");

    let mut client_roots = RootCertStore::empty();
    client_roots
        .add(client_cert.der().clone())
        .expect("trust client cert");
    let verifier = WebPkiClientVerifier::builder(Arc::new(client_roots))
        .build()
        .expect("client verifier");
    let server_crypto = RustlsServerConfig::builder()
        .with_client_cert_verifier(verifier)
        .with_single_cert(
            vec![server_cert.der().clone()],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(server_key.serialize_der())),
        )
        .expect("server crypto");

    let mut server_roots = RootCertStore::empty();
    server_roots
        .add(server_cert.der().clone())
        .expect("trust server cert");
    let client_crypto = RustlsClientConfig::builder()
        .with_root_certificates(server_roots.clone())
        .with_client_auth_cert(
            vec![client_cert.der().clone()],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(client_key.serialize_der())),
        )
        .expect("client crypto");
    let no_cert_client_crypto = RustlsClientConfig::builder()
        .with_root_certificates(server_roots)
        .with_no_client_auth();

    (
        quinn::ServerConfig::with_crypto(Arc::new(
            QuicServerConfig::try_from(server_crypto).expect("QUIC server config"),
        )),
        quinn::ClientConfig::new(Arc::new(
            QuicClientConfig::try_from(client_crypto).expect("QUIC client config"),
        )),
        quinn::ClientConfig::new(Arc::new(
            QuicClientConfig::try_from(no_cert_client_crypto).expect("QUIC client config"),
        )),
    )
}

async fn recv_raw_frame<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Option<DcpFrame>> {
    let mut header = [0_u8; 4];
    match reader.read_exact(&mut header).await {
        Ok(_) => {}
        Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
        Err(error) => return Err(error),
    }
    let len = u32::from_be_bytes(header) as usize;
    if len == 0 || len > 16 * 1024 * 1024 {
        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            "bad frame length",
        ));
    }
    let mut payload = vec![0_u8; len];
    reader.read_exact(&mut payload).await?;
    Ok(Some(DcpFrame::decode(payload.as_slice()).map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, e)
    })?))
}

async fn send_raw_frame<W: AsyncWrite + Unpin>(
    writer: &mut W,
    frame: &DcpFrame,
) -> std::io::Result<()> {
    let payload = frame.encode_to_vec();
    let len = u32::try_from(payload.len())
        .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "frame too large"))?;
    writer.write_all(&len.to_be_bytes()).await?;
    writer.write_all(&payload).await?;
    writer.flush().await
}

async fn raw_metrics_sample(
    addr: SocketAddr,
    job_names: Vec<String>,
    predicate: impl Fn(&MetricSample) -> bool,
) -> MetricSample {
    let mut stream = TcpStream::connect(addr).await.expect("raw client connects");
    send_raw_frame(
        &mut stream,
        &DcpFrame::hello(Hello::new("raw-metrics-client", ClientKind::Cli)),
    )
    .await
    .expect("write hello");
    let hello = recv_raw_frame(&mut stream)
        .await
        .expect("read hello response")
        .expect("hello response");
    assert_ok_response(hello, 0);

    let request_id = 1;
    send_raw_frame(
        &mut stream,
        &DcpFrame::request(Request {
            request_id,
            deadline_ms: 1_000,
            command: Some(request::Command::SubscribeMetrics(SubscribeMetrics {
                interval_ms: 5,
                job_names,
                local_only: false,
            })),
        }),
    )
    .await
    .expect("write metrics subscribe");
    let response = recv_raw_frame(&mut stream)
        .await
        .expect("read subscribe response")
        .expect("subscribe response");
    assert_ok_response(response, request_id);

    let deadline = Instant::now() + Duration::from_secs(3);
    loop {
        match tokio::time::timeout(Duration::from_millis(200), recv_raw_frame(&mut stream)).await {
            Ok(Ok(Some(frame))) => {
                if let Some(dcp_frame::Frame::Metric(metric)) = frame.frame
                    && let Some(sample) = metric.sample
                    && predicate(&sample)
                {
                    return sample;
                }
            }
            Ok(Ok(None)) => panic!("metrics stream closed"),
            Ok(Err(error)) => panic!("metrics stream read failed: {error}"),
            Err(_) => {}
        }
        assert!(
            Instant::now() < deadline,
            "matching metric sample timed out"
        );
    }
}

fn assert_ok_response(frame: DcpFrame, request_id: u64) {
    match frame.frame {
        Some(dcp_frame::Frame::Response(response)) => {
            assert_eq!(response.request_id, request_id);
            assert_eq!(response.status, ResponseStatus::Ok as i32);
        }
        _ => panic!("expected response frame"),
    }
}

fn stream_metric(id: u64, name: &str, elements_through: u64) -> StreamMetric {
    StreamMetric {
        id,
        name: name.to_owned(),
        elements_through,
        restarts: 0,
        state: "Running".to_owned(),
        started_at_ms: 1,
        state_changed_at_ms: 1,
        finished_at_ms: None,
        uptime_ms: 1,
    }
}

fn node_metric(node_id: &str) -> NodeMetric {
    NodeMetric {
        node_id: node_id.to_owned(),
        cpu_percent: 0.0,
        cpu_ms: 0,
        rss_bytes: 1,
        mem_total_bytes: Some(1),
        threads: 1,
        sample_ts_ms: 1,
    }
}

fn stream_names(sample: &MetricSample) -> Vec<String> {
    let mut names = sample
        .streams
        .iter()
        .map(|metric| metric.name.clone())
        .collect::<Vec<_>>();
    names.sort();
    names
}

#[tokio::test]
async fn event_subscription_terminates_on_server_disconnect() {
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("bind listener");
    let addr = listener.local_addr().expect("listener addr");

    let server_task = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.expect("accept");
        let hello = recv_raw_frame(&mut stream)
            .await
            .expect("read hello")
            .expect("hello frame");
        assert!(matches!(
            hello.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Hello(_))
        ));
        let hello_resp = datum_agent::dcp::proto::Response::ok(0, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(hello_resp))
            .await
            .expect("write hello response");

        let subscribe = recv_raw_frame(&mut stream)
            .await
            .expect("read subscribe")
            .expect("subscribe frame");
        assert!(matches!(
            subscribe.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Request(_))
        ));

        let subscribe_resp = datum_agent::dcp::proto::Response::ok(1, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(subscribe_resp))
            .await
            .expect("write subscribe response");

        drop(stream);
    });

    let client = DcpClient::connect_tcp(addr, Hello::new("test-client", ClientKind::Cli))
        .await
        .expect("client connects");

    let mut events = client.subscribe_events().await.expect("events subscribe");

    server_task.await.expect("server task");

    let deadline = Instant::now() + Duration::from_secs(5);
    loop {
        match tokio::time::timeout(Duration::from_millis(100), events.recv()).await {
            Ok(None) => break,
            Ok(Some(_)) => {
                if Instant::now() > deadline {
                    panic!("subscription still receiving events after server disconnect");
                }
            }
            Err(_elapsed) => {
                if Instant::now() > deadline {
                    panic!("subscription recv() never completed after server disconnect");
                }
            }
        }
    }
    assert!(
        events.disconnect_reason().is_some(),
        "disconnect reason should be set after server drops connection"
    );
}

#[tokio::test]
async fn metric_subscription_terminates_on_server_disconnect() {
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("bind listener");
    let addr = listener.local_addr().expect("listener addr");

    let server_task = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.expect("accept");
        let hello = recv_raw_frame(&mut stream)
            .await
            .expect("read hello")
            .expect("hello frame");
        assert!(matches!(
            hello.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Hello(_))
        ));
        let hello_resp = datum_agent::dcp::proto::Response::ok(0, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(hello_resp))
            .await
            .expect("write hello response");

        let subscribe = recv_raw_frame(&mut stream)
            .await
            .expect("read subscribe")
            .expect("subscribe frame");
        assert!(matches!(
            subscribe.frame,
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Request(_))
        ));

        let subscribe_resp = datum_agent::dcp::proto::Response::ok(1, Vec::new());
        send_raw_frame(&mut stream, &DcpFrame::response(subscribe_resp))
            .await
            .expect("write subscribe response");

        drop(stream);
    });

    let client = DcpClient::connect_tcp(addr, Hello::new("test-client", ClientKind::Cli))
        .await
        .expect("client connects");

    let mut metrics = client
        .subscribe_metrics(50)
        .await
        .expect("metrics subscribe");

    server_task.await.expect("server task");

    let deadline = Instant::now() + Duration::from_secs(5);
    loop {
        match tokio::time::timeout(Duration::from_millis(100), metrics.recv()).await {
            Ok(None) => break,
            Ok(Some(_)) => {
                if Instant::now() > deadline {
                    panic!("subscription still receiving metrics after server disconnect");
                }
            }
            Err(_elapsed) => {
                if Instant::now() > deadline {
                    panic!("subscription recv() never completed after server disconnect");
                }
            }
        }
    }
    assert!(
        metrics.disconnect_reason().is_some(),
        "disconnect reason should be set after server drops connection"
    );
}

#[tokio::test]
async fn metric_subscription_cancel_notifies_server() {
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("bind listener");
    let addr = listener.local_addr().expect("listener addr");

    let server_task = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.expect("accept");
        recv_raw_frame(&mut stream)
            .await
            .expect("read hello")
            .expect("hello frame");
        send_raw_frame(
            &mut stream,
            &DcpFrame::response(datum_agent::dcp::proto::Response::ok(0, Vec::new())),
        )
        .await
        .expect("write hello response");

        let subscribe = recv_raw_frame(&mut stream)
            .await
            .expect("read subscribe")
            .expect("subscribe frame");
        let subscription_id = match subscribe.frame {
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Request(request)) => {
                assert!(matches!(
                    request.command,
                    Some(datum_agent::dcp::proto::request::Command::SubscribeMetrics(
                        _
                    ))
                ));
                request.request_id
            }
            _ => panic!("expected metrics subscribe request"),
        };
        send_raw_frame(
            &mut stream,
            &DcpFrame::response(datum_agent::dcp::proto::Response::ok(
                subscription_id,
                Vec::new(),
            )),
        )
        .await
        .expect("write subscribe response");

        let unsubscribe = recv_raw_frame(&mut stream)
            .await
            .expect("read unsubscribe")
            .expect("unsubscribe frame");
        match unsubscribe.frame {
            Some(datum_agent::dcp::proto::dcp_frame::Frame::Request(request)) => {
                match request.command {
                    Some(datum_agent::dcp::proto::request::Command::UnsubscribeMetrics(
                        unsubscribe,
                    )) => assert_eq!(unsubscribe.subscription_id, subscription_id),
                    _ => panic!("expected metrics unsubscribe command"),
                }
            }
            _ => panic!("expected metrics unsubscribe request"),
        }
    });

    let client = DcpClient::connect_tcp(addr, Hello::new("test-client", ClientKind::Cli))
        .await
        .expect("client connects");
    client
        .subscribe_metrics(50)
        .await
        .expect("metrics subscribe")
        .cancel()
        .await
        .expect("metrics cancel");

    tokio::time::timeout(Duration::from_secs(1), server_task)
        .await
        .expect("unsubscribe reaches server")
        .expect("server task");
}