otelite_api/api/

traces.rs

use super::validate_time_range;
use crate::server::{AppState, QueryCache};
use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    response::IntoResponse,
    Json,
};
use otelite_core::api::{
    ErrorResponse, LogEntry, LogsResponse, SpanEntry, TraceDetail, TraceEntry, TracesResponse,
};
use otelite_core::query::{Operator, QueryPredicate, QueryValue};
use otelite_core::storage::QueryParams;
use serde::{Deserialize, Serialize};

/// Query parameters for trace listing
#[derive(Debug, Deserialize, Serialize, utoipa::IntoParams)]
pub struct TracesQuery {
    /// Filter by trace ID
    #[serde(default)]
    pub trace_id: Option<String>,

    /// Filter by service name
    #[serde(default)]
    pub service: Option<String>,

    /// Filter by resource attribute (format: key=value)
    #[serde(default)]
    pub resource: Option<String>,

    /// Full-text search in span names
    #[serde(default)]
    pub search: Option<String>,

    /// Filter by session ID (session.id attribute)
    #[serde(default)]
    pub session_id: Option<String>,

    /// Filter by gen_ai.conversation.id attribute
    #[serde(default)]
    pub conversation_id: Option<String>,

    /// Filter by model (matches gen_ai.request.model)
    #[serde(default)]
    pub model: Option<String>,

    /// JSON-encoded list of attribute filters: [{"key":"x","op":"=","value":"y"}, ...]
    /// Supported ops: "=" (Equal), "!=" (NotEqual)
    #[serde(default)]
    pub attrs: Option<String>,

    /// Structured query predicate; multiple predicates joined with `&&`,
    /// e.g. `name = "claude_code.llm_request" && attributes.attempt >= 2`.
    /// A malformed expression returns 400.
    #[serde(default)]
    pub query: Option<String>,

    /// Start time (Unix timestamp in nanoseconds)
    #[serde(default)]
    pub start_time: Option<i64>,

    /// End time (Unix timestamp in nanoseconds)
    #[serde(default)]
    pub end_time: Option<i64>,

    /// Maximum number of results (default: 100, max: 1000)
    #[serde(default = "default_limit")]
    pub limit: usize,

    /// Offset for pagination
    #[serde(default)]
    pub offset: usize,
}

fn default_limit() -> usize {
    100
}

