ImageApi/src/otel.rs

use actix_web::HttpRequest;
use actix_web::http::header::HeaderMap;
use opentelemetry::global::{BoxedSpan, BoxedTracer};
use opentelemetry::propagation::TextMapPropagator;
use opentelemetry::trace::{
    Span, SpanBuilder, SpanContext, SpanKind, Status, TraceContextExt, Tracer, TracerProvider,
};
use opentelemetry::{Context, KeyValue, global};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use opentelemetry_appender_log::OpenTelemetryLogBridge;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::logs::{BatchLogProcessor, SdkLoggerProvider};
use opentelemetry_sdk::propagation::TraceContextPropagator;

pub fn global_tracer() -> BoxedTracer {
    global::tracer("image-server")
}

/// Install the process-wide `tracing` subscriber plus the OTel span pipeline.
///
/// Called in BOTH debug and release builds (the old split made spans a no-op
/// in dev). Behaviour is driven by the environment, not the build profile:
/// - `OTLP_OTLS_ENDPOINT` set -> spans are batch-exported over OTLP (the
///   release / staging path).
/// - unset                    -> spans are created with valid contexts (so
///   outgoing `traceparent` injection still works) but exported nowhere; the
///   console `fmt` layer is what shows them.
///
/// The console `fmt` layer is always present, so a release build run locally
/// also prints spans to stdout. `log` records are NOT routed through this
/// subscriber (the `tracing-log` bridge is off) — they keep their existing
/// path: `env_logger` in debug, OTel logs in release.
pub fn init_tracing() {
    // W3C propagator for the whole process: the request middleware
    // (tracing-actix-web) and the outgoing `traceparent` helpers both read
    // the *global* propagator, so it must be set before any request lands.
    global::set_text_map_propagator(TraceContextPropagator::new());

    let resources = Resource::builder()
        .with_attributes([
            KeyValue::new("service.name", "image-server"),
            KeyValue::new("service.version", env!("CARGO_PKG_VERSION")),
        ])
        .build();

    let mut provider_builder =
        opentelemetry_sdk::trace::SdkTracerProvider::builder().with_resource(resources);

    if let Ok(endpoint) = std::env::var("OTLP_OTLS_ENDPOINT") {
        let span_exporter = opentelemetry_otlp::SpanExporter::builder()
            .with_tonic()
            .with_endpoint(endpoint)
            .build()
            .expect("OTLP span exporter");
        provider_builder = provider_builder.with_batch_exporter(span_exporter);
    }
    let tracer_provider = provider_builder.build();

    // The rest of the codebase resolves spans through the global provider
    // (`global_tracer()`, `global::tracer("db")`), so install it there; the
    // subscriber's OTel layer uses the same provider's tracer.
    let tracer = tracer_provider.tracer("image-server");
    global::set_tracer_provider(tracer_provider);

    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .with(tracing_opentelemetry::OpenTelemetryLayer::new(tracer))
        .with(tracing_subscriber::fmt::layer().with_target(false))
        .try_init()
        .expect("failed to install tracing subscriber");
}

#[allow(dead_code)]
pub fn init_logs() {
    let otlp_exporter = opentelemetry_otlp::LogExporter::builder()
        .with_tonic()
        .with_endpoint(std::env::var("OTLP_OTLS_ENDPOINT").unwrap())
        .build()
        .unwrap();

    let exporter = opentelemetry_stdout::LogExporter::default();

    let resources = Resource::builder()
        .with_attributes([
            KeyValue::new("service.name", "image-server"),
            KeyValue::new("service.version", env!("CARGO_PKG_VERSION")),
        ])
        .build();

    let log_provider = SdkLoggerProvider::builder()
        .with_log_processor(BatchLogProcessor::builder(exporter).build())
        .with_log_processor(BatchLogProcessor::builder(otlp_exporter).build())
        .with_resource(resources)
        .build();

    let otel_log_appender = OpenTelemetryLogBridge::new(&log_provider);
    log::set_boxed_logger(Box::new(otel_log_appender)).expect("Unable to set boxed logger");
    //TODO: Still set this with the env? Ideally we still have a clean/simple local logger for local dev
    log::set_max_level(log::LevelFilter::Info);
}

struct HeaderExtractor<'a>(&'a HeaderMap);

impl<'a> opentelemetry::propagation::Extractor for HeaderExtractor<'a> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).and_then(|v| v.to_str().ok())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(|k| k.as_str()).collect()
    }
}
/// Start a client span for a single LLM call. Every provider funnels its
/// LlmClient methods through this so the span shape (name, `SpanKind::Client`,
/// provider / operation / model attributes) is identical across
/// openrouter / llamacpp / ollama. `operation` is one of `chat`,
/// `chat_stream`, `embeddings`, `describe_image`, `generate`.
pub fn start_llm_request_span(
    cx: &Context,
    provider: &str,
    operation: &str,
    model: &str,
) -> BoxedSpan {
    let builder = SpanBuilder::from_name("llm.request").with_kind(SpanKind::Client);
    let mut span = global_tracer().build_with_context(builder, cx);
    span.set_attribute(KeyValue::new("llm.provider", provider.to_string()));
    span.set_attribute(KeyValue::new("llm.operation", operation.to_string()));
    span.set_attribute(KeyValue::new("llm.model", model.to_string()));
    span
}

