feat(rust-gateway): sniff upload MIME by magic bytes and preserve reducto ids

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm-rust/crates/ai-gateway/src/constants.rs

@@ -5,6 +5,11 @@
 //! modules. Env-overridable tunables keep their `DEFAULT_*` value here; the env
 //! read + fallback happens at the host/config layer.
 
+use litellm_core::constants::{
+    MIME_APPLICATION_OCTET_STREAM, MIME_APPLICATION_PDF, MIME_IMAGE_BMP, MIME_IMAGE_GIF,
+    MIME_IMAGE_JPEG, MIME_IMAGE_PNG, MIME_IMAGE_TIFF, MIME_IMAGE_WEBP,
+};
+
 /// Default LiteLLM control-plane base URL for request-log egress when
 /// `LITELLM_PROXY_BASE_URL` is unset.
 pub(crate) const DEFAULT_PROXY_BASE_URL: &str = "http://localhost:4000";
@@ -28,18 +33,18 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
 /// Provider attributed to realtime sessions in the logging payload.
 pub(crate) const DEFAULT_PROVIDER: &str = "openai";
 
-pub(crate) const DEFAULT_UPLOAD_MIME_TYPE: &str = "application/octet-stream";
+pub(crate) const DEFAULT_UPLOAD_MIME_TYPE: &str = MIME_APPLICATION_OCTET_STREAM;
 
 pub(crate) const MAX_OCR_REQUEST_BYTES: usize = 100 * 1024 * 1024;
 
 pub(crate) const OCR_UPLOAD_MIME_BY_EXTENSION: &[(&str, &str)] = &[
-    ("pdf", "application/pdf"),
-    ("png", "image/png"),
-    ("jpg", "image/jpeg"),
-    ("jpeg", "image/jpeg"),
-    ("gif", "image/gif"),
-    ("webp", "image/webp"),
-    ("tiff", "image/tiff"),
-    ("tif", "image/tiff"),
-    ("bmp", "image/bmp"),
+    ("pdf", MIME_APPLICATION_PDF),
+    ("png", MIME_IMAGE_PNG),
+    ("jpg", MIME_IMAGE_JPEG),
+    ("jpeg", MIME_IMAGE_JPEG),
+    ("gif", MIME_IMAGE_GIF),
+    ("webp", MIME_IMAGE_WEBP),
+    ("tiff", MIME_IMAGE_TIFF),
+    ("tif", MIME_IMAGE_TIFF),
+    ("bmp", MIME_IMAGE_BMP),
 ];

litellm-rust/crates/ai-gateway/src/routes/ocr/mod.rs

