From dd2e6c17bc19b839ff0a5e8500dfe40a7612f430 Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Fri, 18 Sep 2026 21:56:13 +0000
Subject: [PATCH] fix(rust): preserve OCR callback headers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm-rust/crates/callbacks-legacy/src/adapter.rs  | 12 ++++++++++--
 .../crates/callbacks-legacy/src/callbacks.rs         |  3 +++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs
index 7204207cd62..53aa6ce9d2e 100644
--- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs
+++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs
@@ -44,6 +44,7 @@ pub struct LegacyLogging {
     response: Option<Py<PyAny>>,
     error: Option<Py<PyBaseException>>,
     body: Option<Py<PyDict>>,
+    headers: Option<Py<PyDict>>,
     context: Option<RequestContext>,
     asynchronous: bool,
     internal: bool,
@@ -77,6 +78,7 @@ impl LegacyLogging {
             response: None,
             error: None,
             body: None,
+            headers: None,
             context: None,
             asynchronous,
             internal: false,
@@ -245,6 +247,7 @@ impl PythonLifecycle for LegacyLogging {
             headers.set_item(name, value)?;
         }
         self.body = Some(body.clone().unbind());
+        self.headers = Some(headers.clone().unbind());
         self.context = Some(context.clone());
         self.logger()?.pre_call(
             py,
@@ -299,8 +302,13 @@ impl PythonLifecycle for LegacyLogging {
                     .as_ref()
                     .and_then(|context| context.api_key.as_ref())
                     .map(|api_key| api_key.expose());
-                self.logger()?
-                    .post_call(py, &raw.body, api_key, self.body.as_ref())?;
+                self.logger()?.post_call(
+                    py,
+                    &raw.body,
+                    api_key,
+                    self.body.as_ref(),
+                    self.headers.as_ref(),
+                )?;
                 Ok(LifecycleStep::Done)
             }
             (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => {
diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs
index 9464f1d6612..9fcfe98368e 100644
--- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs
+++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs
@@ -38,6 +38,7 @@ pub trait LegacyCallbacks {
         original_response: &str,
         api_key: Option<&str>,
         body: Option<&Py<PyDict>>,
+        headers: Option<&Py<PyDict>>,
     ) -> PyResult<()>;
 
     fn defers_async_logging(&self, py: Python<'_>) -> bool;
@@ -150,9 +151,11 @@ impl LegacyCallbacks for PythonLogger {
         original_response: &str,
         api_key: Option<&str>,
         body: Option<&Py<PyDict>>,
+        headers: Option<&Py<PyDict>>,
     ) -> PyResult<()> {
         let additional = PyDict::new(py);
         additional.set_item("complete_input_dict", body)?;
+        additional.set_item("headers", headers)?;
         Logging::PostCall.call(
             py,
             (self.object(py), original_response, api_key, &additional),