/// The OpenTelemetry [`Context`] to parent this request's operation spans
/// under. Prefers the request's root span — created by the
/// `tracing-actix-web` middleware, which has already extracted any incoming
/// W3C `traceparent` — so the per-request operation span nests *under* the
/// root span (one trace tree) instead of duplicating it as a sibling. Falls
/// back to parsing the request's `traceparent` header when no root span is
/// active (e.g. a handler invoked directly in a test without the middleware).
pub fn extract_context_from_request(req: &HttpRequest) -> Context {
    use tracing_opentelemetry::OpenTelemetrySpanExt as _;
    let current = tracing::Span::current();
    let cx = current.context();
    if cx.span().span_context().is_valid() {
        return cx;
    }
    let propagator = TraceContextPropagator::new();
    propagator.extract(&HeaderExtractor(req.headers()))
}

/// Build the W3C trace-context headers to attach to an outgoing HTTP
/// request so a downstream service (llama-swap, Ollama) can parent its own
/// spans onto this trace. Returns an empty vec when the context carries no
/// valid span — e.g. debug builds, where the global tracer is a no-op and
/// injecting an all-zero traceparent would produce garbage parent spans.
pub fn traceparent_headers(context: &Context) -> Vec<(String, String)> {
    if !context.span().span_context().is_valid() {
        return Vec::new();
    }
    let mut headers: std::collections::HashMap<String, String> = std::collections::HashMap::new();
    TraceContextPropagator::new().inject_context(context, &mut headers);
    // The propagator writes an empty tracestate header when there is no
    // state — drop empty-valued headers rather than sending noise.
    headers.into_iter().filter(|(_, v)| !v.is_empty()).collect()
}

/// Attach the W3C traceparent header to an outgoing request so a downstream
/// service parents its own spans onto this trace. No-op when the span context
/// is not valid — e.g. debug builds, where the global tracer is a no-op and an
/// all-zero traceparent would produce garbage parent spans.
pub fn with_traceparent(
    builder: reqwest::RequestBuilder,
    span_context: &SpanContext,
) -> reqwest::RequestBuilder {
    let mut b = builder;
    for (k, v) in
        traceparent_headers(&Context::new().with_remote_span_context(span_context.clone()))
    {
        b = b.header(k, v);
    }
    b
}

/// Record the HTTP status code on the span (success or failure path).
pub fn set_http_status(span: Option<&mut BoxedSpan>, status: u16) {
    if let Some(s) = span {
        s.set_attribute(KeyValue::new("http.status_code", status as i64));
    }
}

/// Record an HTTP failure on the span: status code plus an error status.
pub fn record_http_failure(span: Option<&mut BoxedSpan>, status: u16, message: &str) {
    if let Some(s) = span {
        s.set_attribute(KeyValue::new("http.status_code", status as i64));
        s.set_status(Status::error(message.to_string()));
    }
}

pub fn trace_db_call<F, O>(
    context: &Context,
    query_type: &str,
    operation: &str,
    func: F,
) -> anyhow::Result<O>
where
    F: FnOnce(&mut BoxedSpan) -> anyhow::Result<O>,
{
    let tracer = global::tracer("db");
    let mut span = tracer
        .span_builder(format!("db.{}.{}", query_type, operation))
        .with_attributes(vec![
            KeyValue::new("db.query_type", query_type.to_string().clone()),
            KeyValue::new("db.operation", operation.to_string().clone()),
        ])
        .start_with_context(&tracer, context);

    let result = func(&mut span);
    match &result {
        Ok(_) => {
            span.set_status(Status::Ok);
        }
        Err(e) => span.set_status(Status::error(e.to_string())),
    }

    result
}

#[cfg(test)]
mod traceparent_tests {
    use super::*;
    use opentelemetry::trace::{
        SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState,
    };

    /// A context carrying a span context with known ids — no global
    /// provider needed, so tests can't pollute process state.
    fn context_with_flags(flags: TraceFlags) -> Context {
        let trace_id = TraceId::from_bytes([
            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
            0x0f, 0x10,
        ]);
        let span_id = SpanId::from_bytes([0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]);
        let span_context = SpanContext::new(trace_id, span_id, flags, false, TraceState::default());
        Context::new().with_remote_span_context(span_context)
    }

    #[test]
    fn traceparent_headers_sampled_context() {
        let cx = context_with_flags(TraceFlags::SAMPLED);
        let headers = traceparent_headers(&cx);
        assert_eq!(
            headers,
            vec![(
                "traceparent".to_string(),
                "00-0102030405060708090a0b0c0d0e0f10-0102030405060708-01".to_string(),
            )]
        );
    }

    #[test]
    fn traceparent_headers_not_sampled_context() {
        // A valid-but-unsampled context still propagates; the flags byte
        // carries the state.
        let cx = context_with_flags(TraceFlags::NOT_SAMPLED);
        let headers = traceparent_headers(&cx);
        assert_eq!(
            headers,
            vec![(
                "traceparent".to_string(),
                "00-0102030405060708090a0b0c0d0e0f10-0102030405060708-00".to_string(),
            )]
        );
    }

    #[test]
    fn traceparent_headers_invalid_context_is_empty() {
        // No span (debug-build no-op tracer) → no header, never an
        // all-zero traceparent.
        let cx = Context::new();
        assert!(traceparent_headers(&cx).is_empty());
    }
}