klieo_otel/

lib.rs

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
//! `klieo-otel` — OpenTelemetry OTLP exporter wiring for klieo agent
//! runtimes.
//!
//! The klieo runtime is instrumented with `tracing` spans on every
//! [`Agent::run_steps`](https://docs.rs/klieo-core) call, every LLM
//! invocation, and every tool dispatch. Field names follow the
//! [OpenTelemetry GenAI semantic
//! conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/),
//! so any OTel-aware backend (Tempo, Jaeger, Honeycomb, Datadog, …)
//! ingests them once you wire up an exporter.
//!
//! This crate is that exporter. One function call:
//!
//! ```rust,no_run
//! use klieo_otel::{install, OtelConfig, Protocol};
//! use std::time::Duration;
//!
//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
//! let _guard = install(
//!     OtelConfig::new("my-agent")
//!         .service_version(env!("CARGO_PKG_VERSION"))
//!         .gen_ai_system("ollama")
//!         .endpoint("http://localhost:4317")
//!         .protocol(Protocol::Grpc)
//!         .timeout(Duration::from_secs(5)),
//! )?;
//! // run agents here — spans flow to the configured collector
//! // until `_guard` is dropped.
//! # Ok(()) }
//! ```
//!
//! ## Defaults
//!
//! - **Endpoint** — `OTEL_EXPORTER_OTLP_ENDPOINT` env var (the
//!   standard OTel SDK convention) is honoured if set; otherwise the
//!   value passed to [`OtelConfig::endpoint`] is used; otherwise
//!   `http://localhost:4317`.
//! - **Protocol** — gRPC (`Protocol::Grpc`). Switch with
//!   [`OtelConfig::protocol`].
//! - **Resource attributes** — `service.name`, `service.version`,
//!   `gen_ai.system` are emitted on every span, per OTel semantic
//!   conventions.
//!
//! ## Guard semantics
//!
//! [`OtelGuard`]'s `Drop` impl force-flushes any in-flight batch and
//! shuts the provider down cleanly. Bind it to a binding that lives
//! for the lifetime of your process (e.g. in `main`):
//!
//! ```rust,no_run
//! # use klieo_otel::{install, OtelConfig};
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! let _otel = install(OtelConfig::new("my-agent"))?;
//! // … run server … the guard drops at end of main, flushing spans.
//! # Ok(()) }
//! ```
//!
//! Prefix the binding with `_` (not `_otel = …` then unused) to
//! suppress the unused-variable warning while keeping the binding
//! alive — `let _ = …` would drop the guard immediately and discard
//! every export.
//!
//! ## Composition with existing subscribers
//!
//! [`install`] sets a global `tracing` subscriber that combines:
//!
//! - an `env-filter` layer respecting `RUST_LOG`,
//! - the OpenTelemetry layer that exports spans over OTLP,
//! - a human-readable `fmt` layer on stderr.
//!
//! If you already call `tracing_subscriber::fmt().init()` (or any
//! other subscriber installer) earlier in `main`, [`install`] will
//! return [`OtelError::AlreadyInitialized`]. Call this crate **first**
//! (or skip the other call) to avoid a double-install.

mod config;
mod error;
mod guard;

#[cfg(any(test, feature = "test-utils"))]
pub mod test_utils;

pub use config::{OtelConfig, Protocol};
pub use error::OtelError;
pub use guard::OtelGuard;

pub use tracing_opentelemetry::OpenTelemetrySpanExt;

use opentelemetry::trace::TracerProvider as _;
use opentelemetry::KeyValue;
use opentelemetry_otlp::{SpanExporter, WithExportConfig, WithHttpConfig, WithTonicConfig};
use opentelemetry_sdk::{trace::TracerProvider as SdkTracerProvider, Resource};
use tonic::metadata::{MetadataKey, MetadataMap, MetadataValue};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::EnvFilter;

/// Standard OTel SDK environment variable: when set, overrides
/// [`OtelConfig::endpoint`].
pub const OTEL_EXPORTER_OTLP_ENDPOINT_ENV: &str = "OTEL_EXPORTER_OTLP_ENDPOINT";

