feat: improve trace ingestion and trace details

litellm-rust/Cargo.toml

@@ -17,6 +17,7 @@ litellm-core = { path = "crates/core" }
 litellm-gateway-mcp = { path = "crates/gateway-mcp" }
 litellm-gateway = { path = "crates/gateway" }
 litellm-gateway-inference = { path = "crates/gateway-inference" }
+litellm-gateway-traces = { path = "crates/gateway-traces" }
 litellm-gateway-auth = { path = "crates/gateway-auth" }
 litellm-gateway-management = { path = "crates/gateway-management" }
 litellm-gateway-ui = { path = "crates/gateway-ui" }
@@ -81,6 +82,7 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "mul
 qdrant-client = { version = "1.19.0", default-features = false }
 uuid = { version = "1", features = ["v4"] }
 rstest = "0.26.1"
+wiremock = "0.6.5"
 rstest_reuse = "0.7.0"
 rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
 rustify = "=0.7.0"
@@ -115,6 +117,8 @@ time = { version = "0.3.53", features = ["parsing"] }
 criterion = "0.8.2"
 fancy-regex = "0.19.2"
 veil = "0.3.0"
+prost = "0.14.4"
+opentelemetry-proto = "0.33"
 
 [profile.release]
 opt-level = 3

litellm-rust/crates/cache-azure-blob/Cargo.toml

@@ -26,4 +26,4 @@ litellm-cache-testing.workspace = true
 rstest.workspace = true
 serde_json.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
-wiremock = "0.6.5"
+wiremock.workspace = true

litellm-rust/crates/cache-gcs/Cargo.toml

@@ -21,4 +21,4 @@ litellm-cache-testing.workspace = true
 rstest.workspace = true
 serde_json.workspace = true
 tokio.workspace = true
-wiremock = "0.6.5"
+wiremock.workspace = true

litellm-rust/crates/cache-response/Cargo.toml

@@ -21,4 +21,4 @@ redis = "1.7.0"
 redis-test = "1.0.4"
 rstest.workspace = true
 tokio.workspace = true
-wiremock = "0.6.5"
+wiremock.workspace = true

litellm-rust/crates/cache-s3/Cargo.toml

@@ -23,6 +23,6 @@ tokio.workspace = true
 litellm-http = { workspace = true, features = ["test-support"] }
 litellm-cache-testing.workspace = true
 rstest.workspace = true
-wiremock = "0.6.5"
+wiremock.workspace = true
 serde_json.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

litellm-rust/crates/core/Cargo.toml

@@ -47,4 +47,4 @@ litellm-host-native.workspace = true
 litellm-llms = { workspace = true, features = ["test-support"] }
 rstest.workspace = true
 rstest_reuse.workspace = true
-wiremock = "0.6.5"
+wiremock.workspace = true

litellm-rust/crates/gateway-inference/Cargo.toml

@@ -30,4 +30,4 @@ futures-util.workspace = true
 tokio = { workspace = true, features = ["io-util"] }
 rstest.workspace = true
 tower = { version = "0.5.3", features = ["util"] }
-wiremock = "0.6.5"
+wiremock.workspace = true

litellm-rust/crates/gateway-traces/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "litellm-gateway-traces"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+axum.workspace = true
+flate2.workspace = true
+litellm-traces.workspace = true
+
+[dev-dependencies]
+prost.workspace = true
+opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] }
+rstest.workspace = true
+tokio.workspace = true
+tower = { version = "0.5", features = ["util"] }

litellm-rust/crates/gateway-traces/src/lib.rs