/// Handler for GET /api/traces
#[utoipa::path(
    get,
    path = "/api/traces",
    params(TracesQuery),
    responses(
        (status = 200, description = "List of traces matching query", body = TracesResponse),
        (status = 400, description = "Invalid time range", body = ErrorResponse),
        (status = 500, description = "Internal server error", body = ErrorResponse)
    ),
    tag = "traces"
)]
pub async fn list_traces(
    State(state): State<AppState>,
    Query(params): Query<TracesQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    validate_time_range(params.start_time, params.end_time)?;

    // Check cache first
    let cache_key = QueryCache::make_key(&params);
    if let Some(cached) = state.cache.traces.get(&cache_key) {
        return Ok((
            StatusCode::OK,
            [("content-type", "application/json")],
            cached,
        )
            .into_response());
    }

    // Validate and cap limit
    let limit = params.limit.min(1000);

    // Build query parameters
    let mut query = QueryParams {
        start_time: params.start_time,
        end_time: params.end_time,
        trace_id: params.trace_id.clone(),
        ..Default::default()
    };

    if let Some(ref sid) = params.session_id {
        if !sid.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "session.id".to_string(),
                operator: Operator::Equal,
                value: QueryValue::String(sid.clone()),
            });
        }
    }

    if let Some(ref cid) = params.conversation_id {
        if !cid.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "gen_ai.conversation.id".to_string(),
                operator: Operator::Equal,
                value: QueryValue::String(cid.clone()),
            });
        }
    }

    if let Some(ref model) = params.model {
        if !model.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "gen_ai.request.model".to_string(),
                operator: Operator::Equal,
                value: QueryValue::String(model.clone()),
            });
        }
    }

    if let Some(ref attrs_json) = params.attrs {
        if !attrs_json.is_empty() {
            #[derive(serde::Deserialize)]
            struct AttrFilter {
                key: String,
                op: String,
                value: Option<String>,
            }
            if let Ok(filters) = serde_json::from_str::<Vec<AttrFilter>>(attrs_json) {
                for f in filters {
                    let op = match f.op.as_str() {
                        "=" => Operator::Equal,
                        "!=" => Operator::NotEqual,
                        _ => continue,
                    };
                    if let Some(v) = f.value {
                        query.predicates.push(QueryPredicate {
                            field: f.key,
                            operator: op,
                            value: QueryValue::String(v),
                        });
                    }
                }
            }
        }
    }

    // `search` filters on span names and `service`/`resource` on resource
    // attributes — both via the structured-predicate path so they constrain
    // trace selection exactly as any other predicate does.
    if let Some(ref s) = params.search {
        if !s.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "name".to_string(),
                operator: Operator::Contains,
                value: QueryValue::String(s.clone()),
            });
        }
    }
    if let Some(ref svc) = params.service {
        if !svc.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "resource.service.name".to_string(),
                operator: Operator::Equal,
                value: QueryValue::String(svc.clone()),
            });
        }
    }
    if let Some(ref kv) = params.resource {
        if let Some((k, v)) = kv.split_once('=') {
            if !k.is_empty() && !v.is_empty() {
                query.predicates.push(QueryPredicate {
                    field: format!("resource.{k}"),
                    operator: Operator::Equal,
                    value: QueryValue::String(v.to_string()),
                });
            }
        }
    }

    // Structured `query` predicate(s). Parse errors are a 400, never a
    // silently-ignored filter.
    if let Some(ref q) = params.query {
        if !q.is_empty() {
            query.predicates.extend(super::parse_query_param(q)?);
        }
    }

    // Per-trace summaries straight from storage (#251): the old path
    // fetched and parsed every span of the selected traces (2.9M spans
    // for a 1 h window on a production database) to derive these eight
    // fields; the summaries are computed in SQL instead.
    let mut trace_entries: Vec<TraceEntry> = state
        .storage
        .query_trace_summaries(&query, limit)
        .await
        .map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(ErrorResponse::storage_error(format!(
                    "query trace summaries: {}",
                    e
                ))),
            )
        })?;

    // Sort by start time (newest first)
    trace_entries.sort_by_key(|b| std::cmp::Reverse(b.start_time));

    // Apply pagination
    let total = trace_entries.len();
    let paginated_traces: Vec<TraceEntry> = trace_entries
        .into_iter()
        .skip(params.offset)
        .take(limit)
        .collect();

    let response = TracesResponse {
        traces: paginated_traces,
        total,
        limit,
        offset: params.offset,
    };

    // Cache the response
    if let Ok(json) = serde_json::to_string(&response) {
        state.cache.traces.insert(cache_key, json.clone());
        Ok((StatusCode::OK, [("content-type", "application/json")], json).into_response())
    } else {
        Ok(Json(response).into_response())
    }
}

/// Handler for GET /api/traces/:trace_id
#[utoipa::path(
    get,
    path = "/api/traces/{trace_id}",
    params(
        ("trace_id" = String, Path, description = "Trace ID")
    ),
    responses(
        (status = 200, description = "Trace details with all spans", body = TraceDetail),
        (status = 404, description = "Trace not found", body = ErrorResponse),
        (status = 500, description = "Internal server error", body = ErrorResponse)
    ),
    tag = "traces"
)]
pub async fn get_trace(
    State(state): State<AppState>,
    Path(trace_id): Path<String>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    // Query all spans for this trace
    let query = QueryParams {
        trace_id: Some(trace_id.clone()),
        limit: Some(1000), // Max spans per trace
        ..Default::default()
    };

    let spans = state.storage.query_spans(&query).await.map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse::storage_error(format!(
                "query trace spans: {}",
                e
            ))),
        )
    })?;

    if spans.is_empty() {
        return Err((
            StatusCode::NOT_FOUND,
            Json(ErrorResponse::not_found(format!("Trace '{}'", trace_id))),
        ));
    }

    let start_time = spans.iter().map(|s| s.start_time).min().unwrap_or(0);
    let end_time = spans.iter().map(|s| s.end_time).max().unwrap_or(0);
    let duration = end_time - start_time;

    let service_names: Vec<String> = {
        let mut names: Vec<String> = spans
            .iter()
            .filter_map(|s| s.resource.as_ref())
            .filter_map(|r| r.attributes.get("service.name"))
            .cloned()
            .collect::<std::collections::HashSet<_>>()
            .into_iter()
            .collect();
        names.sort();
        names
    };

    let span_entries: Vec<SpanEntry> = spans.into_iter().map(SpanEntry::from).collect();

    let span_count = span_entries.len();

    let trace_detail = TraceDetail {
        trace_id,
        spans: span_entries,
        start_time,
        end_time,
        duration,
        span_count,
        service_names,
    };

    Ok(Json(trace_detail))
}