@@ -129,7 +129,7 @@ fn error_response(error: &CoreError) -> Response {
         CoreError::Auth(_) => (
             StatusCode::UNAUTHORIZED,
             "authentication_error",
-            error.to_string(),
+            "authentication failed".to_string(),
         ),
         CoreError::Routing(_) => (StatusCode::NOT_FOUND, "not_found_error", error.to_string()),
         CoreError::Http { status, .. } => {
@@ -197,6 +197,61 @@ mod tests {
         String::from_utf8_lossy(&request).into_owned()
     }
 
+    async fn read_full_http_request(socket: &mut TcpStream) -> String {
+        let mut request = Vec::new();
+        let mut buffer = [0_u8; 4096];
+        let mut body_start: Option<usize> = None;
+        let mut content_length = 0_usize;
+        loop {
+            let n = socket.read(&mut buffer).await.expect("reads request");
+            if n == 0 {
+                break;
+            }
+            request.extend_from_slice(&buffer[..n]);
+            if body_start.is_none() {
+                if let Some(pos) = request.windows(4).position(|window| window == b"\r\n\r\n") {
+                    let start = pos + 4;
+                    body_start = Some(start);
+                    let headers = String::from_utf8_lossy(&request[..pos]).to_ascii_lowercase();
+                    content_length = headers
+                        .lines()
+                        .find_map(|line| line.strip_prefix("content-length:"))
+                        .and_then(|value| value.trim().parse::<usize>().ok())
+                        .unwrap_or(0);
+                }
+            }
+            if let Some(start) = body_start {
+                if request.len() >= start + content_length {
+                    break;
+                }
+            }
+        }
+        String::from_utf8_lossy(&request).into_owned()
+    }
+
+    async fn spawn_mock_upstream_capture() -> (String, tokio::task::JoinHandle<String>) {
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("binds upstream");
+        let addr = listener.local_addr().expect("upstream addr");
+        let handle = tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.expect("accepts one request");
+            let request = read_full_http_request(&mut socket).await;
+            let body = r#"{"pages":[{"index":0,"markdown":"hello ocr"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
+            let response = format!(
+                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+                body.len(),
+                body
+            );
+            socket
+                .write_all(response.as_bytes())
+                .await
+                .expect("writes response");
+            request
+        });
+        (format!("http://{addr}"), handle)
+    }
+
     async fn spawn_mock_upstream() -> (String, tokio::task::JoinHandle<String>) {
         let listener = TcpListener::bind("127.0.0.1:0")
             .await
@@ -374,6 +429,67 @@ mod tests {
         assert!(upstream_request.starts_with("POST"), "{upstream_request}");
     }
 
+    #[tokio::test]
+    async fn multipart_unnamed_octet_stream_pdf_is_sniffed() {
+        let (upstream, upstream_handle) = spawn_mock_upstream_capture().await;
+        let addr = serve(app_with_deployment(&upstream)).await;
+
+        let form = reqwest::multipart::Form::new()
+            .text("model", "rust-ocr-mistral")
+            .part(
+                "file",
+                reqwest::multipart::Part::bytes(b"%PDF-1.7 minimal pdf bytes".to_vec())
+                    .mime_str("application/octet-stream")
+                    .expect("mime"),
+            );
+
+        let response = reqwest::Client::new()
+            .post(format!("http://{addr}/v1/ocr"))
+            .bearer_auth(MASTER_KEY)
+            .multipart(form)
+            .send()
+            .await
+            .expect("request sent");
+
+        assert_eq!(response.status(), reqwest::StatusCode::OK);
+        let upstream_request = upstream_handle.await.expect("upstream served");
+        assert!(
+            upstream_request.contains("data:application/pdf;base64,"),
+            "unnamed octet-stream PDF must be sniffed to application/pdf: {upstream_request}"
+        );
+    }
+
+    #[tokio::test]
+    async fn multipart_unnamed_octet_stream_image_is_sniffed() {
+        let (upstream, upstream_handle) = spawn_mock_upstream_capture().await;
+        let addr = serve(app_with_deployment(&upstream)).await;
+
+        let png = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x01];
+        let form = reqwest::multipart::Form::new()
+            .text("model", "rust-ocr-mistral")
+            .part(
+                "file",
+                reqwest::multipart::Part::bytes(png)
+                    .mime_str("application/octet-stream")
+                    .expect("mime"),
+            );
+
+        let response = reqwest::Client::new()
+            .post(format!("http://{addr}/v1/ocr"))
+            .bearer_auth(MASTER_KEY)
+            .multipart(form)
+            .send()
+            .await
+            .expect("request sent");
+
+        assert_eq!(response.status(), reqwest::StatusCode::OK);
+        let upstream_request = upstream_handle.await.expect("upstream served");
+        assert!(
+            upstream_request.contains("data:image/png;base64,"),
+            "unnamed octet-stream PNG must be sniffed to image/png: {upstream_request}"
+        );
+    }
+
     #[tokio::test]
     async fn upstream_error_status_is_propagated_without_leaking_provider_body() {
         let upstream = spawn_mock_upstream_error().await;

litellm-rust/crates/ai-gateway/src/routes/ocr/transport.rs

@@ -3,6 +3,7 @@ use std::time::Duration;
 use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
 use base64::Engine;
 use litellm_core::error::CoreError;
+use litellm_core::ocr::mime::sniff_mime;
 use litellm_core::CoreResult;
 use serde::Deserialize;
 use serde_json::{json, Map, Value};
@@ -31,14 +32,6 @@ impl OcrDocument {
                 ))
             }
         };