@@ -0,0 +1,125 @@
+use std::{future::Future, io::Read, sync::Arc};
+
+use axum::{
+    Router,
+    body::to_bytes,
+    extract::{Request, State},
+    http::{StatusCode, header},
+    response::{IntoResponse, Response},
+    routing::post,
+};
+use flate2::read::MultiGzDecoder;
+use litellm_traces::{DecodeError, DecodedSpan, decode_otlp};
+
+pub const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
+
+fn otlp_max_body_bytes() -> usize {
+    std::env::var("OTLP_MAX_BODY_BYTES")
+        .ok()
+        .and_then(|value| value.trim().parse().ok())
+        .unwrap_or(MAX_BODY_BYTES)
+}
+
+pub trait SpanSink: Send + Sync + 'static {
+    fn write(&self, spans: Vec<DecodedSpan>) -> impl Future<Output = Result<(), ()>> + Send;
+}
+
+pub fn router<S: SpanSink>(sink: S) -> Router {
+    Router::new()
+        .route("/v1/traces", post(ingest::<S>))
+        .with_state(Arc::new(sink))
+}
+
+async fn ingest<S: SpanSink>(State(sink): State<Arc<S>>, request: Request) -> Response {
+    let max_body_bytes = otlp_max_body_bytes();
+    let content_type = request
+        .headers()
+        .get(header::CONTENT_TYPE)
+        .and_then(|value| value.to_str().ok())
+        .map(str::to_owned);
+    let encoding = request
+        .headers()
+        .get(header::CONTENT_ENCODING)
+        .map(|value| value.to_str().map(str::to_owned));
+    let response_type = if content_type.as_deref().is_some_and(|value| {
+        value
+            .split(';')
+            .next()
+            .is_some_and(|value| value.trim().eq_ignore_ascii_case("application/json"))
+    }) {
+        "application/json"
+    } else {
+        "application/x-protobuf"
+    };
+    let result = async {
+        let encoding = encoding.transpose().map_err(|_| StatusCode::BAD_REQUEST)?;
+        let body = to_bytes(request.into_body(), max_body_bytes.saturating_add(1))
+            .await
+            .map_err(|_| StatusCode::PAYLOAD_TOO_LARGE)?;
+        let payload = decode_content_encoding(&body, encoding.as_deref(), max_body_bytes)?;
+        let spans =
+            decode_otlp(&payload, content_type.as_deref()).map_err(|error| match error {
+                DecodeError::TooLarge => StatusCode::PAYLOAD_TOO_LARGE,
+                DecodeError::InvalidPayload => StatusCode::BAD_REQUEST,
+            })?;
+        sink.write(spans)
+            .await
+            .map_err(|_| StatusCode::SERVICE_UNAVAILABLE)
+    }
+    .await;
+    match result {
+        Ok(()) if response_type == "application/json" => {
+            ([(header::CONTENT_TYPE, response_type)], "{}").into_response()
+        }
+        Ok(()) => ([(header::CONTENT_TYPE, response_type)], "").into_response(),
+        Err(status) => status.into_response(),
+    }
+}
+
+fn decode_content_encoding(
+    body: &[u8],
+    encoding: Option<&str>,
+    max_body_bytes: usize,
+) -> Result<Vec<u8>, StatusCode> {
+    if body.len() > max_body_bytes {
+        return Err(StatusCode::PAYLOAD_TOO_LARGE);
+    }
+    match encoding {
+        None => Ok(body.to_vec()),
+        Some(value) if value.eq_ignore_ascii_case("identity") => Ok(body.to_vec()),
+        Some(value) if value.eq_ignore_ascii_case("gzip") => {
+            let mut payload = Vec::new();
+            MultiGzDecoder::new(body)
+                .take((max_body_bytes as u64).saturating_add(1))
+                .read_to_end(&mut payload)
+                .map_err(|_| StatusCode::BAD_REQUEST)?;
+            if payload.len() > max_body_bytes {
+                return Err(StatusCode::PAYLOAD_TOO_LARGE);
+            }
+            Ok(payload)
+        }
+        _ => Err(StatusCode::BAD_REQUEST),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::Write;
+
+    use flate2::{Compression, write::GzEncoder};
+    use rstest::rstest;
+
+    use super::*;
+
+    #[rstest]
+    fn decompressed_size_limit_rejects_expansion() {
+        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
+        encoder.write_all(&vec![b' '; MAX_BODY_BYTES + 1]).unwrap();
+        let compressed = encoder.finish().unwrap();
+        assert!(compressed.len() < MAX_BODY_BYTES);
+        assert!(matches!(
+            decode_content_encoding(&compressed, Some("gzip"), MAX_BODY_BYTES),
+            Err(StatusCode::PAYLOAD_TOO_LARGE)
+        ));
+    }
+}