/// Install a global OpenTelemetry tracing subscriber.
///
/// Builds an OTLP exporter, attaches it to a Tokio-batched
/// [`SdkTracerProvider`], wires it through `tracing-opentelemetry`,
/// and installs the result as the global `tracing` subscriber.
///
/// Returns an [`OtelGuard`] that flushes and shuts the provider down
/// when dropped. Bind it to a `_guard` (or longer-lived) binding so
/// it survives until process shutdown.
///
/// # Errors
///
/// - [`OtelError::Build`] — the OTLP exporter could not be constructed
///   (typically a malformed endpoint or unsupported protocol/feature
///   combination).
/// - [`OtelError::AlreadyInitialized`] — a global `tracing`
///   subscriber is already in place. Call this crate before any other
///   subscriber initializer.
///
/// # Panics
///
/// Does not panic. All failure paths return an [`OtelError`].
pub fn install(config: OtelConfig) -> Result<OtelGuard, OtelError> {
    let endpoint = resolve_endpoint(&config);
    // W8 / W2.A12 / CWE-319: refuse cleartext (any non-https) export to
    // non-loopback collectors. Spans carry `gen_ai.prompt` /
    // `gen_ai.completion` attributes — PII / customer prompts — and must not
    // transit in cleartext. Loopback collectors stay allowed so local dev works.
    if config::is_plaintext_remote_otlp(&endpoint) {
        return Err(OtelError::PlaintextRemote(endpoint));
    }
    let provider = build_provider(&config, &endpoint)?;
    let tracer = provider.tracer(config.service_name.clone());

    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
    let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
    let fmt_layer = tracing_subscriber::fmt::layer().with_writer(std::io::stderr);

    tracing_subscriber::registry()
        .with(env_filter)
        .with(otel_layer)
        .with(fmt_layer)
        .try_init()
        .map_err(|e| OtelError::AlreadyInitialized(e.to_string()))?;

    // Promote our provider to the global so that components reading
    // `opentelemetry::global::tracer(...)` see the same exporter.
    opentelemetry::global::set_tracer_provider(provider.clone());

    Ok(OtelGuard::new(provider))
}

/// Install klieo's OpenTelemetry exporter, tolerating an already-installed
/// `tracing` subscriber.
///
/// Same as [`install`] when no global `tracing` subscriber has been set
/// yet — installs klieo's full registry (env-filter + OTel layer + fmt).
///
/// When a global subscriber is already set, the subscriber install step
/// is skipped silently (Rust's `tracing` API exposes no way to uninstall
/// a previously-set global default). The OTel tracer provider is still
/// promoted to the OpenTelemetry global so that any third-party
/// subscriber wired up with `tracing-opentelemetry::layer().with_tracer(
/// opentelemetry::global::tracer(...))` will route klieo's spans through
/// the OTLP exporter. Callers get a working [`OtelGuard`] either way and
/// the flush-on-drop semantics are preserved.
///
/// Use cases:
/// - Integration tests where another fixture pre-installs `tracing_subscriber::fmt()`.
/// - Plug-in loaders or libraries where klieo runs after host
///   initialisation.
/// - Test harnesses with mixed `klieo-otel::install` / `install_or_replace`
///   invocations across test cases.
///
/// Prefer [`install`] when you control the binary's startup order so
/// the env-filter + fmt + OTel layer all reach `tracing` events.
///
/// # Errors
///
/// - [`OtelError::Build`] — the OTLP exporter could not be constructed.
/// - [`OtelError::PlaintextRemote`] — endpoint is plaintext-HTTP to a
///   non-loopback collector (same redaction-defence policy as [`install`]).
///
/// Never returns [`OtelError::AlreadyInitialized`].
pub fn install_or_replace(config: OtelConfig) -> Result<OtelGuard, OtelError> {
    let endpoint = resolve_endpoint(&config);
    if config::is_plaintext_remote_otlp(&endpoint) {
        return Err(OtelError::PlaintextRemote(endpoint));
    }
    let provider = build_provider(&config, &endpoint)?;
    let tracer = provider.tracer(config.service_name.clone());

    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
    let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
    let fmt_layer = tracing_subscriber::fmt::layer().with_writer(std::io::stderr);

    // Best-effort subscriber install. If a subscriber is already set,
    // `try_init` errors and we degrade to provider-only mode.
    let subscriber_init_result = tracing_subscriber::registry()
        .with(env_filter)
        .with(otel_layer)
        .with(fmt_layer)
        .try_init();
    if let Err(err) = subscriber_init_result {
        tracing::debug!(error = %err, "klieo-otel install_or_replace: subscriber already set; routing via OTel global tracer provider only");
    }

    opentelemetry::global::set_tracer_provider(provider.clone());

    Ok(OtelGuard::new(provider))
}

/// Resolve the OTLP endpoint, preferring the standard env var.
pub(crate) fn resolve_endpoint(config: &OtelConfig) -> String {
    if let Ok(env) = std::env::var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV) {
        if !env.trim().is_empty() {
            return env;
        }
    }
    config.endpoint.clone()
}