/// Handler for GET /api/traces/{trace_id}/logs
///
/// Returns all logs associated with a trace, enabling single-round-trip
/// trace→log correlation without a separate API call.
#[utoipa::path(
    get,
    path = "/api/traces/{trace_id}/logs",
    params(
        ("trace_id" = String, Path, description = "Trace ID"),
        ("limit" = Option<usize>, Query, description = "Maximum number of log results (default: 100, max: 1000)"),
        ("search" = Option<String>, Query, description = "Full-text search in log body")
    ),
    responses(
        (status = 200, description = "Logs associated with the trace", body = LogsResponse),
        (status = 500, description = "Internal server error", body = ErrorResponse)
    ),
    tag = "traces"
)]
pub async fn get_trace_logs(
    State(state): State<AppState>,
    Path(trace_id): Path<String>,
    Query(params): Query<TraceLogsQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    let limit = params.limit.unwrap_or(100).min(1000);

    let query = QueryParams {
        trace_id: Some(trace_id.clone()),
        search_text: params.search.clone().filter(|s| !s.is_empty()),
        limit: Some(limit),
        ..Default::default()
    };

    let logs = state.storage.query_logs(&query).await.map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse::storage_error(format!(
                "query logs for trace {}: {}",
                trace_id, e
            ))),
        )
    })?;

    let log_entries: Vec<LogEntry> = logs.into_iter().map(LogEntry::from).collect();
    let total = log_entries.len();

    Ok(Json(LogsResponse {
        logs: log_entries,
        total,
        limit,
        offset: 0,
    }))
}

/// Query parameters for GET /api/traces/{trace_id}/logs
#[derive(Debug, Deserialize, utoipa::IntoParams)]
pub struct TraceLogsQuery {
    /// Maximum number of log results (default: 100, max: 1000)
    #[serde(default)]
    pub limit: Option<usize>,
    /// Full-text search in log body (e.g. "api_request_body" to find request metadata)
    #[serde(default)]
    pub search: Option<String>,
}

/// Export format for traces
#[derive(Debug, Deserialize, utoipa::IntoParams)]
pub struct ExportQuery {
    /// Export format: "json"
    #[serde(default = "default_format")]
    pub format: String,

    /// Same filters as TracesQuery
    #[serde(flatten)]
    pub filters: TracesQuery,
}

fn default_format() -> String {
    "json".to_string()
}

/// Handler for GET /api/traces/export
#[utoipa::path(
    get,
    path = "/api/traces/export",
    params(ExportQuery),
    responses(
        (status = 200, description = "Exported traces in JSON format"),
        (status = 400, description = "Invalid format parameter or time range", body = ErrorResponse),
        (status = 500, description = "Internal server error", body = ErrorResponse)
    ),
    tag = "traces"
)]
pub async fn export_traces(
    State(state): State<AppState>,
    Query(params): Query<ExportQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    validate_time_range(params.filters.start_time, params.filters.end_time)?;

    // Build query parameters (no limit for export, but cap at 10000)
    let query = QueryParams {
        start_time: params.filters.start_time,
        end_time: params.filters.end_time,
        limit: Some(10000),
        trace_id: params.filters.trace_id.clone().filter(|s| !s.is_empty()),
        ..Default::default()
    };

    let spans = state.storage.query_spans(&query).await.map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse::storage_error(format!(
                "export traces: {}",
                e
            ))),
        )
    })?;

    match params.format.as_str() {
        "json" => {
            let span_entries: Vec<SpanEntry> = spans.into_iter().map(SpanEntry::from).collect();

            Ok((
                [
                    ("Content-Type", "application/json"),
                    (
                        "Content-Disposition",
                        "attachment; filename=\"traces.json\"",
                    ),
                ],
                Json(span_entries),
            )
                .into_response())
        },
        _ => Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse::bad_request(
                "Invalid format parameter. Use 'json'",
            )),
        )),
    }
}