-        if url.starts_with("reducto://") {
-            return Err(CoreError::InvalidRequest(
-                "reducto:// file IDs are not accepted through the OCR API; upload the file in \
-                 the same request via multipart/form-data with a 'file' field, or pass an \
-                 inline base64 data URI as the document URL"
-                    .to_string(),
-            ));
-        }
         Ok(json!({ "type": field, field: url }))
     }
 }
@@ -95,7 +88,7 @@ fn mime_from_filename(filename: &str) -> Option<&'static str> {
         .map(|(_, mime)| *mime)
 }
 
-fn resolve_upload_mime(content_type: Option<&str>, filename: Option<&str>) -> String {
+fn resolve_upload_mime(bytes: &[u8], content_type: Option<&str>, filename: Option<&str>) -> String {
     let declared = content_type
         .and_then(|value| value.split(';').next())
         .map(str::trim)
@@ -103,6 +96,9 @@ fn resolve_upload_mime(content_type: Option<&str>, filename: Option<&str>) -> St
     if let Some(declared) = declared {
         return declared.to_string();
     }
+    if let Some(sniffed) = sniff_mime(bytes) {
+        return sniffed.to_string();
+    }
     filename
         .and_then(mime_from_filename)
         .unwrap_or(DEFAULT_UPLOAD_MIME_TYPE)
@@ -119,7 +115,7 @@ pub fn build_upload_document(
             "uploaded file is empty".to_string(),
         ));
     }
