From 882d37b16f1867401faa149afbb00bf2e4184c4d Mon Sep 17 00:00:00 2001 From: Nicholas Velten Date: Mon, 6 Apr 2026 02:44:48 -0300 Subject: [PATCH] feat: OTLP metrics and logs export (closes #103) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add dual-export for metrics and structured logs alongside the existing OTLP trace pipeline. Both pipelines are opt-in via config and fail gracefully when the collector is unreachable. Config: telemetry: otlp_endpoint: "http://otel-collector:4317" export_metrics: true # mirrors Prometheus counters to OTLP export_logs: true # bridges all tracing:: events to OTLP logs Changes: - Cargo.toml: add metrics+logs features to opentelemetry/otlp/sdk; add opentelemetry-appender-tracing for the tracing→OTLP log bridge - config.rs: add export_metrics and export_logs fields (both default false) - metrics.rs: mirror arbitus.requests.total and arbitus.tokens.total to the global OTel meter (no-op when no provider is installed) - arbitus.rs: build_otel_metrics(), build_otel_logs(), updated OtelGuard to shut down all three providers on drop; init_tracing() uses Option to compose pipelines without combinatorial match arms Tests: - 3 unit tests for TelemetryConfig deserialization (defaults, all fields, partial) - 3 unit tests for metrics.rs OTel no-op behaviour - 5 integration tests: gateway starts/operates correctly with unreachable collector; Prometheus /metrics unaffected; blocked tools still work - e2e.sh section 21: health, request handling, and /metrics with OTLP config Co-Authored-By: Claude Sonnet 4.6 --- Cargo.lock | 13 +++++ Cargo.toml | 7 ++- src/bin/arbitus.rs | 138 +++++++++++++++++++++++++++++++++++--------- src/config.rs | 49 +++++++++++++++- src/metrics.rs | 68 ++++++++++++++++++++++ tests/common/mod.rs | 18 ++++++ tests/e2e.sh | 64 ++++++++++++++++++-- tests/otlp.rs | 80 +++++++++++++++++++++++++ 8 files changed, 402 insertions(+), 35 deletions(-) create mode 100644 tests/otlp.rs diff --git a/Cargo.lock b/Cargo.lock index d72d343..84f2be6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -114,6 +114,7 @@ dependencies = [ "libc", "lru", "opentelemetry", + "opentelemetry-appender-tracing", "opentelemetry-otlp", "opentelemetry_sdk", "percent-encoding", @@ -1964,6 +1965,18 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "opentelemetry-appender-tracing" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14f62d9a23c680ab91c74605f5006110768eb67600bb654937fef5c852fb8ec7" +dependencies = [ + "opentelemetry", + "tracing", + "tracing-core", + "tracing-subscriber", +] + [[package]] name = "opentelemetry-otlp" version = "0.26.0" diff --git a/Cargo.toml b/Cargo.toml index 93052c2..1702054 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,9 +40,10 @@ futures-util = "0.3" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } tracing-opentelemetry = "0.27" -opentelemetry = { version = "0.26", features = ["trace"] } -opentelemetry_sdk = { version = "0.26", features = ["rt-tokio"] } -opentelemetry-otlp = { version = "0.26", features = ["trace", "grpc-tonic"] } +opentelemetry = { version = "0.26", features = ["trace", "metrics", "logs"] } +opentelemetry_sdk = { version = "0.26", features = ["rt-tokio", "metrics", "logs"] } +opentelemetry-otlp = { version = "0.26", features = ["trace", "metrics", "logs", "grpc-tonic"] } +opentelemetry-appender-tracing = "0.26" jsonwebtoken = { version = "10", features = ["rust_crypto"] } subtle = "2" chrono = { version = "0.4", default-features = false, features = ["clock"] } diff --git a/src/bin/arbitus.rs b/src/bin/arbitus.rs index 77797f1..4f99596 100644 --- a/src/bin/arbitus.rs +++ b/src/bin/arbitus.rs @@ -37,7 +37,7 @@ use regex::Regex; use rusqlite::{Connection, types::Value}; use std::{collections::HashMap, sync::Arc, time::Duration}; use tokio::sync::watch; -use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; +use tracing_subscriber::{EnvFilter, Layer, layer::SubscriberExt, util::SubscriberInitExt}; // ── CLI definition ───────────────────────────────────────────────────────────── @@ -1198,11 +1198,20 @@ audits: [] // ── OpenTelemetry ────────────────────────────────────────────────────────────── -struct OtelGuard; +struct OtelGuard { + metrics_provider: Option, + logger_provider: Option, +} impl Drop for OtelGuard { fn drop(&mut self) { opentelemetry::global::shutdown_tracer_provider(); + if let Some(p) = self.metrics_provider.take() { + let _ = p.shutdown(); + } + if let Some(p) = self.logger_provider.take() { + let _ = p.shutdown(); + } } } @@ -1213,35 +1222,64 @@ fn init_tracing(telemetry: Option<&TelemetryConfig>) -> Option { let tracer = telemetry.and_then(|tel| match build_otel_tracer(tel) { Ok(t) => Some(t), Err(e) => { - eprintln!("warn: OTel init failed: {e}"); + eprintln!("warn: OTel traces init failed: {e}"); None } }); - let has_otel = tracer.is_some(); - - match (json, tracer) { - (true, Some(t)) => tracing_subscriber::registry() - .with(filter) - .with(tracing_subscriber::fmt::layer().json()) - .with(tracing_opentelemetry::layer().with_tracer(t)) - .init(), - (true, None) => tracing_subscriber::registry() - .with(filter) - .with(tracing_subscriber::fmt::layer().json()) - .init(), - (false, Some(t)) => tracing_subscriber::registry() - .with(filter) - .with(tracing_subscriber::fmt::layer()) - .with(tracing_opentelemetry::layer().with_tracer(t)) - .init(), - (false, None) => tracing_subscriber::registry() - .with(filter) - .with(tracing_subscriber::fmt::layer()) - .init(), - } + let metrics_provider = + telemetry + .filter(|t| t.export_metrics) + .and_then(|tel| match build_otel_metrics(tel) { + Ok(p) => { + opentelemetry::global::set_meter_provider(p.clone()); + Some(p) + } + Err(e) => { + eprintln!("warn: OTel metrics init failed: {e}"); + None + } + }); + + let logger_provider = + telemetry + .filter(|t| t.export_logs) + .and_then(|tel| match build_otel_logs(tel) { + Ok(p) => Some(p), + Err(e) => { + eprintln!("warn: OTel logs init failed: {e}"); + None + } + }); + + let has_traces = tracer.is_some(); + let otel_trace_layer = tracer.map(|t| tracing_opentelemetry::layer().with_tracer(t)); + let otel_log_layer = logger_provider + .as_ref() + .map(opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new); + + let fmt_layer = if json { + tracing_subscriber::fmt::layer().json().boxed() + } else { + tracing_subscriber::fmt::layer().boxed() + }; - if has_otel { Some(OtelGuard) } else { None } + tracing_subscriber::registry() + .with(filter) + .with(fmt_layer) + .with(otel_trace_layer) + .with(otel_log_layer) + .init(); + + let any_otel = has_traces || metrics_provider.is_some() || logger_provider.is_some(); + if any_otel { + Some(OtelGuard { + metrics_provider, + logger_provider, + }) + } else { + None + } } fn build_otel_tracer(tel: &TelemetryConfig) -> anyhow::Result { @@ -1268,3 +1306,51 @@ fn build_otel_tracer(tel: &TelemetryConfig) -> anyhow::Result anyhow::Result { + use opentelemetry::KeyValue; + use opentelemetry_otlp::WithExportConfig; + use opentelemetry_sdk::Resource; + + let resource = Resource::new(vec![KeyValue::new( + "service.name", + tel.service_name.clone(), + )]); + + opentelemetry_otlp::new_pipeline() + .metrics(opentelemetry_sdk::runtime::Tokio) + .with_exporter( + opentelemetry_otlp::new_exporter() + .tonic() + .with_endpoint(&tel.otlp_endpoint), + ) + .with_resource(resource) + .build() + .map_err(|e| anyhow::anyhow!("OTLP metrics pipeline: {e}")) +} + +fn build_otel_logs( + tel: &TelemetryConfig, +) -> anyhow::Result { + use opentelemetry::KeyValue; + use opentelemetry_otlp::WithExportConfig; + use opentelemetry_sdk::Resource; + + let resource = Resource::new(vec![KeyValue::new( + "service.name", + tel.service_name.clone(), + )]); + + opentelemetry_otlp::new_pipeline() + .logging() + .with_resource(resource) + .with_exporter( + opentelemetry_otlp::new_exporter() + .tonic() + .with_endpoint(&tel.otlp_endpoint), + ) + .install_batch(opentelemetry_sdk::runtime::Tokio) + .map_err(|e| anyhow::anyhow!("OTLP logs pipeline: {e}")) +} diff --git a/src/config.rs b/src/config.rs index 0fda1f8..5f1dabb 100644 --- a/src/config.rs +++ b/src/config.rs @@ -538,8 +538,8 @@ fn default_agent_claim() -> String { // ── Telemetry ───────────────────────────────────────────────────────────────── -/// OpenTelemetry tracing configuration. -/// When set, spans are exported to the configured OTLP endpoint. +/// OpenTelemetry configuration. +/// When set, spans, metrics, and/or logs are exported to the configured OTLP endpoint. #[derive(Debug, Deserialize, Clone)] pub struct TelemetryConfig { /// OTLP gRPC endpoint (e.g. `http://localhost:4317`). @@ -547,6 +547,12 @@ pub struct TelemetryConfig { /// `service.name` resource attribute. Defaults to `"arbitus"`. #[serde(default = "default_service_name")] pub service_name: String, + /// Export metrics via OTLP alongside the Prometheus `/metrics` endpoint. + #[serde(default)] + pub export_metrics: bool, + /// Export structured logs via OTLP (bridges all `tracing` events). + #[serde(default)] + pub export_logs: bool, } fn default_service_name() -> String { @@ -1040,4 +1046,43 @@ mod tests { ); assert!(cfg.validate().is_ok()); } + + // ── TelemetryConfig ─────────────────────────────────────────────────────── + + #[test] + fn telemetry_config_defaults() { + let yaml = r#" +otlp_endpoint: "http://localhost:4317" +"#; + let cfg: TelemetryConfig = serde_yaml::from_str(yaml).unwrap(); + assert_eq!(cfg.service_name, "arbitus"); + assert!(!cfg.export_metrics); + assert!(!cfg.export_logs); + } + + #[test] + fn telemetry_config_all_fields() { + let yaml = r#" +otlp_endpoint: "http://otel-collector:4317" +service_name: "my-service" +export_metrics: true +export_logs: true +"#; + let cfg: TelemetryConfig = serde_yaml::from_str(yaml).unwrap(); + assert_eq!(cfg.otlp_endpoint, "http://otel-collector:4317"); + assert_eq!(cfg.service_name, "my-service"); + assert!(cfg.export_metrics); + assert!(cfg.export_logs); + } + + #[test] + fn telemetry_config_partial_flags() { + let yaml = r#" +otlp_endpoint: "http://localhost:4317" +export_metrics: true +"#; + let cfg: TelemetryConfig = serde_yaml::from_str(yaml).unwrap(); + assert!(cfg.export_metrics); + assert!(!cfg.export_logs); + } } diff --git a/src/metrics.rs b/src/metrics.rs index 798e15c..8a4dccf 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -1,3 +1,4 @@ +use opentelemetry::{KeyValue, global}; use prometheus::{Counter, CounterVec, Encoder, Opts, Registry, TextEncoder}; pub struct GatewayMetrics { @@ -60,6 +61,18 @@ impl GatewayMetrics { pub fn record(&self, agent: &str, outcome: &str) { self.requests.with_label_values(&[agent, outcome]).inc(); + // Mirror to OTLP when a global meter provider is installed (no-op otherwise). + global::meter("arbitus") + .u64_counter("arbitus.requests.total") + .with_description("Total requests processed by arbitus") + .init() + .add( + 1, + &[ + KeyValue::new("agent", agent.to_string()), + KeyValue::new("outcome", outcome.to_string()), + ], + ); } /// Record estimated token usage for a single request. @@ -71,11 +84,33 @@ impl GatewayMetrics { self.tokens .with_label_values(&[agent, "input"]) .inc_by(f64::from(input_tokens)); + global::meter("arbitus") + .f64_counter("arbitus.tokens.total") + .with_description("Estimated tokens processed by arbitus") + .init() + .add( + f64::from(input_tokens), + &[ + KeyValue::new("agent", agent.to_string()), + KeyValue::new("direction", "input"), + ], + ); } if output_tokens > 0 { self.tokens .with_label_values(&[agent, "output"]) .inc_by(f64::from(output_tokens)); + global::meter("arbitus") + .f64_counter("arbitus.tokens.total") + .with_description("Estimated tokens processed by arbitus") + .init() + .add( + f64::from(output_tokens), + &[ + KeyValue::new("agent", agent.to_string()), + KeyValue::new("direction", "output"), + ], + ); } } @@ -133,4 +168,37 @@ mod tests { assert!(rendered.contains(r#"agent="cursor""#)); assert!(rendered.contains(r#"agent="claude""#)); } + + // ── OTel dual-export (no-op without provider) ───────────────────────────── + + #[test] + fn record_does_not_panic_without_otel_provider() { + // No global OTel meter provider installed — calls must be silent no-ops. + let m = GatewayMetrics::new().unwrap(); + m.record("cursor", "allowed"); + m.record("cursor", "blocked"); + // Prometheus counter still incremented + let rendered = m.render(); + assert!(rendered.contains("arbitus_requests_total")); + } + + #[test] + fn record_tokens_does_not_panic_without_otel_provider() { + let m = GatewayMetrics::new().unwrap(); + m.record_tokens("cursor", 100, 200); + let rendered = m.render(); + assert!(rendered.contains("arbitus_tokens_total")); + } + + #[test] + fn prometheus_metrics_unaffected_by_otel_calls() { + let m = GatewayMetrics::new().unwrap(); + m.record("agent-a", "allowed"); + m.record("agent-a", "allowed"); + m.record("agent-a", "blocked"); + let rendered = m.render(); + // Two allowed + one blocked — Prometheus counters must reflect this. + assert!(rendered.contains(r#"outcome="allowed""#)); + assert!(rendered.contains(r#"outcome="blocked""#)); + } } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 7d26626..7ec812e 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -305,6 +305,24 @@ pub async fn harness_streamable(config_snippet: &str) -> Harness { harness_inner(config_snippet, "streamable_http", "type: stdout").await } +/// Like `harness` but injects a `telemetry:` section pointing to the given OTLP endpoint. +pub async fn harness_with_telemetry( + config_snippet: &str, + otlp_endpoint: &str, + export_metrics: bool, + export_logs: bool, +) -> Harness { + let telemetry = format!( + "telemetry:\n otlp_endpoint: \"{otlp_endpoint}\"\n export_metrics: {export_metrics}\n export_logs: {export_logs}\n" + ); + harness_inner( + &format!("{telemetry}{config_snippet}"), + "http", + "type: stdout", + ) + .await +} + async fn harness_inner(config_snippet: &str, transport_type: &str, audit_config: &str) -> Harness { let (dummy_port, dummy_abort) = start_dummy().await; let gw_port = free_port().await; diff --git a/tests/e2e.sh b/tests/e2e.sh index 12b62e3..72439d3 100755 --- a/tests/e2e.sh +++ b/tests/e2e.sh @@ -136,9 +136,9 @@ chmod +x tests/mock-server.sh # 2. Cleanup cleanup() { echo -e "\n${YELLOW}🧹 Cleaning up processes and temp files...${NC}" - kill $DUMMY_PID $ARBITUS_PID $NODE_PID $STREAMABLE_PID 2>/dev/null || true - fuser -k 3000/tcp 4001/tcp 4002/tcp 5000/tcp 2>/dev/null || true - rm -rf concurrent_results/ tests/mock-server.sh output-stdio.jsonl tests/fixtures/gateway-hotreload.yml tests/fixtures/gateway-e2e-ip.yml tests/fixtures/gateway-e2e-streamable.yml *.log hitl_resp.txt webhook.log tests/node_helper.js tests/policy.rego tests/fixtures/gateway-verify.yml + kill $DUMMY_PID $ARBITUS_PID $NODE_PID $STREAMABLE_PID $OTLP_PID 2>/dev/null || true + fuser -k 3000/tcp 4001/tcp 4002/tcp 4003/tcp 5000/tcp 2>/dev/null || true + rm -rf concurrent_results/ tests/mock-server.sh output-stdio.jsonl tests/fixtures/gateway-hotreload.yml tests/fixtures/gateway-e2e-ip.yml tests/fixtures/gateway-e2e-streamable.yml tests/fixtures/gateway-e2e-otlp.yml *.log hitl_resp.txt webhook.log tests/node_helper.js tests/policy.rego tests/fixtures/gateway-verify.yml } trap cleanup EXIT @@ -464,10 +464,66 @@ else fail "invalidated session returned $AFTER_DEL (expected 404)"; fi kill $STREAMABLE_PID 2>/dev/null || true +echo -e "\n${CYAN}📡 21. OTLP METRICS + LOGS${NC}" +echo " Starting gateway with OTLP metrics+logs configured (collector unreachable)..." + +cat << 'EOF' > tests/fixtures/gateway-e2e-otlp.yml +transport: + type: http + addr: "127.0.0.1:4003" + upstream: "http://127.0.0.1:3000/mcp" +telemetry: + otlp_endpoint: "http://127.0.0.1:19999" + service_name: "arbitus-e2e" + export_metrics: true + export_logs: true +agents: + cursor: + allowed_tools: ["echo"] + rate_limit: 100 +EOF + +./target/debug/arbitus tests/fixtures/gateway-e2e-otlp.yml >> arbitus.log 2>&1 & +OTLP_PID=$! + +# Wait for health +for i in {1..40}; do + if curl -s http://127.0.0.1:4003/health | grep -q "ok\|healthy\|status"; then break; fi + sleep 0.1 +done + +echo " Testing gateway is healthy with OTLP configured (collector unreachable)..." +HEALTH=$(curl -s http://127.0.0.1:4003/health) +if echo "$HEALTH" | grep -qi "ok\|healthy\|status"; then pass "gateway healthy with unreachable OTLP collector" +else fail "gateway unhealthy with OTLP configured: $HEALTH"; fi + +echo " Testing request handling is unaffected by OTLP config..." +INIT_RESP=$(curl -s -i -X POST http://127.0.0.1:4003/mcp \ + -H "Content-Type: application/json" \ + -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"cursor","version":"1.0"}}}') +SID=$(echo "$INIT_RESP" | grep -i "mcp-session-id:" | awk '{print $2}' | tr -d '\r') +if [[ -n "$SID" ]]; then pass "initialize succeeds with OTLP configured" +else fail "initialize failed with OTLP configured"; fi + +ECHO_RESP=$(curl -s -X POST http://127.0.0.1:4003/mcp \ + -H "Content-Type: application/json" \ + -H "mcp-session-id: $SID" \ + -d '{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"echo","arguments":{"text":"otlp-e2e"}}}') +if echo "$ECHO_RESP" | grep -q "echo: otlp-e2e"; then pass "tools/call succeeds with OTLP configured" +else fail "tools/call failed with OTLP configured: $ECHO_RESP"; fi + +echo " Testing Prometheus /metrics still works alongside OTLP..." +METRICS=$(curl -s http://127.0.0.1:4003/metrics) +if echo "$METRICS" | grep -q "arbitus_requests_total"; then pass "Prometheus metrics endpoint intact with OTLP configured" +else fail "Prometheus metrics broken when OTLP is configured"; fi + +kill $OTLP_PID 2>/dev/null || true +rm -f tests/fixtures/gateway-e2e-otlp.yml + # ── Final summary ────────────────────────────────────────────────────────────── echo "" if [[ $FAILURES -eq 0 ]]; then - echo -e "${MAGENTA}🏆 ALL 20 SECTIONS PASSED${NC}" + echo -e "${MAGENTA}🏆 ALL 21 SECTIONS PASSED${NC}" else echo -e "${RED}✗ $FAILURES ASSERTION(S) FAILED${NC}" exit 1 diff --git a/tests/otlp.rs b/tests/otlp.rs new file mode 100644 index 0000000..f3501c1 --- /dev/null +++ b/tests/otlp.rs @@ -0,0 +1,80 @@ +mod common; + +use common::*; +use serde_json::json; + +// ── OTLP pipeline resilience ────────────────────────────────────────────────── +// +// These tests verify that the gateway starts and operates correctly when OTLP +// is configured, even when the collector is unreachable. The pipelines are +// asynchronous and batch-based, so a missing collector only causes background +// export errors — it must never crash the gateway or block request handling. + +#[tokio::test] +async fn gateway_starts_with_export_metrics_and_unreachable_collector() { + // Port 19999 is almost certainly not listening — collector unreachable. + let h = harness_with_telemetry(DEFAULT_CONFIG, "http://127.0.0.1:19999", true, false).await; + // Gateway must be healthy and serve requests normally. + let (sid, body) = h.init("cursor").await; + assert!(body["result"]["serverInfo"].is_object()); + assert!(!sid.is_empty()); +} + +#[tokio::test] +async fn gateway_starts_with_export_logs_and_unreachable_collector() { + let h = harness_with_telemetry(DEFAULT_CONFIG, "http://127.0.0.1:19999", false, true).await; + let (sid, _) = h.init("cursor").await; + assert!(!sid.is_empty()); +} + +#[tokio::test] +async fn gateway_starts_with_all_otlp_pipelines_and_unreachable_collector() { + let h = harness_with_telemetry(DEFAULT_CONFIG, "http://127.0.0.1:19999", true, true).await; + let (sid, _) = h.init("cursor").await; + assert!(!sid.is_empty()); +} + +// ── Prometheus /metrics unaffected ─────────────────────────────────────────── + +#[tokio::test] +async fn prometheus_metrics_still_work_when_otlp_is_configured() { + let h = harness_with_telemetry(DEFAULT_CONFIG, "http://127.0.0.1:19999", true, true).await; + + // Generate some traffic so counters are non-zero. + let (sid, _) = h.init("cursor").await; + h.json(Some(&sid), list_body()).await; + h.json(Some(&sid), call_body("echo", json!({"text": "otlp-test"}))) + .await; + + // /metrics must still respond with Prometheus text format. + let metrics = h + .client + .get(h.url("/metrics")) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + + assert!( + metrics.contains("arbitus_requests_total"), + "Prometheus counter missing: {metrics}" + ); +} + +// ── request handling unaffected by OTLP config ─────────────────────────────── + +#[tokio::test] +async fn tool_call_blocked_correctly_with_otlp_configured() { + let h = harness_with_telemetry(DEFAULT_CONFIG, "http://127.0.0.1:19999", true, true).await; + // cursor only has `echo` in allowed_tools + let (sid, _) = h.init("cursor").await; + let body = h + .json(Some(&sid), call_body("delete_database", json!({}))) + .await; + assert!( + body["error"].is_object(), + "expected blocked response; got: {body}" + ); +}