fix(ocr): keep inlined remote documents when the host restores retained fields

litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs

@@ -134,7 +134,66 @@ fn invalid_api_base() -> crate::ocr::Error {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::{Arc, Mutex};
+
+    use base64::{Engine, engine::general_purpose::STANDARD};
+    use serde_json::json;
+
     use super::*;
+    use crate::ocr::test_support::{
+        MockResponse, RetainedFieldsHost, mock_server, perform_ocr, wire_request_with_document,
+        with_hooks, with_source,
+    };
+
+    #[tokio::test]
+    async fn remote_image_stays_inlined_when_the_host_restores_retained_fields() {
+        let (base, seen, server) = mock_server(vec![
+            MockResponse::png(json!("pixels")),
+            MockResponse::json(json!({"pages":[{"index":0,"markdown":{"content":"hello"}}]})),
+        ])
+        .await;
+        let source = format!("{base}/scan.png");
+        let retained_fields = Arc::new(Mutex::new(Vec::new()));
+        let request = with_hooks(
+            with_source(
+                wire_request_with_document(
+                    "azure_ai/cohere-parse-v5",
+                    &base,
+                    json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}),
+                    json!({}),
+                ),
+                &source,
+            ),
+            Arc::new(RetainedFieldsHost {
+                original_document: json!({"type":"image_url","image_url":source}),
+                retained_fields: retained_fields.clone(),
+            }),
+        );
+
+        let result = perform_ocr(request).await.unwrap();
+        server.await.unwrap();
+        assert_eq!(result.pages[0].markdown, "hello");
+        assert!(
+            !retained_fields
+                .lock()
+                .unwrap()
+                .iter()
+                .any(|name| name == "document")
+        );
+        let requests = seen.lock().unwrap();
+        assert_eq!(requests.len(), 2);
+        assert!(requests[0].starts_with("GET /scan.png "));
+        assert!(requests[1].starts_with("POST /providers/cohere/v2/parse "));
+        let body: Value =
+            serde_json::from_str(requests[1].split_once("\r\n\r\n").unwrap().1).unwrap();
+        assert_eq!(
+            body["document"],
+            json!({
+                "type":"image_url",
+                "image_url":format!("data:image/png;base64,{}", STANDARD.encode(br#""pixels""#))
+            })
+        );
+    }
 
     #[test]
     fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() {

litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs

@@ -304,12 +304,16 @@ mod tests {
         );
     }
 
-    use std::sync::Arc;
+    use std::sync::{Arc, Mutex};
 
+    use base64::{Engine, engine::general_purpose::STANDARD};
     use serde_json::json;
 
     use crate::ocr::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
-    use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+    use crate::ocr::test_support::{
+        MockResponse, RetainedFieldsHost, mock_server, perform_ocr, wire_request, with_hooks,
+        with_source,
+    };
 
     #[tokio::test]
     async fn facade_executes_azure_mistral_with_prepared_auth() {
@@ -374,6 +378,54 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn remote_document_stays_inlined_when_the_host_restores_retained_fields() {
+        let (base, seen, server) = mock_server(vec![
+            MockResponse::json(json!({"page":"bytes"})),
+            MockResponse::json(json!({
+                "pages":[{"index":0,"markdown":"hello"}],
+                "usage_info":{"pages_processed":1}
+            })),
+        ])
+        .await;
+        let source = format!("{base}/scan.pdf");
+        let retained_fields = Arc::new(Mutex::new(Vec::new()));
+        let request = with_hooks(
+            with_source(wire_request("azure_ai/model", &base, json!({})), &source),
+            Arc::new(RetainedFieldsHost {
+                original_document: json!({"type":"document_url","document_url":source}),
+                retained_fields: retained_fields.clone(),
+            }),
+        );
+
+        let result = perform_ocr(request).await.unwrap();
+        server.await.unwrap();
+        assert_eq!(result.pages[0].markdown, "hello");
+        assert!(
+            !retained_fields
+                .lock()
+                .unwrap()
+                .iter()
+                .any(|name| name == "document")
+        );
+        let requests = seen.lock().unwrap();
+        assert_eq!(requests.len(), 2);
+        assert!(requests[0].starts_with("GET /scan.pdf "));
+        assert!(requests[1].starts_with("POST /providers/mistral/azure/ocr "));
+        let body: Value =
+            serde_json::from_str(requests[1].split_once("\r\n\r\n").unwrap().1).unwrap();
+        assert_eq!(
+            body["document"],
+            json!({
+                "type":"document_url",
+                "document_url":format!(
+                    "data:application/json;base64,{}",
+                    STANDARD.encode(br#"{"page":"bytes"}"#)
+                )
+            })
+        );
+    }
+
     struct ReplaceBodyDocument;
 
     impl OcrHooks for ReplaceBodyDocument {

litellm-rust/crates/core/src/llms/vertex_ai/ocr/transformation.rs

@@ -220,10 +220,15 @@ mod tests {
         );
     }
 
+    use std::sync::{Arc, Mutex};
+
     use litellm_auth::InputSource;
     use serde_json::{Value, json};
 
-    use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+    use crate::ocr::test_support::{
+        MockResponse, RetainedFieldsHost, mock_server, perform_ocr, wire_request, with_hooks,
+        with_source,
+    };
 
     fn request_body(request: &str) -> Value {
         serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
@@ -269,6 +274,50 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn remote_document_stays_inlined_when_the_host_restores_retained_fields() {
+        let (base, seen, server) = mock_server(vec![
+            MockResponse::json(json!({"page":"bytes"})),
+            MockResponse::json(json!({"pages":[]})),
+        ])
+        .await;
+        let source = format!("{base}/scan.pdf");
+        let retained_fields = Arc::new(Mutex::new(Vec::new()));
+        let request = with_hooks(
+            with_source(
+                wire_request(
+                    "vertex_ai/mistral-ocr-maas",
+                    &base,
+                    json!({"vertex_project":"project-1","vertex_location":"europe-west4"}),
+                ),
+                &source,
+            ),
+            Arc::new(RetainedFieldsHost {
+                original_document: json!({"type":"document_url","document_url":source}),
+                retained_fields: retained_fields.clone(),
+            }),
+        );
+
+        perform_ocr(request).await.unwrap();
+        server.await.unwrap();
+        assert!(
+            !retained_fields
+                .lock()
+                .unwrap()
+                .iter()
+                .any(|name| name == "document")
+        );
+        let requests = seen.lock().unwrap();
+        assert_eq!(requests.len(), 2);
+        assert!(requests[0].starts_with("GET /scan.pdf "));
+        assert!(
+            request_body(&requests[1])["document"]["document_url"]
+                .as_str()
+                .unwrap()
+                .starts_with("data:application/json;base64,")
+        );
+    }
+
     #[tokio::test]
     async fn supplied_authorization_is_forwarded_without_a_static_token() {
         let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;

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

@@ -27,12 +27,7 @@ where
         .keys()
         .filter(|name| composed.get(*name).is_some())
         .cloned()
-        .chain(
-            composed
-                .get("document")
-                .is_some()
-                .then(|| "document".to_string()),
-        )
+        .chain(retains_document(&composed, &request.document).then(|| "document".to_string()))
         .collect();
     let (body, headers) = if request.hooks.intercepts_requests() {
         let changed = request
@@ -106,6 +101,10 @@ pub(crate) async fn guardrail_document(
     Ok((document, changed.headers))
 }
 
+fn retains_document(composed: &Value, document: &OcrDocument) -> bool {
+    body_document(composed).is_ok_and(|sent| sent.source() == document.source())
+}
+
 pub(crate) fn body_document(body: &Value) -> Result<OcrDocument, super::Error> {
     let document = body
         .get("document")

litellm-rust/crates/core/tests/ocr.rs

@@ -8,7 +8,10 @@ use super::hooks::{
     OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
     OcrPreCallRequest,
 };
-use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+use super::test_support::{
+    MockResponse, RetainedFieldsHost, mock_server, perform_ocr, wire_request, with_hooks,
+    with_source,
+};
 use super::wire::{OcrWireRequest, decode_request};
 use super::{
     NativeOutcome, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHost,
@@ -102,6 +105,36 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() {
     );
 }
 
+#[tokio::test]
+async fn facade_keeps_the_host_document_when_the_url_is_sent_unchanged() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
+    let source = format!("{base}/scan.pdf");
+    let original_document = json!({
+        "type":"document_url",
+        "document_url":source,
+        "document_name":"scan.pdf"
+    });
+    let retained_fields = Arc::new(Mutex::new(Vec::new()));
+    let request = with_hooks(
+        with_source(wire_request("mistral/model", &base, json!({})), &source),
+        Arc::new(RetainedFieldsHost {
+            original_document: original_document.clone(),
+            retained_fields: retained_fields.clone(),
+        }),
+    );
+
+    perform_ocr(request).await.unwrap();
+    server.await.unwrap();
+    assert_eq!(
+        *retained_fields.lock().unwrap(),
+        vec!["document".to_string()]
+    );
+    let requests = seen.lock().unwrap();
+    assert_eq!(requests.len(), 1);
+    let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap();
+    assert_eq!(body["document"], original_document);
+}
+
 #[tokio::test]
 async fn facade_executes_direct_mistral_once() {
     let (base, seen, server) = mock_server(vec![MockResponse::json(json!({

litellm-rust/crates/core/tests/ocr/support.rs

@@ -4,6 +4,7 @@ use serde_json::{Value, json};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpListener;
 
+use crate::ocr::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
 use crate::ocr::wire::{OcrWireRequest, decode_request};
 use crate::ocr::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
 
@@ -22,9 +23,23 @@ pub(crate) async fn perform_ocr(
 }
 
 pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest {
+    wire_request_with_document(
+        model,
+        base,
+        json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
+        options,
+    )
+}
+
+pub(crate) fn wire_request_with_document(
+    model: &str,
+    base: &str,
+    document: Value,
+    options: Value,
+) -> LiteLLMOcrRequest {
     decode_request(OcrWireRequest {
         model: model.into(),
-        document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
+        document,
         api_key: Some("test-key".into()),
         api_base: Some(base.into()),
         custom_llm_provider: None,
@@ -50,6 +65,41 @@ pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOc
     request.with_document(document.into())
 }
 
+pub(crate) fn with_hooks(
+    request: LiteLLMOcrRequest,
+    hooks: Arc<dyn OcrHooks>,
+) -> LiteLLMOcrRequest {
+    LiteLLMOcrRequest { hooks, ..request }
+}
+
+pub(crate) struct RetainedFieldsHost {
+    pub original_document: Value,
+    pub retained_fields: Arc<Mutex<Vec<String>>>,
+}
+
+impl OcrHooks for RetainedFieldsHost {
+    fn intercepts_requests(&self) -> bool {
+        true
+    }
+
+    fn during_call(
+        &self,
+        mut request: OcrDuringCallRequest,
+    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
+        Box::pin(async move {
+            *self.retained_fields.lock().unwrap() = request.retained_fields.clone();
+            if request
+                .retained_fields
+                .iter()
+                .any(|name| name == "document")
+            {
+                request.body["document"] = self.original_document.clone();
+            }
+            Ok(request)
+        })
+    }
+}
+
 pub(crate) struct MockResponse {
     pub status: u16,
     pub headers: Vec<(&'static str, String)>,
@@ -64,6 +114,13 @@ impl MockResponse {
             body,
         }
     }
+
+    pub fn png(body: Value) -> Self {
+        Self {
+            headers: vec![("Content-Type", "image/png".into())],
+            ..Self::json(body)
+        }
+    }
 }
 
 pub(crate) async fn mock_server(
@@ -104,6 +161,15 @@ pub(crate) async fn mock_server(
                 .unwrap()
                 .push(String::from_utf8_lossy(&bytes).into_owned());
             let body = serde_json::to_vec(&response.body).unwrap();
+            let default_content_type = if response
+                .headers
+                .iter()
+                .any(|(name, _)| name.eq_ignore_ascii_case("content-type"))
+            {
+                ""
+            } else {
+                "Content-Type: application/json\r\n"
+            };
             let headers = response
                 .headers
                 .into_iter()
@@ -112,8 +178,9 @@ pub(crate) async fn mock_server(
                 })
                 .collect::<String>();
             let head = format!(
-                "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n",
+                "HTTP/1.1 {} OK\r\n{}Content-Length: {}\r\nConnection: close\r\n{}\r\n",
                 response.status,
+                default_content_type,
                 body.len(),
                 headers
             );