-    let mime = resolve_upload_mime(content_type, filename);
+    let mime = resolve_upload_mime(&bytes, content_type, filename);
     let data_uri = format!("data:{mime};base64,{}", BASE64_STANDARD.encode(&bytes));
     let field = if mime.starts_with("image/") {
         "image_url"
@@ -221,15 +217,13 @@ mod tests {
     }
 
     #[test]
-    fn rejects_reducto_file_id_over_json() {
-        let err = parse_json_body(
-            br#"{"model":"m","document":{"type":"document_url","document_url":"reducto://abc"}}"#,
+    fn preserves_reducto_file_id_over_json() {
+        let call = parse_json_body(
+            br#"{"model":"m","document":{"type":"document_url","document_url":"reducto://abc123"}}"#,
         )
-        .expect_err("reducto id rejected");
-        match err {
-            CoreError::InvalidRequest(message) => assert!(message.contains("reducto://")),
-            other => panic!("expected InvalidRequest, got {other:?}"),
-        }
+        .expect("reducto id preserved");
+        assert_eq!(call.document["type"], "document_url");
+        assert_eq!(call.document["document_url"], "reducto://abc123");
     }
 
     #[test]
@@ -255,6 +249,36 @@ mod tests {
             .starts_with("data:application/pdf;base64,"));
     }
 
+    #[test]
+    fn upload_sniffs_pdf_from_bytes_when_unnamed_octet_stream() {
+        let document = build_upload_document(
+            b"%PDF-1.7 minimal".to_vec(),
+            None,
+            Some("application/octet-stream"),
+        )
+        .expect("builds document");
+        assert_eq!(document["type"], "document_url");
+        assert!(document["document_url"]
+            .as_str()
+            .expect("data uri")
+            .starts_with("data:application/pdf;base64,"));
+    }
+
+    #[test]
+    fn upload_sniffs_png_from_bytes_when_unnamed_and_no_content_type() {
+        let document = build_upload_document(
+            vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A, 0x00],
+            None,
+            None,
+        )
+        .expect("builds document");
+        assert_eq!(document["type"], "image_url");
+        assert!(document["image_url"]
+            .as_str()
+            .expect("data uri")
+            .starts_with("data:image/png;base64,"));
+    }
+
     #[test]
     fn upload_infers_mime_from_filename_when_octet_stream() {
         let document = build_upload_document(

litellm-rust/crates/core/src/constants.rs

@@ -0,0 +1,8 @@
+pub const MIME_APPLICATION_PDF: &str = "application/pdf";
+pub const MIME_APPLICATION_OCTET_STREAM: &str = "application/octet-stream";
+pub const MIME_IMAGE_PNG: &str = "image/png";
+pub const MIME_IMAGE_JPEG: &str = "image/jpeg";
+pub const MIME_IMAGE_GIF: &str = "image/gif";
+pub const MIME_IMAGE_WEBP: &str = "image/webp";
+pub const MIME_IMAGE_TIFF: &str = "image/tiff";
+pub const MIME_IMAGE_BMP: &str = "image/bmp";

litellm-rust/crates/core/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod constants;
 pub mod error;
 pub mod ocr;
 pub mod providers;

litellm-rust/crates/core/src/ocr/mime.rs

@@ -0,0 +1,93 @@
+use crate::constants::{
+    MIME_APPLICATION_PDF, MIME_IMAGE_BMP, MIME_IMAGE_GIF, MIME_IMAGE_JPEG, MIME_IMAGE_PNG,
+    MIME_IMAGE_TIFF, MIME_IMAGE_WEBP,
+};
+
+pub fn sniff_mime(bytes: &[u8]) -> Option<&'static str> {
+    if bytes.starts_with(b"%PDF") {
+        return Some(MIME_APPLICATION_PDF);
+    }
+    if bytes.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
+        return Some(MIME_IMAGE_PNG);
+    }
+    if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) {
+        return Some(MIME_IMAGE_JPEG);
+    }
+    if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
+        return Some(MIME_IMAGE_GIF);
+    }
+    if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
+        return Some(MIME_IMAGE_WEBP);
+    }
+    if bytes.starts_with(&[0x49, 0x49, 0x2A, 0x00]) || bytes.starts_with(&[0x4D, 0x4D, 0x00, 0x2A])
+    {
+        return Some(MIME_IMAGE_TIFF);
+    }
+    if bytes.starts_with(b"BM") {
+        return Some(MIME_IMAGE_BMP);
+    }
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn sniffs_pdf() {
+        assert_eq!(sniff_mime(b"%PDF-1.7\n..."), Some(MIME_APPLICATION_PDF));
+    }
+
+    #[test]
+    fn sniffs_png() {
+        assert_eq!(
+            sniff_mime(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A, 0x00]),
+            Some(MIME_IMAGE_PNG)
+        );
+    }
+
+    #[test]
+    fn sniffs_jpeg() {
+        assert_eq!(sniff_mime(&[0xFF, 0xD8, 0xFF, 0xE0]), Some(MIME_IMAGE_JPEG));
+    }
+
+    #[test]
+    fn sniffs_gif() {
+        assert_eq!(sniff_mime(b"GIF87a...."), Some(MIME_IMAGE_GIF));
+        assert_eq!(sniff_mime(b"GIF89a...."), Some(MIME_IMAGE_GIF));
+    }
+
+    #[test]
+    fn sniffs_webp() {
+        let mut bytes = b"RIFF".to_vec();
+        bytes.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
+        bytes.extend_from_slice(b"WEBP");
+        assert_eq!(sniff_mime(&bytes), Some(MIME_IMAGE_WEBP));
+    }
+
+    #[test]
+    fn does_not_sniff_riff_without_webp() {
+        let mut bytes = b"RIFF".to_vec();
+        bytes.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
+        bytes.extend_from_slice(b"WAVE");
+        assert_eq!(sniff_mime(&bytes), None);
+    }
+
+    #[test]
+    fn sniffs_tiff_both_byte_orders() {
+        assert_eq!(sniff_mime(&[0x49, 0x49, 0x2A, 0x00]), Some(MIME_IMAGE_TIFF));
+        assert_eq!(sniff_mime(&[0x4D, 0x4D, 0x00, 0x2A]), Some(MIME_IMAGE_TIFF));
+    }
+
+    #[test]
+    fn sniffs_bmp() {
+        assert_eq!(sniff_mime(b"BM...."), Some(MIME_IMAGE_BMP));
+    }
+
+    #[test]
+    fn returns_none_for_unknown_and_truncated() {
+        assert_eq!(sniff_mime(&[0x00, 0x01, 0x02, 0x03]), None);
+        assert_eq!(sniff_mime(b""), None);
+        assert_eq!(sniff_mime(&[0x89, b'P', b'N', b'G']), None);
+    }
+}

litellm-rust/crates/core/src/ocr/mod.rs

@@ -1,2 +1,3 @@
+pub mod mime;
 pub mod transformation;
 pub mod types;