Publish harness and TUI open-source

initial sync from the monorepo

crates/common/xai-computer-hub-sdk/src/metric_donate.rs

@@ -0,0 +1,396 @@
+//! Forward the process's Prometheus metrics to the connected server over
+//! the WebSocket transport (`metrics.donate`).
+//!
+//! A [`MetricDonationReporter`] periodically snapshots the default
+//! Prometheus registry via [`prometheus::gather`], converts the
+//! `MetricFamily` set to native OTLP metrics (Counter→Sum, Gauge→Gauge,
+//! Histogram→Histogram, labels preserved, cumulative temporality), and
+//! pumps the batch over the shared [`crate::donate_pump`]. Because it
+//! gathers the whole registry, every current and future metric is
+//! exported with zero per-metric wiring. Metrics are **process-aggregate**
+//! — [`ToolServer::donate_metrics`] requires no bound session.
+
+use std::sync::{Arc, LazyLock};
+use std::time::Duration;
+
+use arc_swap::ArcSwapOption;
+use base64::Engine as _;
+use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest;
+use opentelemetry_proto::tonic::common::v1::KeyValue;
+use opentelemetry_proto::tonic::metrics::v1::{
+    AggregationTemporality, Gauge, Histogram, HistogramDataPoint, Metric, NumberDataPoint,
+    ResourceMetrics, ScopeMetrics, Sum, metric, number_data_point,
+};
+use opentelemetry_proto::tonic::resource::v1::Resource;
+use prometheus::proto::{MetricFamily, MetricType};
+use prost::Message as _;
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+use xai_tool_protocol::{MAX_DONATION_BYTES, MAX_METRICS_PER_DONATION};
+
+use crate::donate_pump::{
+    PENDING_FLUSHES, PumpMsg, make_resource, now_unix_nanos, run_pump, string_kv,
+};
+use crate::server::ToolServer;
+
+/// How often the reporter snapshots the registry. The server re-stamps
+/// attribution; cumulative temporality means missed ticks only delay
+/// freshness, never lose monotonic state.
+const DEFAULT_GATHER_INTERVAL: Duration = Duration::from_secs(60);
+
+fn labels_to_kv(labels: &[prometheus::proto::LabelPair]) -> Vec<KeyValue> {
+    labels
+        .iter()
+        .map(|l| string_kv(l.name(), l.value().to_owned()))
+        .collect()
+}
+
+fn number_point(metric: &prometheus::proto::Metric, value: f64, now: u64) -> NumberDataPoint {
+    NumberDataPoint {
+        attributes: labels_to_kv(metric.get_label()),
+        time_unix_nano: now,
+        value: Some(number_data_point::Value::AsDouble(value)),
+        ..Default::default()
+    }
+}
+
+/// Prometheus histogram buckets are **cumulative** (`le` counts); OTLP
+/// wants per-bucket counts plus an implicit `+Inf` bucket, so the
+/// cumulative counts are differenced here.
+fn histogram_point(metric: &prometheus::proto::Metric, now: u64) -> HistogramDataPoint {
+    let hist = metric.get_histogram();
+    let mut bucket_counts = Vec::new();
+    let mut explicit_bounds = Vec::new();
+    let mut prev = 0u64;
+    for bucket in hist.get_bucket() {
+        let cumulative = bucket.cumulative_count();
+        bucket_counts.push(cumulative.saturating_sub(prev));
+        explicit_bounds.push(bucket.upper_bound());
+        prev = cumulative;
+    }
+    let total = hist.get_sample_count();
+    bucket_counts.push(total.saturating_sub(prev));
+    HistogramDataPoint {
+        attributes: labels_to_kv(metric.get_label()),
+        time_unix_nano: now,
+        count: total,
+        sum: Some(hist.get_sample_sum()),
+        bucket_counts,
+        explicit_bounds,
+        ..Default::default()
+    }
+}
+
+/// Convert a gathered `MetricFamily` set to OTLP metrics. Summaries and
+/// untyped families are skipped defensively (none registered today).
+fn convert_families(families: &[MetricFamily]) -> Vec<Metric> {
+    let now = now_unix_nanos();
+    let cumulative = AggregationTemporality::Cumulative as i32;
+    let mut out = Vec::new();
+    for family in families {
+        let name = family.name().to_owned();
+        let data = match family.get_field_type() {
+            MetricType::COUNTER => metric::Data::Sum(Sum {
+                data_points: family
+                    .get_metric()
+                    .iter()
+                    .map(|m| number_point(m, m.get_counter().value(), now))
+                    .collect(),
+                aggregation_temporality: cumulative,
+                is_monotonic: true,
+            }),
+            MetricType::GAUGE => metric::Data::Gauge(Gauge {
+                data_points: family
+                    .get_metric()
+                    .iter()
+                    .map(|m| number_point(m, m.get_gauge().value(), now))
+                    .collect(),
+            }),
+            MetricType::HISTOGRAM => metric::Data::Histogram(Histogram {
+                data_points: family
+                    .get_metric()
+                    .iter()
+                    .map(|m| histogram_point(m, now))
+                    .collect(),
+                aggregation_temporality: cumulative,
+            }),
+            MetricType::SUMMARY | MetricType::UNTYPED => continue,
+        };
+        out.push(Metric {
+            name,
+            data: Some(data),
+            ..Default::default()
+        });
+    }
+    out
+}
+
+/// Encodes batches of OTLP metrics onto the pump channel. Chunks at
+/// [`MAX_METRICS_PER_DONATION`], drops payloads over
+/// [`MAX_DONATION_BYTES`], and never blocks.
+#[derive(Clone)]
+struct MetricExporter {
+    tx: mpsc::Sender<PumpMsg>,
+    resource: Resource,
+}
+
+impl MetricExporter {
+    fn export(&self, mut metrics: Vec<Metric>) {
+        while !metrics.is_empty() {
+            let chunk = if metrics.len() > MAX_METRICS_PER_DONATION {
+                let rest = metrics.split_off(MAX_METRICS_PER_DONATION);
+                std::mem::replace(&mut metrics, rest)
+            } else {
+                std::mem::take(&mut metrics)
+            };
+            let request = ExportMetricsServiceRequest {
+                resource_metrics: vec![ResourceMetrics {
+                    resource: Some(self.resource.clone()),
+                    scope_metrics: vec![ScopeMetrics {
+                        metrics: chunk,
+                        ..Default::default()
+                    }],
+                    schema_url: String::new(),
+                }],
+            };
+            let bytes = request.encode_to_vec();
+            if bytes.len() > MAX_DONATION_BYTES {
+                tracing::debug!(
+                    len = bytes.len(),
+                    "dropping oversized metric donation payload"
+                );
+                continue;
+            }
+            let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
+            if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
+                tracing::debug!("metric donation queue full; dropping metric batch");
+            }
+        }
+    }
+
+    fn gather_and_send(&self) {
+        let metrics = convert_families(&prometheus::gather());
+        if metrics.is_empty() {
+            return;
+        }
+        self.export(metrics);
+    }
+}
+
+/// Process-global handle to the active exporter so [`gather_and_send`]
+/// can drive a final teardown gather without a reference.
+static ACTIVE_METRIC_EXPORTER: LazyLock<ArcSwapOption<MetricExporter>> =
+    LazyLock::new(ArcSwapOption::empty);
+
+/// Final registry gather onto the active metric pump. Called from
+/// `ToolServer` teardown before the pump drain so a crash-y shutdown
+/// captures the latest values.
+pub(crate) fn gather_and_send() {
+    if let Some(exporter) = ACTIVE_METRIC_EXPORTER.load_full() {
+        exporter.gather_and_send();
+    }
+}
+
+/// Drop the process-global exporter on teardown so its pump `Sender` is released
+/// and the metric pump can wind down. Called from `flush_donations_inner` after
+/// the final [`gather_and_send`] (and alongside clearing the stored pump
+/// senders), so a dropped `ToolServer` doesn't leak the pump task.
+pub(crate) fn clear_active_exporter() {
+    ACTIVE_METRIC_EXPORTER.store(None);
+}
+
+/// Periodic registry gatherer spawned by
+/// [`ToolServer::metric_donation_reporter`]. Internal: constructed and run
+/// only by `metric_donation_reporter`; not part of the crate's public API.
+pub(crate) struct MetricDonationReporter {
+    exporter: MetricExporter,
+    interval: Duration,
+    shutdown: CancellationToken,
+}
+
+impl MetricDonationReporter {
+    async fn run(self) {
+        let mut ticker = tokio::time::interval(self.interval);
+        loop {
+            tokio::select! {
+                _ = ticker.tick() => self.exporter.gather_and_send(),
+                // Stop on teardown so this task (and the pump `tx` clone it
+                // holds) doesn't outlive `ToolServer::shutdown` and keep
+                // gathering/sending forever.
+                _ = self.shutdown.cancelled() => break,
+            }
+        }
+    }
+}
+
+/// Shutdown fence: drains queued metric donations before the connection
+/// closes.
+pub struct MetricDonationPump {
+    tx: mpsc::Sender<PumpMsg>,
+}
+
+impl MetricDonationPump {
+    /// Resolves once every payload queued before this call has had a
+    /// send attempt.
+    pub async fn drain(&self) {
+        crate::donate_pump::drain_via(&self.tx).await;
+    }
+}
+
+impl ToolServer {
+    /// Post-connect entry point: spawn the metric donation pump (wiring
+    /// [`ToolServer::donate_metrics`]) plus the periodic registry
+    /// gatherer, and return a drain handle. Activates on server presence
+    /// (no env flag). `service_name` must be server-allowlisted.
+    pub fn metric_donation_reporter(&self, service_name: impl Into<String>) -> MetricDonationPump {
+        let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
+        let server = self.downgrade();
+        tokio::spawn(run_pump(rx, move |payload: String| {
+            let server = server.clone();
+            async move {
+                let Some(server) = server.upgrade() else {
+                    return (false, payload);
+                };
+                let ok = server.donate_metrics(&payload).await.is_ok();
+                (ok, payload)
+            }
+        }));
+        self.set_metric_donation_pump(tx.clone());
+
+        let exporter = MetricExporter {
+            tx: tx.clone(),
+            resource: make_resource(service_name.into()),
+        };
+        ACTIVE_METRIC_EXPORTER.store(Some(Arc::new(exporter.clone())));
+        tokio::spawn(
+            MetricDonationReporter {
+                exporter,
+                interval: DEFAULT_GATHER_INTERVAL,
+                shutdown: self.shutdown_token(),
+            }
+            .run(),
+        );
+
+        MetricDonationPump { tx }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use opentelemetry_proto::tonic::common::v1::any_value;
+    use prometheus::{Histogram, HistogramOpts, IntCounterVec, IntGauge, Opts, Registry};
+
+    use super::*;
+
+    fn decode(payload: String) -> ExportMetricsServiceRequest {
+        let bytes = base64::engine::general_purpose::STANDARD
+            .decode(payload)
+            .expect("payload must be base64");
+        ExportMetricsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP")
+    }
+
+    fn label_map(attrs: &[KeyValue]) -> std::collections::HashMap<String, String> {
+        attrs
+            .iter()
+            .filter_map(|kv| match kv.value.as_ref().and_then(|v| v.value.clone()) {
+                Some(any_value::Value::StringValue(s)) => Some((kv.key.clone(), s)),
+                _ => None,
+            })
+            .collect()
+    }
+
+    #[test]
+    fn converts_counter_gauge_histogram_with_labels() {
+        let registry = Registry::new();
+
+        let counter =
+            IntCounterVec::new(Opts::new("grok_test_total", "help"), &["reason"]).unwrap();
+        registry.register(Box::new(counter.clone())).unwrap();
+        counter.with_label_values(&["zdr"]).inc_by(5);
+
+        let gauge = IntGauge::new("grok_test_pending", "help").unwrap();
+        registry.register(Box::new(gauge.clone())).unwrap();
+        gauge.set(7);
+
+        let hist = Histogram::with_opts(
+            HistogramOpts::new("grok_test_seconds", "help").buckets(vec![0.5, 1.0]),
+        )
+        .unwrap();
+        registry.register(Box::new(hist.clone())).unwrap();
+        hist.observe(0.25);
+        hist.observe(0.75);
+        hist.observe(5.0);
+
+        let metrics = convert_families(&registry.gather());
+        let by_name: std::collections::HashMap<_, _> =
+            metrics.iter().map(|m| (m.name.clone(), m)).collect();
+
+        // Counter -> Sum (monotonic, cumulative), label preserved.
+        let metric::Data::Sum(sum) = by_name["grok_test_total"].data.as_ref().unwrap() else {
+            panic!("counter must convert to Sum");
+        };
+        assert!(sum.is_monotonic);
+        assert_eq!(
+            sum.aggregation_temporality,
+            AggregationTemporality::Cumulative as i32
+        );
+        let dp = &sum.data_points[0];
+        assert_eq!(dp.value, Some(number_data_point::Value::AsDouble(5.0)));
+        assert_eq!(
+            label_map(&dp.attributes).get("reason").map(String::as_str),
+            Some("zdr")
+        );
+
+        // Gauge -> Gauge.
+        let metric::Data::Gauge(g) = by_name["grok_test_pending"].data.as_ref().unwrap() else {
+            panic!("gauge must convert to Gauge");
+        };
+        assert_eq!(
+            g.data_points[0].value,
+            Some(number_data_point::Value::AsDouble(7.0))
+        );
+
+        // Histogram -> Histogram with cumulative buckets differenced and
+        // a +Inf bucket appended.
+        let metric::Data::Histogram(h) = by_name["grok_test_seconds"].data.as_ref().unwrap() else {
+            panic!("histogram must convert to Histogram");
+        };
+        assert_eq!(
+            h.aggregation_temporality,
+            AggregationTemporality::Cumulative as i32
+        );
+        let hdp = &h.data_points[0];
+        assert_eq!(hdp.count, 3);
+        assert_eq!(hdp.sum, Some(6.0));
+        assert_eq!(hdp.explicit_bounds, vec![0.5, 1.0]);
+        // (<=0.5): 0.25 -> 1 ; (0.5,1.0]: 0.75 -> 1 ; (+Inf): 5.0 -> 1
+        assert_eq!(hdp.bucket_counts, vec![1, 1, 1]);
+    }
+
+    #[test]
+    fn exporter_chunks_at_max_metrics_per_donation() {
+        let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
+        let exporter = MetricExporter {
+            tx,
+            resource: make_resource("test-service".to_owned()),
+        };
+        let metrics = vec![Metric::default(); MAX_METRICS_PER_DONATION + 1];
+        exporter.export(metrics);
+
+        let mut payloads = 0;
+        let mut total = 0;
+        while let Ok(PumpMsg::Payload(p)) = rx.try_recv() {
+            payloads += 1;
+            total += decode(p).resource_metrics[0].scope_metrics[0].metrics.len();
+        }
+        assert_eq!(payloads, 2, "one full chunk + remainder");
+        assert_eq!(total, MAX_METRICS_PER_DONATION + 1);
+    }
+
+    #[test]
+    fn summary_and_untyped_families_are_skipped() {
+        // An empty registry gathers nothing; convert yields nothing.
+        let registry = Registry::new();
+        assert!(convert_families(&registry.gather()).is_empty());
+    }
+}