crank/crates/crank-trace/src/lib.rs

//! Закрытый семантический контракт spans Crank.
//!
//! Этот crate не настраивает subscriber и не знает об OTLP. Он ограничивает
//! имена и атрибуты стадий статическим словарём, чтобы продуктовые crate не
//! могли случайно экспортировать пользовательские данные.

use std::future::Future;

use opentelemetry::{
    Context,
    trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState},
};
use tracing::{Instrument, Span, field::Empty, info_span};
use tracing_opentelemetry::OpenTelemetrySpanExt;

pub fn set_parent_from_trace_context(span: &Span, context: &crank_core::TraceContext) -> bool {
    let mut parts = context.traceparent().split('-');
    let (Some("00"), Some(trace_id), Some(parent_id), Some(flags), None) = (
        parts.next(),
        parts.next(),
        parts.next(),
        parts.next(),
        parts.next(),
    ) else {
        return false;
    };
    let (Ok(trace_id), Ok(parent_id)) = (TraceId::from_hex(trace_id), SpanId::from_hex(parent_id))
    else {
        return false;
    };
    let flags = match flags {
        "00" => TraceFlags::default(),
        "01" => TraceFlags::SAMPLED,
        _ => return false,
    };
    let parent = SpanContext::new(trace_id, parent_id, flags, true, TraceState::default());
    span.set_parent(Context::new().with_remote_span_context(parent))
        .is_ok()
}

pub fn trace_context_for_span(span: &Span) -> Option<crank_core::TraceContext> {
    let context = span.context();
    let span_context = context.span().span_context().clone();
    span_context.is_valid().then(|| {
        crank_core::TraceContext::from_span_parts(
            &span_context.trace_id().to_string(),
            &span_context.span_id().to_string(),
            span_context.is_sampled(),
        )
        .ok()
    })?
}

macro_rules! stage_span {
    ($name:literal) => {
        info_span!(
            target: "crank::trace",
            $name,
            outcome = Empty,
            error.category = Empty,
        )
    };
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Stage {
    McpRateLimit,
    McpAccessCheck,
    McpCatalogLoad,
    McpToolsResolve,
    ApprovalCheck,
    RuntimeExecute,
    RuntimeArgumentsMap,
    RuntimeIdempotency,
    UpstreamHttp,
    RuntimeResponseTransform,
    AuthResolve,
    ApprovalRecovery,
    HistoryWrite,
    DbQuery,
}

impl Stage {
    pub fn span(self) -> Span {
        match self {
            Self::McpRateLimit => stage_span!("mcp.rate_limit"),
            Self::McpAccessCheck => stage_span!("mcp.access.check"),
            Self::McpCatalogLoad => stage_span!("mcp.catalog.load"),
            Self::McpToolsResolve => stage_span!("mcp.tools.resolve"),
            Self::ApprovalCheck => stage_span!("approval.check"),
            Self::RuntimeExecute => stage_span!("runtime.execute"),
            Self::RuntimeArgumentsMap => stage_span!("runtime.arguments.map"),
            Self::RuntimeIdempotency => stage_span!("runtime.idempotency"),
            Self::UpstreamHttp => stage_span!("upstream.http"),
            Self::RuntimeResponseTransform => stage_span!("runtime.response.transform"),
            Self::AuthResolve => stage_span!("auth.resolve"),
            Self::ApprovalRecovery => stage_span!("approval.recovery"),
            Self::HistoryWrite => stage_span!("history.write"),
            Self::DbQuery => info_span!(
                target: "crank::trace",
                "db.query",
                outcome = Empty,
                error.category = Empty,
                db.system = "postgresql",
                db.operation = Empty,
            ),
        }
    }
}

pub async fn observe_db_query<T, E>(
    operation: DbOperation,
    future: impl Future<Output = Result<T, E>>,
) -> Result<T, E> {
    let span = operation.span();
    let result = future.instrument(span.clone()).await;
    match &result {
        Ok(_) => StageOutcome::Success.record(&span),
        Err(_) => {
            StageOutcome::Error.record(&span);
            ErrorCategory::Database.record(&span);
        }
    }
    result
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StageOutcome {
    Success,
    Error,
    Allowed,
    Denied,
    Required,
    Replay,
    Execute,
    Skipped,
    CacheHit,
}

impl StageOutcome {
    pub fn record(self, span: &Span) {
        span.record("outcome", self.as_str());
    }

    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Success => "success",
            Self::Error => "error",
            Self::Allowed => "allowed",
            Self::Denied => "denied",
            Self::Required => "required",
            Self::Replay => "replay",
            Self::Execute => "execute",
            Self::Skipped => "skipped",
            Self::CacheHit => "cache_hit",
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ErrorCategory {
    Access,
    RateLimit,
    Catalog,
    Approval,
    Idempotency,
    Schema,
    Mapping,
    Upstream,
    Transformation,
    History,
    Database,
    Concurrency,
    Configuration,
    Internal,
}

impl ErrorCategory {
    pub fn record(self, span: &Span) {
        span.record("error.category", self.as_str());
    }

    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Access => "access",
            Self::RateLimit => "rate_limit",
            Self::Catalog => "catalog",
            Self::Approval => "approval",
            Self::Idempotency => "idempotency",
            Self::Schema => "schema",
            Self::Mapping => "mapping",
            Self::Upstream => "upstream",
            Self::Transformation => "transformation",
            Self::History => "history",
            Self::Database => "database",
            Self::Concurrency => "concurrency",
            Self::Configuration => "configuration",
            Self::Internal => "internal",
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DbOperation {
    MachineAccessRead,
    MachineAccessTouch,
    CatalogLoad,
    ApprovalRead,
    ApprovalWrite,
    AuthProfileRead,
    SecretRead,
    SecretTouch,
    InvocationHistoryWrite,
}

impl DbOperation {
    pub fn span(self) -> Span {
        let span = Stage::DbQuery.span();
        span.record("db.operation", self.as_str());
        span
    }

    pub const fn as_str(self) -> &'static str {
        match self {
            Self::MachineAccessRead => "machine_access.read",
            Self::MachineAccessTouch => "machine_access.touch",
            Self::CatalogLoad => "catalog.load",
            Self::ApprovalRead => "approval.read",
            Self::ApprovalWrite => "approval.write",
            Self::AuthProfileRead => "auth_profile.read",
            Self::SecretRead => "secret.read",
            Self::SecretTouch => "secret.touch",
            Self::InvocationHistoryWrite => "invocation_history.write",
        }
    }
}