/// Convert configured `(name, value)` headers into a gRPC [`MetadataMap`].
///
/// Header names are lowercased (gRPC metadata keys are case-insensitive ASCII).
/// A name or value that cannot be represented as gRPC metadata returns an error
/// so the exporter build fails closed — an auth header is never silently
/// dropped, which would ship spans unauthenticated.
fn headers_to_metadata(headers: &[(String, String)]) -> Result<MetadataMap, OtelError> {
    let mut metadata = MetadataMap::with_capacity(headers.len());
    for (name, value) in headers {
        let key = MetadataKey::from_bytes(name.to_ascii_lowercase().as_bytes())
            .map_err(|e| OtelError::Build(format!("invalid OTLP header name {name:?}: {e}")))?;
        let value = MetadataValue::try_from(value.as_str()).map_err(|e| {
            OtelError::Build(format!("invalid OTLP header value for {name:?}: {e}"))
        })?;
        metadata.insert(key, value);
    }
    Ok(metadata)
}

/// Build the [`SdkTracerProvider`] from the resolved endpoint.
pub(crate) fn build_provider(
    config: &OtelConfig,
    endpoint: &str,
) -> Result<SdkTracerProvider, OtelError> {
    // `OtelConfig.headers` carries auth tokens (e.g. `authorization: Bearer …`,
    // `x-honeycomb-team: …`) for remote collectors. They MUST reach the
    // exporter — a dropped header means a 401 / silent telemetry loss, and any
    // `gen_ai.*` PII on the span would have shipped unauthenticated. Wire them
    // per transport (gRPC metadata / HTTP headers); an unrepresentable header
    // fails the build closed rather than silently dropping it.
    let exporter = match config.protocol {
        Protocol::Grpc => {
            let mut builder = SpanExporter::builder()
                .with_tonic()
                .with_endpoint(endpoint)
                .with_timeout(config.timeout);
            if !config.headers.is_empty() {
                builder = builder.with_metadata(headers_to_metadata(&config.headers)?);
            }
            builder
                .build()
                .map_err(|e| OtelError::Build(e.to_string()))?
        }
        Protocol::HttpProtobuf => {
            let mut builder = SpanExporter::builder()
                .with_http()
                .with_endpoint(endpoint)
                .with_timeout(config.timeout);
            if !config.headers.is_empty() {
                builder = builder.with_headers(config.headers.iter().cloned().collect());
            }
            builder
                .build()
                .map_err(|e| OtelError::Build(e.to_string()))?
        }
    };

    let mut attrs = vec![
        KeyValue::new("service.name", config.service_name.clone()),
        KeyValue::new("gen_ai.system", config.gen_ai_system.clone()),
    ];
    if let Some(version) = &config.service_version {
        attrs.push(KeyValue::new("service.version", version.clone()));
    }

    let provider = SdkTracerProvider::builder()
        .with_batch_exporter(exporter, opentelemetry_sdk::runtime::Tokio)
        .with_resource(Resource::new(attrs))
        .build();

    Ok(provider)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn headers_to_metadata_carries_valid_entries() {
        let md = headers_to_metadata(&[
            ("authorization".into(), "Bearer abc".into()),
            ("x-honeycomb-team".into(), "team-1".into()),
        ])
        .expect("valid headers");
        assert_eq!(
            md.get("authorization").unwrap().to_str().unwrap(),
            "Bearer abc"
        );
        assert_eq!(
            md.get("x-honeycomb-team").unwrap().to_str().unwrap(),
            "team-1"
        );
    }

    #[test]
    fn headers_to_metadata_fails_closed_on_invalid_name() {
        // A space is not a valid metadata key; the build must fail rather than
        // silently drop the (possibly auth) header.
        let err = headers_to_metadata(&[("bad name".into(), "v".into())]).unwrap_err();
        assert!(matches!(err, OtelError::Build(_)), "got {err:?}");
    }

    #[test]
    fn headers_to_metadata_fails_closed_on_invalid_value() {
        // A newline in a metadata value is rejected.
        let err =
            headers_to_metadata(&[("authorization".into(), "Bearer \nx".into())]).unwrap_err();
        assert!(matches!(err, OtelError::Build(_)), "got {err:?}");
    }

    #[test]
    fn config_builder_round_trip() {
        let cfg = OtelConfig::new("svc")
            .endpoint("http://collector:4317")
            .protocol(Protocol::Grpc)
            .timeout(Duration::from_secs(7))
            .service_version("1.2.3")
            .gen_ai_system("ollama");
        assert_eq!(cfg.service_name, "svc");
        assert_eq!(cfg.endpoint, "http://collector:4317");
        assert!(matches!(cfg.protocol, Protocol::Grpc));
        assert_eq!(cfg.timeout, Duration::from_secs(7));
        assert_eq!(cfg.service_version.as_deref(), Some("1.2.3"));
        assert_eq!(cfg.gen_ai_system, "ollama");
    }

    /// `OTEL_EXPORTER_OTLP_ENDPOINT` is process-global, so any test
    /// that mutates it must hold this lock to avoid racing with
    /// sibling tests that read or mutate it. Cargo runs tests in
    /// parallel by default.
    static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

    #[test]
    fn env_var_overrides_endpoint_default() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let cfg = OtelConfig::new("svc").endpoint("http://default:4317");

        let prior = std::env::var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV).ok();
        std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, "http://override:4317");
        let resolved = resolve_endpoint(&cfg);
        match prior {
            Some(v) => std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, v),
            None => std::env::remove_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV),
        }

        assert_eq!(resolved, "http://override:4317");
    }

    #[test]
    fn empty_env_var_falls_back_to_config() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let cfg = OtelConfig::new("svc").endpoint("http://default:4317");
        let prior = std::env::var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV).ok();
        std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, "   ");
        let resolved = resolve_endpoint(&cfg);
        match prior {
            Some(v) => std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, v),
            None => std::env::remove_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV),
        }
        assert_eq!(resolved, "http://default:4317");
    }

    #[test]
    fn unset_env_var_falls_back_to_config() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let cfg = OtelConfig::new("svc").endpoint("http://default:4317");
        let prior = std::env::var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV).ok();
        std::env::remove_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV);
        let resolved = resolve_endpoint(&cfg);
        if let Some(v) = prior {
            std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, v);
        }
        assert_eq!(resolved, "http://default:4317");
    }

    /// `install_or_replace` shares the redaction-defence policy of
    /// `install`: plaintext http:// to a non-loopback collector still
    /// errors with [`OtelError::PlaintextRemote`].
    #[test]
    fn install_or_replace_refuses_plaintext_remote_endpoint() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let prior = std::env::var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV).ok();
        std::env::remove_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV);
        let cfg = OtelConfig::new("svc").endpoint("http://collector.example.com:4317");
        let result = install_or_replace(cfg);
        if let Some(v) = prior {
            std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, v);
        }
        assert!(matches!(result, Err(OtelError::PlaintextRemote(_))));
    }

    /// `install_or_replace` returns `Ok` even when a global tracing
    /// subscriber is already set — installing one before the call
    /// proves we do not return [`OtelError::AlreadyInitialized`].
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn install_or_replace_succeeds_after_global_subscriber_set() {
        let result = {
            let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
            let prior = std::env::var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV).ok();
            std::env::remove_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV);

            // Pre-install a basic subscriber. We don't care whether THIS
            // call wins or loses (parallel tests may have set one first)
            // — we just need a default to be in place when
            // install_or_replace runs.
            let _ = tracing_subscriber::fmt().with_test_writer().try_init();

            // Loopback endpoint passes the plaintext-redaction check.
            let cfg = OtelConfig::new("svc")
                .endpoint("http://127.0.0.1:1")
                .timeout(Duration::from_millis(50));
            let r = install_or_replace(cfg);

            if let Some(v) = prior {
                std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, v);
            }
            r
        };

        let guard = result.expect("install_or_replace must not return AlreadyInitialized");
        // Drop the guard on a blocking thread (mirrors guard_drop test).
        tokio::task::spawn_blocking(move || drop(guard))
            .await
            .expect("drop must not panic");
    }

    /// W8 / W2.A12 / CWE-319: install must refuse plaintext http:// to
    /// non-loopback collectors. The check fires before any subscriber
    /// install side-effects, so re-running this test in a process that
    /// has already installed a subscriber is safe.
    #[test]
    fn install_refuses_plaintext_remote_endpoint() {
        let _g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
        let prior = std::env::var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV).ok();
        std::env::remove_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV);
        let cfg = OtelConfig::new("svc").endpoint("http://collector.example.com:4317");
        let result = install(cfg);
        if let Some(v) = prior {
            std::env::set_var(OTEL_EXPORTER_OTLP_ENDPOINT_ENV, v);
        }
        assert!(matches!(result, Err(OtelError::PlaintextRemote(_))));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn guard_drop_invokes_flush_and_shutdown() {
        // Build a provider pointing at a deliberately unreachable
        // endpoint with a tiny timeout. The provider construction
        // itself must succeed; flush/shutdown on drop must not panic,
        // even if no spans were ever exported.
        //
        // We need >=2 worker threads so the batch processor task can
        // make progress while the synchronous `force_flush` blocks.
        let cfg = OtelConfig::new("svc")
            .endpoint("http://127.0.0.1:1") // unreachable
            .timeout(Duration::from_millis(50));
        let provider = build_provider(&cfg, &cfg.endpoint).expect("build provider");
        let guard = OtelGuard::new(provider);
        let marker = guard.shutdown_marker();
        assert!(!marker.load(std::sync::atomic::Ordering::SeqCst));
        // Dropping the guard must not panic and must mark itself flushed.
        // Run drop on a blocking thread so the synchronous force_flush
        // can wait on the batch worker without starving the runtime.
        tokio::task::spawn_blocking(move || drop(guard))
            .await
            .expect("drop must not panic");
        assert!(marker.load(std::sync::atomic::Ordering::SeqCst));
    }
}